[auto parallel] Support shard to partial reshard function. (#59292)

FeixLiu · web-flow · commit e7ac9cd0fd31 · 2023-11-27T17:54:45.000+08:00
diff --git a/paddle/fluid/pybind/auto_parallel_py.cc b/paddle/fluid/pybind/auto_parallel_py.cc
@@ -47,6 +47,7 @@
 #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/s_to_p_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
@@ -215,6 +216,10 @@ void BindAutoParallel(py::module *m) {
       *m, "SToSReshardFunction", ReshardFunction)
       .def(py::init<>());
 
+  py::class_<phi::distributed::SToPReshardFunction>(
+      *m, "SToPReshardFunction", ReshardFunction)
+      .def(py::init<>());
+
   py::class_<phi::distributed::PToSReshardFunction>(
       *m, "PToSReshardFunction", ReshardFunction)
       .def(py::init<>());
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/CMakeLists.txt b/paddle/phi/core/distributed/auto_parallel/reshard/CMakeLists.txt
@@ -9,6 +9,7 @@ collect_srcs(
   p_to_r_reshard_function.cc
   s_to_s_reshard_function.cc
   p_to_s_reshard_function.cc
+  s_to_p_reshard_function.cc
   nd_mesh_reshard_function.cc
   same_status_reshard_function.cc
   reshard_function_registry.cc)
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.cc b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.cc
@@ -20,6 +20,7 @@
 #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/s_to_p_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
@@ -45,7 +46,13 @@ std::vector<std::unique_ptr<ReshardFunction>>& GetReshardFunctionList() {
   return func_list;
 }
 
+// NOTE: be aware of the registration order of the reshard function.
+// Higher priority will be granted to the reshard function
+// which was registered earlier.
+// Reshard function with higher priority will be evoked
+// when more than one reshard function satisfy the request.
 REGISTER_RESHARD_FUNC(SToRReshardFunction);
+REGISTER_RESHARD_FUNC(SToPReshardFunction);
 REGISTER_RESHARD_FUNC(SToRReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(RToSReshardFunction);
 REGISTER_RESHARD_FUNC(RToSReshardFunctionCrossMesh);
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.cc
@@ -0,0 +1,67 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_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/s_to_r_reshard_function.h"
+
+#include "glog/logging.h"
+#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/reshard_utils.h"
+#include "paddle/phi/kernels/reduce_scatter_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi {
+namespace distributed {
+
+bool SToPReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_partial());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void SToPReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call SToPReshardFunction Eval";
+
+  // step 1, create tmp dist attr and tmp dist tensor
+  TensorDistAttr tmp_attr(out_dist_attr);
+  DistTensor tmp_tensor;
+  tmp_attr.clean_partial_status();
+
+  // step 2, do s to r reshard on `in` to `tmp`
+  SToRReshardFunction s_to_r;
+  s_to_r.Eval(dev_ctx, in, tmp_attr, &tmp_tensor);
+
+  // step 3, do r to p reshard on `tmp` to `out`
+  RToPReshardFunction r_to_p;
+  r_to_p.Eval(dev_ctx, tmp_tensor, out_dist_attr, out);
+}
+
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.h
@@ -0,0 +1,35 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class SToPReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToPReshard"; }
+};
+
+}  // namespace distributed
+}  // namespace phi
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -96,6 +96,9 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
   py_test_modules(test_reshard_s_to_r MODULES test_reshard_s_to_r)
   set_tests_properties(test_reshard_s_to_r
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_reshard_s_to_p MODULES test_reshard_s_to_p)
+  set_tests_properties(test_reshard_s_to_p
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
   py_test_modules(test_reshard_p_to_s MODULES test_reshard_p_to_s)
   set_tests_properties(test_reshard_p_to_s
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
diff --git a/test/auto_parallel/reshard_s_to_p.py b/test/auto_parallel/reshard_s_to_p.py
@@ -0,0 +1,59 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestReshardSToR:
+    def __init__(self):
+        self._shape = eval(os.getenv("shape"))
+        self._dtype = os.getenv("dtype")
+        self._seeds = eval(os.getenv("seeds"))
+        self._backend = os.getenv("backend")
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+
+        a = paddle.ones(self._shape)
+
+        in_placement = [dist.Shard(0)]
+        input_tensor = dist.shard_tensor(
+            a, mesh=self._mesh, placements=in_placement
+        )
+        assert input_tensor._local_shape[0] == self._shape[0] // 2
+
+        out = dist.reshard(
+            input_tensor,
+            mesh=self._mesh,
+            placements=[dist.Partial(dist.ReduceType.kRedSum)],
+        )
+
+        if dist.get_rank() == 0:
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+        else:
+            zeros = paddle.zeros(self._shape)
+            np.testing.assert_equal(out._local_value().numpy(), zeros.numpy())
+        assert np.equal(out.shape, input_tensor.shape).all()
+        assert np.equal(out._local_shape, input_tensor.shape).all()
+
+
+if __name__ == '__main__':
+    TestReshardSToR().run_test_case()
diff --git a/test/auto_parallel/semi_auto_parallel_shard_optimizer.py b/test/auto_parallel/semi_auto_parallel_shard_optimizer.py
@@ -148,8 +148,7 @@ def run_test_case(self):
         if self._backend == "gpu":
             self.test_adamw_mp()
             self.test_adamw_shard_optimizer(stage1=True)
-            # A problem has to be addressed if not shard batch.
-            # self.test_adamw_shard_optimizer(stage1=False)
+            self.test_adamw_shard_optimizer(stage1=False)
 
 
 if __name__ == '__main__':
diff --git a/test/auto_parallel/test_reshard_s_to_p.py b/test/auto_parallel/test_reshard_s_to_p.py
@@ -0,0 +1,44 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestReshardSToP(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=120)
+        self._default_envs = {
+            "shape": "(10, 20)",
+            "dtype": "float32",
+            "seeds": "1234",
+        }
+        self._changeable_envs = {
+            "backend": ["cpu", "gpu"],
+        }
+
+    def test_reshard_s_to_r(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "reshard_s_to_p.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
