[XPU] fix several unit tests (#71966)

houj04 · web-flow · commit 61a8f272583d · 2025-03-31T10:57:59.000+08:00
diff --git a/paddle/phi/kernels/xpu/compare_kernel.cc b/paddle/phi/kernels/xpu/compare_kernel.cc
@@ -34,6 +34,11 @@ void XPUCompareKernelImpl(
                       bool*,
                       const std::vector<int64_t>&,
                       const std::vector<int64_t>&)> func) {
+  auto* out_data = dev_ctx.template Alloc<bool>(out);
+  if (out->numel() == 0) {
+    return;
+  }
+
   auto x_shape = common::vectorize<int64_t>(x.dims());
   auto y_shape = common::vectorize<int64_t>(y.dims());
 
@@ -46,7 +51,6 @@ void XPUCompareKernelImpl(
 
   auto x_data = reinterpret_cast<const XPUType*>(x.data<T>());
   auto y_data = reinterpret_cast<const XPUType*>(y.data<T>());
-  auto* out_data = dev_ctx.template Alloc<bool>(out);
 
   int ret =
       func(dev_ctx.x_context(), x_data, y_data, out_data, x_shape, y_shape);
diff --git a/paddle/phi/kernels/xpu/set_value_kernel.cc b/paddle/phi/kernels/xpu/set_value_kernel.cc
@@ -139,6 +139,11 @@ void SetValueImpl(const Context& dev_ctx,
   int r = XPU_SUCCESS;
   out->Resize(in.dims());
   dev_ctx.template Alloc<T>(out);
+
+  if (in.numel() == 0) {
+    return;
+  }
+
   r = xpu::copy(dev_ctx.x_context(),
                 reinterpret_cast<const XPUType*>(in.data<T>()),
                 reinterpret_cast<XPUType*>(out->data<T>()),
diff --git a/paddle/phi/kernels/xpu/stack_kernel.cc b/paddle/phi/kernels/xpu/stack_kernel.cc
@@ -28,6 +28,15 @@ void StackKernel(const Context& dev_ctx,
   if (axis < 0) {
     axis += x[0]->dims().size() + 1;
   }
+
+  // zero sized tensor case
+  if (x[0]->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    auto out_dims = out->dims();
+    out->Resize(out_dims);
+    return;
+  }
+
   dev_ctx.template Alloc<T>(out);
   auto& dim = x[0]->dims();
   std::vector<int> xdims;
diff --git a/test/collective/fleet/dygraph_group_sharded_stage1_fp16.py b/test/collective/fleet/dygraph_group_sharded_stage1_fp16.py
@@ -148,7 +148,7 @@ def train_mlp(
     )
 
     if sharding_stage == 1:
-        model.to(device="gpu")
+        model.to(device="xpu" if paddle.core.is_compiled_with_xpu() else "gpu")
 
     if not use_pure_fp16:
         for param in model.parameters():
diff --git a/test/collective/fleet/dygraph_group_sharded_stage2_bf16.py b/test/collective/fleet/dygraph_group_sharded_stage2_bf16.py
@@ -51,7 +51,10 @@ def train_mlp(
     )
 
     if sharding_stage != "dp":
-        group = paddle.distributed.new_group([0, 1], backend="nccl")
+        group = paddle.distributed.new_group(
+            [0, 1],
+            backend="bkcl" if paddle.core.is_compiled_with_xpu() else "nccl",
+        )
     scaler = None
     if test_scaler:
         assert sharding_stage == 2
@@ -89,7 +92,7 @@ def train_mlp(
         model = paddle.DataParallel(model)
 
     if sharding_stage == 2:
-        model.to(device="gpu")
+        model.to(device="xpu" if paddle.core.is_compiled_with_xpu() else "gpu")
 
     if not use_pure_bf16:
         for param in model.parameters():
@@ -192,7 +195,10 @@ def _compare_bf16_o1_vs_o2(acc_steps=1):
             use_pure_bf16=False,
             acc_steps=acc_steps,
         )
-        paddle.device.cuda.synchronize()
+        if paddle.core.is_compiled_with_xpu():
+            paddle.device.xpu.synchronize()
+        else:
+            paddle.device.cuda.synchronize()
         o2_losses, model_param_dict_o2, optimizer_state_dict_o2 = train_mlp(
             mlp2,
             sharding_stage=2,
@@ -201,7 +207,10 @@ def _compare_bf16_o1_vs_o2(acc_steps=1):
             acc_steps=acc_steps,
             use_main_grad=True,
         )
-        paddle.device.cuda.synchronize()
+        if paddle.core.is_compiled_with_xpu():
+            paddle.device.xpu.synchronize()
+        else:
+            paddle.device.cuda.synchronize()
         np.testing.assert_array_equal(o1_losses, o2_losses)
         # compare_state_dict(model_param_dict_o1, model_param_dict_o2, optimizer_state_dict_o2)
         return o1_losses, o2_losses
diff --git a/test/collective/fleet/dygraph_group_sharded_stage3_bf16.py b/test/collective/fleet/dygraph_group_sharded_stage3_bf16.py
@@ -47,7 +47,10 @@ def train_mlp(
     )
 
     if sharding_stage != "dp":
-        group = paddle.distributed.new_group([0, 1], backend="nccl")
+        group = paddle.distributed.new_group(
+            [0, 1],
+            backend="bkcl" if paddle.core.is_compiled_with_xpu() else "nccl",
+        )
     scaler = None
     if test_scaler:
         assert sharding_stage == 3
@@ -70,7 +73,7 @@ def train_mlp(
         ]
 
     if sharding_stage == 3:
-        model.to(device="gpu")
+        model.to(device="xpu" if paddle.core.is_compiled_with_xpu() else "gpu")
 
     if not use_pure_bf16:
         for param in model.parameters():
diff --git a/test/indexing/test_getitem_appendix.py b/test/indexing/test_getitem_appendix.py
@@ -244,10 +244,6 @@ def test_indexing(self):
         self.accuracy_check(x[True], [42])
 
 
-@unittest.skipIf(
-    paddle.core.is_compiled_with_xpu(),
-    "There are some bugs on XPU.",
-)
 class TestOSizeTensorIndexing(unittest.TestCase):
     def accuracy_check(self, paddle_t, numpy_array):
         np.testing.assert_allclose(paddle_t, numpy_array)
diff --git a/test/indexing/test_setitem_appendix.py b/test/indexing/test_setitem_appendix.py
@@ -249,10 +249,6 @@ def test_indexing(self):
         self.accuracy_check(x, 99)
 
 
-@unittest.skipIf(
-    paddle.core.is_compiled_with_xpu(),
-    "There are some bugs on XPU.",
-)
 class TestOSizeTensorIndexing(unittest.TestCase):
     def accuracy_check(self, paddle_t, numpy_array):
         np.testing.assert_allclose(paddle_t, numpy_array)
diff --git a/tools/xpu/disable_ut_xpu_kl3.local b/tools/xpu/disable_ut_xpu_kl3.local
@@ -214,12 +214,7 @@ test_auto_parallel_data_unshard_deprecated
 test_auto_parallel_parallelizer_deprecated
 test_auto_parallel_save_load_deprecated
 test_dygraph_group_sharded_api_for_eager
-test_dygraph_sharding_stage1_fp16
-test_dygraph_sharding_stage2_bf16
-test_dygraph_sharding_stage3_bf16
-test_parallel_dygraph_no_sync_gradient_check
 test_parallel_dygraph_pipeline_parallel_sync_send
-test_parallel_dygraph_sep_parallel
 test_parallel_dygraph_sharding_parallel
 test_collective_allgather_xpu
 test_collective_allreduce_xpu
diff --git a/tools/xpu/multi_card_ut_xpu_kl3.local b/tools/xpu/multi_card_ut_xpu_kl3.local
@@ -21,3 +21,8 @@ test_collective_reduce_xpu
 test_collective_softmax_with_cross_entropy_xpu
 test_llama_1d_xpu
 test_dist_base_xpu
+test_dygraph_sharding_stage1_fp16
+test_dygraph_sharding_stage2_bf16
+test_dygraph_sharding_stage3_bf16
+test_parallel_dygraph_no_sync_gradient_check
+test_parallel_dygraph_sep_parallel

Original file line number	Diff line number	Diff line change
`@@ -148,7 +148,7 @@ def train_mlp(`
`148`	`148`	`)`
`149`	`149`
`150`	`150`	`if sharding_stage == 1:`
`151`		`- model.to(device="gpu")`
	`151`	`+ model.to(device="xpu" if paddle.core.is_compiled_with_xpu() else "gpu")`
`152`	`152`
`153`	`153`	`if not use_pure_fp16:`
`154`	`154`	`for param in model.parameters():`