[XPU] recompute ut for xpu (#69006)

houj04 · web-flow · commit 909eb33a6976 · 2024-11-04T16:53:47.000+08:00
* [XPU] debug for recompute ut

* fix loss scaling, and add xpu support.
diff --git a/test/collective/fleet/test_dygraph_recompute.py b/test/collective/fleet/test_dygraph_recompute.py
@@ -149,7 +149,7 @@ def run_model(
     )
 
     if enable_autocast:
-        scaler = paddle.amp.GradScaler()
+        scaler = paddle.amp.GradScaler(init_loss_scaling=4096)
 
     loss_ = []
     param_ = []
@@ -268,7 +268,9 @@ def test_fc_net_with_fp16(self):
         self.test_base_case(enable_autocast=True, pure_fp16=True)
 
     def test_recompute_kwargs(self):
-        paddle.set_device("gpu")
+        paddle.set_device(
+            "xpu" if paddle.base.core.is_compiled_with_xpu() else "gpu"
+        )
         kwargs = {"is_test": False}
         with self.assertRaises(TypeError):
             loss_ref, param_ref, grad_ref = run_model(
diff --git a/test/collective/fleet/test_dygraph_recompute_for_eager.py b/test/collective/fleet/test_dygraph_recompute_for_eager.py
@@ -188,7 +188,7 @@ def run_model(
     )
 
     if enable_autocast:
-        scaler = paddle.amp.GradScaler()
+        scaler = paddle.amp.GradScaler(init_loss_scaling=4096)
 
     loss_ = []
     param_ = []
@@ -340,7 +340,9 @@ def test_fc_net_with_fp16(self):
         )
 
     def test_recompute_kwargs(self):
-        paddle.set_device("gpu")
+        paddle.set_device(
+            "xpu" if paddle.base.core.is_compiled_with_xpu() else "gpu"
+        )
         pos = paddle.randn(shape=[10, 10], dtype="float32")
         pos.stop_gradient = False
 
