support unbalanced shard to replicate (#61432)

LiYuRio · web-flow · commit b4e13a66bb50 · 2024-02-05T10:52:02.000+08:00
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.cc
@@ -23,15 +23,71 @@
 #include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/all_gather_kernel.h"
 #include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/split_kernel.h"
 
 namespace phi {
 namespace distributed {
 
+namespace {
+
+void ReshardSToRWithPadding(DeviceContext* dev_ctx,
+                            int64_t split_axis,
+                            const std::vector<int64_t>& process_ids,
+                            const DenseTensor& in,
+                            int64_t num_of_padding,
+                            DenseTensor* out) {
+  int64_t num_of_process = process_ids.size();
+  auto dtype = in.dtype();
+
+  // For balanced split to replicate, we need to do all gather first.
+  // If the input value doesn't split on axis 0, we need to split
+  // and concat on specific axis.
+  RESHARD_FUNCTOR_WITH_COMM(
+      dev_ctx, AllGather, dtype, process_ids, in, num_of_process, out);
+
+  if (split_axis != 0 || num_of_padding != 0) {
+    IntArray sections(std::vector<int64_t>(num_of_process, in.dims()[0]));
+
+    std::vector<DenseTensor> split_out_vec;
+    RESHARD_FUNCTOR(dev_ctx,
+                    Split,
+                    dtype,
+                    *out,
+                    sections,
+                    /*split_axis*/ 0,
+                    &split_out_vec);
+
+    if (num_of_padding != 0) {
+      for (int64_t i = num_of_padding; i < num_of_process; ++i) {
+        std::vector<DenseTensor> tmp_out_vec;
+        IntArray tmp_sections(
+            std::vector<int64_t>{in.dims()[split_axis] - 1, 1});
+        RESHARD_FUNCTOR(dev_ctx,
+                        Split,
+                        dtype,
+                        split_out_vec[i],
+                        tmp_sections,
+                        split_axis,
+                        &tmp_out_vec);
+        split_out_vec[i] = tmp_out_vec[0];
+      }
+    }
+
+    // Concat the result after split on correct axis.
+    std::vector<const DenseTensor*> concat_input_vec;
+    for (const auto& tensor : split_out_vec) {
+      concat_input_vec.emplace_back(&tensor);
+    }
+    RESHARD_FUNCTOR(dev_ctx, Concat, dtype, concat_input_vec, split_axis, out);
+  }
+}
+
+}  // namespace
+
 bool SToRReshardFunction::IsSuitable(const DistTensor& in,
                                      const TensorDistAttr& out_dist_attr) {
   const auto& in_dist_attr = in.dist_attr();
-  const auto& in_dims_mapping = in_dist_attr.dims_mapping();
 
   RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
   RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated());
@@ -43,14 +99,6 @@ bool SToRReshardFunction::IsSuitable(const DistTensor& in,
   RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
   RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
 
-  // Ensure the tensor is balanced split, or we need send/recv rather than
-  // all_gather
-  int split_axis = GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->first;
-  int64_t num_of_process = in_process_mesh.size();
-  RESHARD_SHORTCUT_IF_FALSE(in.local_dims()[static_cast<int>(split_axis)] *
-                                num_of_process ==
-                            in.dims()[static_cast<int>(split_axis)]);
-
   return true;
 }
 
@@ -63,58 +111,62 @@ void SToRReshardFunction::Eval(DeviceContext* dev_ctx,
   const auto& in_dims_mapping = in_dist_attr.dims_mapping();
   const auto& in_process_mesh = in_dist_attr.process_mesh();
   const auto& in_process_ids = in_process_mesh.process_ids();
-  auto dtype = in.dtype();
 
-  // Since the precondition ensure the out_process_ids is equal to the
-  // in_process_ids, so the participate process ids mush equal to either
-  // in_process_ids or out_process_ids.
-  RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
-                            AllGather,
-                            dtype,
-                            in_process_ids,
-                            in.value(),
-                            in_process_ids.size(),
-                            GetMutableTensor(out));
   int split_axis = GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->first;
-
-  if (split_axis == 0) {
-    // If the input dist tensor is shard(0), the subsequent split
-    // and concat is unnecessary.
-    SetDistProps(out, in.dims(), out_dist_attr);
+  int64_t num_of_process = in_process_mesh.size();
+  int64_t num_of_padding = in.dims()[split_axis] % num_of_process;
+  bool is_balanced_split = (num_of_padding == 0);
+
+  if (is_balanced_split) {
+    VLOG(3) << "Balanced reshard from shard to replicated";
+    ReshardSToRWithPadding(dev_ctx,
+                           split_axis,
+                           in_process_ids,
+                           in.value(),
+                           num_of_padding,
+                           GetMutableTensor(out));
   } else {
-    // Since the result of all_gather always concat the tensor on axis 0,
-    // first we need to split the result on axis 0,
-    // then we need to concat the split result on input split axis.
-    int64_t default_split_axis = 0;
-    int64_t num_of_process = static_cast<int64_t>(in_process_ids.size());
-
-    IntArray sections(std::vector<int64_t>(
-        num_of_process,
-        in.value().dims()[static_cast<int>(default_split_axis)]));
-    std::vector<DenseTensor> split_out_vec;
-    RESHARD_FUNCTOR(dev_ctx,
-                    Split,
-                    dtype,
-                    out->value(),
-                    sections,
-                    default_split_axis,
-                    &split_out_vec);
-
-    // Concat the result after split on correct axis.
-    std::vector<const DenseTensor*> concat_input_vec;
-    for (const auto& tensor : split_out_vec) {
-      concat_input_vec.emplace_back(&tensor);
+    VLOG(3) << "Unbalanced reshard from shard to replicated";
+    bool need_padding =
+        (in.dims()[split_axis] / num_of_process == in.local_dims()[split_axis]);
+    if (need_padding) {
+      DDim concat_local_shape = in.local_dims();
+      concat_local_shape[split_axis] = 1;
+      IntArray concat_local_shape_int_array(concat_local_shape.Get(),
+                                            concat_local_shape.size());
+      auto dtype = in.dtype();
+
+      DenseTensor concat_local_tensor;
+      RESHARD_FUNCTOR(dev_ctx,
+                      Full,
+                      dtype,
+                      concat_local_shape_int_array,
+                      0,
+                      &concat_local_tensor);
+
+      DenseTensor in_local_tensor = in.value();
+      std::vector<const DenseTensor*> concat_input_vec = {&in_local_tensor,
+                                                          &concat_local_tensor};
+
+      DenseTensor concat_result;
+      RESHARD_FUNCTOR(
+          dev_ctx, Concat, dtype, concat_input_vec, split_axis, &concat_result);
+      ReshardSToRWithPadding(dev_ctx,
+                             split_axis,
+                             in_process_ids,
+                             concat_result,
+                             num_of_padding,
+                             GetMutableTensor(out));
+    } else {
+      ReshardSToRWithPadding(dev_ctx,
+                             split_axis,
+                             in_process_ids,
+                             in.value(),
+                             num_of_padding,
+                             GetMutableTensor(out));
     }
-
-    RESHARD_FUNCTOR(dev_ctx,
-                    Concat,
-                    dtype,
-                    concat_input_vec,
-                    split_axis,
-                    GetMutableTensor(out));
-
-    SetDistProps(out, in.dims(), out_dist_attr);
   }
+  SetDistProps(out, in.dims(), out_dist_attr);
 }
 
 bool SToRReshardFunctionCrossMesh::IsSuitable(
diff --git a/test/auto_parallel/reshard_s_to_r.py b/test/auto_parallel/reshard_s_to_r.py
@@ -47,6 +47,7 @@ def run_test_case(self):
 
         assert np.equal(out.shape, out._local_shape).all()
         assert np.equal(out.shape, input_tensor.shape).all()
+        np.testing.assert_equal(out.numpy(), a.numpy())
 
 
 if __name__ == '__main__':
diff --git a/test/auto_parallel/test_reshard_s_to_r.py b/test/auto_parallel/test_reshard_s_to_r.py
@@ -21,7 +21,7 @@ class TestReshardSToR(test_base.CommunicationTestDistBase):
     def setUp(self):
         super().setUp(num_of_devices=2, timeout=120)
         self._default_envs = {
-            "shape": "(10, 20)",
+            "shape": "(11, 20)",
             "dtype": "float32",
             "seeds": str(self._seeds),
         }

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ class TestReshardSToR(test_base.CommunicationTestDistBase):`
`21`	`21`	`def setUp(self):`
`22`	`22`	`super().setUp(num_of_devices=2, timeout=120)`
`23`	`23`	`self._default_envs = {`
`24`		`- "shape": "(10, 20)",`
	`24`	`+ "shape": "(11, 20)",`
`25`	`25`	`"dtype": "float32",`
`26`	`26`	`"seeds": str(self._seeds),`
`27`	`27`	`}`