[Semi-Auto]Support p_to_s in nd_mesh_reshard (#59138)

ForFishes · web-flow · commit 164c63f517ad · 2023-11-21T16:38:14.000+08:00
* support p_to_s in nd_mesh

* fix typo in nd_mesh_reshard

* fix utest
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.cc
@@ -19,6 +19,7 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
@@ -106,7 +107,8 @@ void SameNdMeshReshardFunction::Eval(phi::DeviceContext* dev_ctx,
     const auto in_partial_status = in_dist_attr.partial_status();
     const auto& out_partial_status = out_dist_attr_orig.partial_status();
     for (const auto& kv : in_partial_status) {
-      if (out_partial_status.count(kv.first) != 0) {
+      if (out_partial_status.count(kv.first) != 0 ||
+          out_dist_attr_orig.is_shard(kv.first)) {
         continue;
       }
       VLOG(3) << "Step1: partial axis " << kv.first;
@@ -216,11 +218,15 @@ void SameNdMeshReshardFunction::Eval(phi::DeviceContext* dev_ctx,
     }
   }
 
-  // 4. Change replicated to shard
+  // 4. Change replicated/partial to shard
   for (int64_t i = first_diff_axis; i >= 0; --i) {
     int64_t out_mesh_axis = out_dist_attr_orig.dims_mapping()[i];
     if (out_mesh_axis != -1) {
-      VLOG(3) << "Step4: out_mesh axis " << out_mesh_axis;
+      const auto& in_partial_status = out->dist_attr().partial_status();
+      bool is_partial = in_partial_status.count(out_mesh_axis) != 0;
+
+      VLOG(3) << "Step4: out_mesh axis : " << out_mesh_axis
+              << "; paratial state :" << is_partial;
       // 4.1 Calculate the dist_attr after this transform
       TensorDistAttr real_out_dist_attr(out->dist_attr());
       std::vector<int64_t> real_dims_mapping =
@@ -246,9 +252,13 @@ void SameNdMeshReshardFunction::Eval(phi::DeviceContext* dev_ctx,
       // 4.5 Change from replicated to shard
       DistTensor tmp_result;
       SetDistProps(out, in_one_dim_dist_attr);
-      RToSReshardFunction func;
-      func.Eval(dev_ctx, *out, out_one_dim_dist_attr, &tmp_result);
-
+      if (is_partial) {
+        PToSReshardFunction func;
+        func.Eval(dev_ctx, *out, out_one_dim_dist_attr, &tmp_result);
+      } else {
+        RToSReshardFunction func;
+        func.Eval(dev_ctx, *out, out_one_dim_dist_attr, &tmp_result);
+      }
       // 4.6 Reset to the right dist attr
       SetValue(out, tmp_result.value());
       SetDistProps(out, real_out_dist_attr);
diff --git a/test/auto_parallel/reshard_nd_mesh.py b/test/auto_parallel/reshard_nd_mesh.py
@@ -197,6 +197,40 @@ def test_shard_to_shard(self, dev_ctx):
 
         assert np.equal(out.shape, input_tensor.shape).all()
 
+    def test_partial_replicate_to_shard_replicated(self, dev_ctx):
+        paddle.seed(self._seeds)
+        a = paddle.randn(self._shape).astype(self._dtype)
+        in_shard_specs = [None for i in range(len(self._shape))]
+        out_shard_specs = [None for i in range(len(self._shape))]
+
+        dist_attr = dist.DistAttr(
+            mesh=self._mesh, sharding_specs=in_shard_specs
+        )
+        dist_attr._set_partial_dims([0])
+
+        out_shard_specs[0] = "x"
+        out_dist_attr = dist.DistAttr(
+            mesh=self._mesh, sharding_specs=out_shard_specs
+        )
+
+        input_tensor = dist.shard_tensor(a, dist_attr=dist_attr)
+
+        reshard_func = core.SameNdMeshReshardFunction()
+        assert reshard_func.is_suitable(input_tensor, out_dist_attr)
+
+        out = reshard_func.eval(dev_ctx, input_tensor, out_dist_attr)
+
+        # check the value of input tensor
+        out_expected_local_tensor_list = paddle.split(
+            a, num_or_sections=self._mesh.shape[0], axis=0
+        )
+        index = dist.get_rank() % self._mesh.shape[0]
+        np.testing.assert_equal(
+            out._local_value().numpy(),
+            out_expected_local_tensor_list[index].numpy(),
+        )
+        assert np.equal(out.shape, input_tensor.shape).all()
+
     def run_test_case(self):
         if self._backend == "cpu":
             paddle.set_device("cpu")
@@ -211,6 +245,10 @@ def run_test_case(self):
         self.test_shard_partial_to_shard_replicated(dev_ctx)
         self.test_shard_partial_to_replicated(dev_ctx)
 
+        if self._backend == "gpu":
+            # reduce_scatter is not supported on CPU
+            self.test_partial_replicate_to_shard_replicated(dev_ctx)
+
 
 if __name__ == '__main__':
     TestReshardNdMesh().run_test_case()
