[XPU] fix fleet unittests (#68542)

* [XPU] fix fleet unittests

* [XPU] fix fleet unittests

* refine: use new default parameter

* revert unnecessary modifications.

* revert unnecessary modifications.

* fix cmakelist

* revert unnecessary modifications.

* fix cmakelist for recompute ut.

Author: houj04

paddle/phi/api/lib/data_transform.cc

@@ -143,6 +143,33 @@ phi::DenseTensor CastDataType(const phi::GPUContext& dev_ctx,
 }
 #endif
 
+#ifdef PADDLE_WITH_XPU
+phi::DenseTensor CastDataType(const phi::XPUContext& dev_ctx,
+                              const phi::DenseTensor& tensor,
+                              DataType dtype) {
+  switch (tensor.dtype()) {
+    case DataType::FLOAT32:
+      return phi::Cast<float>(dev_ctx, tensor, dtype);
+    case DataType::FLOAT64:
+      return phi::Cast<double>(dev_ctx, tensor, dtype);
+    case DataType::INT32:
+      return phi::Cast<int32_t>(dev_ctx, tensor, dtype);
+    case DataType::INT64:
+      return phi::Cast<int64_t>(dev_ctx, tensor, dtype);
+    case DataType::FLOAT16:
+      return phi::Cast<phi::dtype::float16>(dev_ctx, tensor, dtype);
+    case DataType::BOOL:
+      return phi::Cast<bool>(dev_ctx, tensor, dtype);
+    case DataType::UINT8:
+      return phi::Cast<uint8_t>(dev_ctx, tensor, dtype);
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Data type (%s) is not supported when casting data type.",
+          tensor.dtype()));
+  }
+}
+#endif
+
 inline phi::DenseTensor TransDataType(const phi::DenseTensor& tensor,
                                       DataType dtype) {
   auto& pool = phi::DeviceContextPool::Instance();
@@ -161,6 +188,11 @@ inline phi::DenseTensor TransDataType(const phi::DenseTensor& tensor,
     auto* dev_ctx = static_cast<phi::GPUContext*>(pool.Get(tensor.place()));
     return CastDataType(*dev_ctx, tensor, dtype);
 #endif
+#ifdef PADDLE_WITH_XPU
+  } else if (tensor.place().GetType() == phi::AllocationType::XPU) {
+    auto* dev_ctx = static_cast<phi::XPUContext*>(pool.Get(tensor.place()));
+    return CastDataType(*dev_ctx, tensor, dtype);
+#endif
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   } else if (tensor.place().GetType() == phi::AllocationType::CUSTOM) {
     phi::DenseTensor out;

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_optimizer_stage2.py

@@ -74,7 +74,7 @@ def __init__(
         optim,
         group=None,
         offload=False,
-        device="gpu",
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
         pretrain_sync_models=True,
         dp_group=None,
         **kw,
@@ -590,6 +590,12 @@ def _step(self):
                             )
                             .cast(dtype=param.dtype)
                         )
+                    elif self._default_device == "xpu":
+                        param.set_value(
+                            self._master_params[param.name]
+                            .to("xpu:" + str(self.dev_id))
+                            .cast(dtype=param.dtype)
+                        )
                     else:
                         param.set_value(
                             self._master_params[param.name]

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage2.py

@@ -31,6 +31,7 @@
 from paddle import nn
 from paddle.distributed import collective
 from paddle.distributed.utils.log_utils import get_logger
+from paddle.framework import core
 
 from .group_sharded_optimizer_stage2 import GroupShardedOptimizerStage2
 from .group_sharded_storage import GradStorage
@@ -66,7 +67,7 @@ def __init__(
         sync_buffers=False,
         buffer_max_size=2**23,  # 8MB
         auto_refresh_trainable=True,
-        device="gpu",
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
         dp_group=None,
     ):
         super().__init__()

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage3.py

@@ -104,7 +104,7 @@ def __init__(
         optimizer,
         group=None,
         sync_buffers=False,
-        device="gpu",
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
         segment_size=2**20,
         pretrain_sync_models=True,
         offload=False,
@@ -310,7 +310,10 @@ def _clear_gradients(self):
                         paddle.CustomPlace(self._default_device, DEV_ID), True
                     )
                 else:
-                    tmp_var = param.cuda(DEV_ID)
+                    # both GPU and XPU
+                    tmp_var = param.to(
+                        self._default_device + ":" + (str)(DEV_ID)
+                    )
 
                 if (
                     tmp_var.dtype == Type.fp32.value
@@ -1197,7 +1200,8 @@ def _cpu2device(param):
     if DEV in paddle.device.get_all_custom_device_type():
         tmp_p = param.fw_storage._copy_to(paddle.CustomPlace(DEV, DEV_ID), True)
     else:
-        tmp_p = param.fw_storage.cuda(DEV_ID)
+        # both GPU and XPU
+        tmp_p = param.fw_storage.to(DEV + ":" + (str)(DEV_ID))
     if (
         tmp_p.dtype == Type.fp32.value
         and param2dtype[param.name] == Type.fp16.value

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_utils.py

@@ -167,6 +167,8 @@ def _dygraph_clip(self, params_grads):
                 global_norm_var = global_norm_var._copy_to(
                     paddle.CustomPlace(dev_type, dev_id), True
                 )
+            elif dev_type == "xpu":
+                global_norm_var = global_norm_var.to(self._device)
             else:
                 global_norm_var = global_norm_var.cuda(dev_id)
 

test/collective/fleet/CMakeLists.txt

@@ -61,8 +61,13 @@ if((WITH_ROCM) AND LOCAL_ALL_PLAT)
     "PADDLE_DIST_UT_PORT=21204;http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python"
   )
 endif()
-if(WITH_NCCL OR WITH_RCCL)
-  if((WITH_GPU OR WITH_ROCM) AND LOCAL_ALL_PLAT)
+if(WITH_NCCL
+   OR WITH_RCCL
+   OR WITH_XPU_BKCL)
+  if((WITH_GPU
+      OR WITH_ROCM
+      OR WITH_XPU)
+     AND LOCAL_ALL_PLAT)
     bash_test_modules(
       test_parallel_dygraph_mp_layers
       START_BASH
@@ -608,13 +613,19 @@ if((WITH_GPU OR WITH_ROCM) AND LOCAL_ALL_PLAT)
   set_tests_properties(test_imperative_auto_mixed_precision_for_eager
                        PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=DIST")
 endif()
-if((WITH_GPU OR WITH_ROCM) AND LOCAL_ALL_PLAT)
+if((WITH_GPU
+    OR WITH_ROCM
+    OR WITH_XPU)
+   AND LOCAL_ALL_PLAT)
   py_test_modules(
     test_dygraph_recompute_for_eager MODULES test_dygraph_recompute_for_eager
     ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
 endif()
-if((WITH_GPU OR WITH_ROCM) AND LOCAL_ALL_PLAT)
+if((WITH_GPU
+    OR WITH_ROCM
+    OR WITH_XPU)
+   AND LOCAL_ALL_PLAT)
   py_test_modules(
     test_dygraph_recompute MODULES test_dygraph_recompute ENVS
     "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")

test/collective/fleet/dygraph_group_sharded_stage2.py

@@ -99,7 +99,9 @@ def train_mlp(
     scale_fn_test=False,
 ):
     if sharding_stage != "dp":
-        group = paddle.distributed.new_group([0, 1], backend="nccl")
+        group = paddle.distributed.new_group(
+            [0, 1], backend="bkcl" if paddle.is_compiled_with_xpu() else "nccl"
+        )
     if opt_group:
         optimizer = optimizer_setting(
             model=model, use_pure_fp16=use_pure_fp16, opt_group=opt_group
@@ -149,7 +151,7 @@ def train_mlp(
     )
 
     if sharding_stage == 2:
-        model.to(device="gpu")
+        model.to(device="xpu" if paddle.is_compiled_with_xpu() else "gpu")
 
     for eop in range(epoch):
         model.train()
@@ -210,7 +212,10 @@ def test_dp_stage2():
     )
     for i in range(len(dp_params)):
         np.testing.assert_allclose(
-            dp_params[i].numpy(), stage2_params[i].numpy(), rtol=1e-6
+            dp_params[i].numpy(),
+            stage2_params[i].numpy(),
+            rtol=1e-6,
+            atol=1e-8 if paddle.is_compiled_with_xpu() else 0,
         )
 
     # stage2 accumulate grad
@@ -232,7 +237,10 @@ def test_dp_stage2():
     )
     for i in range(len(dp_params)):
         np.testing.assert_allclose(
-            dp_params[i].numpy(), stage2_params[i].numpy(), rtol=1e-6
+            dp_params[i].numpy(),
+            stage2_params[i].numpy(),
+            rtol=1e-6,
+            atol=1e-8 if paddle.is_compiled_with_xpu() else 0,
         )
 
     # save/load model

test/collective/fleet/dygraph_group_sharded_stage2_comm_overlap.py

@@ -98,7 +98,9 @@ def train_mlp(
     test_minimize=False,
 ):
     if sharding_stage != "dp":
-        group = paddle.distributed.new_group([0, 1], backend="nccl")
+        group = paddle.distributed.new_group(
+            [0, 1], backend="bkcl" if paddle.is_compiled_with_xpu() else "nccl"
+        )
     if opt_group:
         optimizer = optimizer_setting(
             model=model, use_pure_fp16=use_pure_fp16, opt_group=opt_group
@@ -140,7 +142,7 @@ def train_mlp(
     )
 
     if sharding_stage == 2:
-        model.to(device="gpu")
+        model.to(device="xpu" if paddle.is_compiled_with_xpu() else "gpu")
 
     for eop in range(epoch):
         model.train()
@@ -166,7 +168,10 @@ def train_mlp(
             optimizer.step()
             optimizer.clear_grad()
 
-    paddle.device.cuda.synchronize()
+    if paddle.is_compiled_with_xpu():
+        paddle.device.xpu.synchronize()
+    else:
+        paddle.device.cuda.synchronize()
 
     if save_model:
         return model, optimizer
@@ -201,7 +206,10 @@ def test_dp_stage2():
     )
     for i in range(len(dp_params)):
         np.testing.assert_allclose(
-            dp_params[i].numpy(), stage2_params[i].numpy(), rtol=1e-6
+            dp_params[i].numpy(),
+            stage2_params[i].numpy(),
+            rtol=1e-6,
+            atol=1e-8 if paddle.is_compiled_with_xpu() else 0,
         )
 
     # stage2 accumulate grad
@@ -223,7 +231,10 @@ def test_dp_stage2():
     )
     for i in range(len(dp_params)):
         np.testing.assert_allclose(
-            dp_params[i].numpy(), stage2_params[i].numpy(), rtol=1e-6
+            dp_params[i].numpy(),
+            stage2_params[i].numpy(),
+            rtol=1e-6,
+            atol=1e-8 if paddle.is_compiled_with_xpu() else 0,
         )
 
     # save/load model

test/collective/fleet/dygraph_group_sharded_stage2_offload.py

@@ -94,7 +94,10 @@ def train_mlp(model, offload=False, test=False):
 
     for dtype in optimizer.param_storages:
         for dst_rank, param_storage in optimizer.param_storages[dtype].items():
-            param_storage.to(device="gpu", dtype=dtype)
+            param_storage.to(
+                device="xpu" if paddle.is_compiled_with_xpu() else "gpu",
+                dtype=dtype,
+            )
 
     return model.parameters()
 

test/collective/fleet/dygraph_group_sharded_stage3.py

@@ -366,10 +366,9 @@ def test_stage2_stage3():
         )
 
     # bfp16
-    nccl_version = core.nccl_version()
-
     if (
-        nccl_version >= 21000
+        paddle.is_compiled_with_xpu()
+        or core.nccl_version() >= 21000
         and paddle.device.cuda.get_device_properties().major >= 8
     ):
         stage2_params = train_mlp(