[AutoParallel] Add transpose spmd rules (#59236)

deepllz · web-flow · commit f9aa87b819d7 · 2023-11-23T05:43:12.000-06:00
diff --git a/paddle/phi/api/yaml/legacy_backward.yaml b/paddle/phi/api/yaml/legacy_backward.yaml
@@ -803,6 +803,7 @@
   infer_meta :
     func : TransposeGradInferMeta
     param : [out_grad, perm]
+    spmd_rule: TransposeGradInferSpmd
   kernel :
     func : transpose_grad
   backward : transpose_double_grad
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -1101,6 +1101,7 @@
   output : Tensor(out)
   infer_meta :
     func : TransposeInferMeta
+    spmd_rule: TransposeInferSpmd
   kernel :
     func : transpose
   inplace : (x -> out)
diff --git a/paddle/phi/infermeta/spmd_rules/transpose.cc b/paddle/phi/infermeta/spmd_rules/transpose.cc
@@ -24,31 +24,35 @@ namespace distributed {
 
 using phi::distributed::auto_parallel::str_join;
 
-////////////////// Utils Functions //////////////////
-std::string GetTransposeOutputNotation(int input_ndim,
-                                       const std::string& x_axes,
-                                       std::vector<int> perm_dims) {
-  // convert the negative dim value to normal dim value
-  for (int i = 0, n = perm_dims.size(); i < n; ++i) {
-    if (perm_dims[i] < 0) {
-      perm_dims[i] = input_ndim + perm_dims[i];
+void BuildEinsumNotation(const size_t x_ndim,
+                         std::vector<int> perm,
+                         std::string* p_x_axes,
+                         std::string* p_out_axes) {
+  std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
+  // get einsum notation for x
+  *p_x_axes = alphabet.substr(0, x_ndim);
+
+  // convert perm
+  for (size_t i = 0; i < x_ndim; i++) {
+    if (perm[i] < 0) {
+      perm[i] += x_ndim;
     }
   }
 
-  std::string out_axes = "";
-  for (int64_t i = 0; i < input_ndim; i++) {
-    out_axes.append(1, x_axes[perm_dims[i]]);
+  // get einsum notation for out
+  *p_out_axes = "";
+  for (size_t i = 0; i < x_ndim; i++) {
+    p_out_axes->append(1, p_x_axes->at(perm[i]));
   }
-
-  return out_axes;
 }
+
 ////////////////// InferMeta(Contains SPMD) Functions //////////////////
 SpmdInfo TransposeInferSpmd(const DistMetaTensor& x,
                             const std::vector<int>& perm) {
   // Step0: Verify input args based on transpose logic
-  auto x_shape = phi::vectorize(x.dims());
-  int x_ndim = x_shape.size();
-  auto x_dist_attr_src = x.dist_attr();
+  std::vector<int64_t> x_shape = phi::vectorize(x.dims());
+  size_t x_ndim = x_shape.size();
+  const TensorDistAttr& x_dist_attr_src = x.dist_attr();
   std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
   PADDLE_ENFORCE_EQ(
       x_ndim,
@@ -57,14 +61,19 @@ SpmdInfo TransposeInferSpmd(const DistMetaTensor& x,
                                    "dims_mapping size [%d] are not matched.",
                                    x_ndim,
                                    x_dims_mapping.size()));
+  // check perm size
+  PADDLE_ENFORCE_EQ(
+      x_ndim,
+      perm.size(),
+      phi::errors::InvalidArgument("The Tensor X's rank [%d] and "
+                                   "perm size [%d] are not matched.",
+                                   x_ndim,
+                                   perm.size()));
 
   // Step1: Build Einsum Notation
-  std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
-  // get einsum notation for input
-  std::string x_axes = alphabet.substr(0, x_ndim);
-
-  // get einsum notation for output
-  std::string out_axes = GetTransposeOutputNotation(x_ndim, x_axes, perm);
+  std::string x_axes;
+  std::string out_axes;
+  BuildEinsumNotation(x_ndim, perm, &x_axes, &out_axes);
 
   // Step2: Sharding Propogation
   // Step2.1: Merge input shardings
@@ -98,11 +107,11 @@ SpmdInfo TransposeInferSpmdReverse(const DistMetaTensor& x,
                                    const DistMetaTensor& out,
                                    const std::vector<int>& perm) {
   // Step0: Verify input args based on transpose logic
-  auto x_shape = phi::vectorize(x.dims());
-  auto out_shape = phi::vectorize(out.dims());
+  const std::vector<int64_t> x_shape = phi::vectorize(x.dims());
+  const std::vector<int64_t> out_shape = phi::vectorize(out.dims());
   int x_ndim = x_shape.size();
   int out_ndim = out_shape.size();
-  auto out_dist_attr_src = out.dist_attr();
+  TensorDistAttr out_dist_attr_src = out.dist_attr();
   std::vector<int64_t> out_dims_mapping = out_dist_attr_src.dims_mapping();
   PADDLE_ENFORCE_EQ(
       out_ndim,
@@ -111,14 +120,26 @@ SpmdInfo TransposeInferSpmdReverse(const DistMetaTensor& x,
                                    "dims_mapping size [%d] are not matched.",
                                    out_ndim,
                                    out_dims_mapping.size()));
+  PADDLE_ENFORCE_EQ(
+      x_ndim,
+      out_ndim,
+      phi::errors::InvalidArgument("The Tensor X's rank [%d] and "
+                                   "Out's rank [%d] are not matched.",
+                                   x_ndim,
+                                   out_ndim));
+  // check perm size
+  PADDLE_ENFORCE_EQ(
+      out_ndim,
+      perm.size(),
+      phi::errors::InvalidArgument("The Tensor Out's rank [%d] and "
+                                   "perm size [%d] are not matched.",
+                                   out_ndim,
+                                   perm.size()));
 
   // Step1: Build Einsum Notation
-  std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
-  // get einsum notation for input
-  std::string x_axes = alphabet.substr(0, x_ndim);
-
-  // get einsum notation for output
-  std::string out_axes = GetTransposeOutputNotation(x_ndim, x_axes, perm);
+  std::string x_axes;
+  std::string out_axes;
+  BuildEinsumNotation(x_ndim, perm, &x_axes, &out_axes);
 
   // Step2: Sharding Propogation
   // Step2.1: merge input shardings
@@ -148,5 +169,38 @@ SpmdInfo TransposeInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr}, {out_dist_attr_src}};
 }
 
+SpmdInfo TransposeGradInferSpmd(const DistMetaTensor& out_grad,
+                                const std::vector<int>& perm) {
+  const std::vector<int64_t> out_grad_shape = phi::vectorize(out_grad.dims());
+  size_t out_grad_ndim = out_grad_shape.size();
+  const std::vector<int64_t> out_grad_dims_mapping =
+      out_grad.dist_attr().dims_mapping();
+  size_t out_grad_dims_mapping_size = out_grad_dims_mapping.size();
+  PADDLE_ENFORCE_EQ(out_grad_ndim,
+                    out_grad_dims_mapping_size,
+                    phi::errors::InvalidArgument(
+                        "The Tensor Out_grad's rank [%d] and "
+                        "Out_grad's dims_mapping size [%d] are not matched.",
+                        out_grad_ndim,
+                        out_grad_dims_mapping_size));
+  size_t perm_size = perm.size();
+  PADDLE_ENFORCE_EQ(out_grad_ndim,
+                    perm_size,
+                    phi::errors::InvalidArgument(
+                        "The Tensor Out_grad's rank [%d] and perm size "
+                        "[%d] are not matched.",
+                        out_grad_ndim,
+                        perm_size));
+  std::vector<int64_t> x_dims_mapping(out_grad_ndim, -1);
+  for (size_t i = 0; i < perm.size(); ++i) {
+    int origin_index = perm[i] >= 0 ? perm[i] : out_grad_ndim + perm[i];
+    x_dims_mapping[origin_index] = out_grad_dims_mapping[i];
+  }
+  TensorDistAttr x_grad_dist_attr = out_grad.dist_attr();
+  x_grad_dist_attr.clean_partial_status();
+  x_grad_dist_attr.set_dims_mapping(x_dims_mapping);
+  return {{out_grad.dist_attr()}, {x_grad_dist_attr}};
+}
+
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/transpose.h b/paddle/phi/infermeta/spmd_rules/transpose.h
@@ -29,5 +29,8 @@ SpmdInfo TransposeInferSpmdReverse(const DistMetaTensor& x,
                                    const DistMetaTensor& out,
                                    const std::vector<int>& perm);
 
+SpmdInfo TransposeGradInferSpmd(const DistMetaTensor& out_grad,
+                                const std::vector<int>& perm);
+
 }  // namespace distributed
 }  // namespace phi
diff --git a/test/auto_parallel/semi_auto_parallel_for_transpose.py b/test/auto_parallel/semi_auto_parallel_for_transpose.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
+
+from semi_auto_parallel_util import SemiAutoParallelTestBase
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestTransposeApiForSemiAutoParallel(SemiAutoParallelTestBase):
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def check_dim_mapping(self, output, expected_dim_mapping):
+        assert (
+            output.dist_attr.dims_mapping == expected_dim_mapping
+        ), f"{output.dist_attr.dims_mapping}  vs {expected_dim_mapping}"
+
+    def test_transpose_shard(self):
+        x_shape = ([10, 6, 8],)
+        x_specs = ([None, 'x', None],)
+        _, output = self.runfunc_and_check(
+            x_shape,
+            x_specs,
+            op_func=paddle.transpose,
+            with_backward=True,
+            perm=[1, 2, -3],
+        )
+        self.check_dim_mapping(output, [0, -1, -1])
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.test_transpose_shard()
+
+
+if __name__ == '__main__':
+    TestTransposeApiForSemiAutoParallel().run_test_case()
diff --git a/test/auto_parallel/test_semi_auto_parallel_basic.py b/test/auto_parallel/test_semi_auto_parallel_basic.py
@@ -158,6 +158,16 @@ def test_triu_api(self):
                 user_defined_envs=envs,
             )
 
+    def test_transpose_api(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_for_transpose.py",
+                user_defined_envs=envs,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -1205,6 +1205,43 @@ TEST(Util, Ctor) {
   EXPECT_TRUE(!PlacementEqual(d, e));
 }
 
+TEST(Transpose, Ctor) {
+  std::vector<int64_t> mesh_shape = {2, 2};
+  std::vector<int64_t> process_ids = {0, 1, 2, 3};
+  std::vector<std::string> dim_names = {"x", "y"};
+  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
+
+  std::vector<int64_t> shape = {6, 8, 10};
+  std::vector<int64_t> dims_mapping = {0, -1, 1};
+
+  TensorDistAttr t_dist_attr = TensorDistAttr();
+  t_dist_attr.set_process_mesh(process_mesh);
+  t_dist_attr.set_dims_mapping(dims_mapping);
+  t_dist_attr.set_dynamic_dims({false, false, false});
+  phi::distributed::DistMetaTensor x =
+      phi::distributed::DistMetaTensor(phi::make_ddim(shape), t_dist_attr);
+  std::vector<int> perm = {1, 2, -3};
+  // test forward
+  phi::distributed::SpmdInfo forward_spmd_info =
+      phi::distributed::TransposeInferSpmd(x, perm);
+  EXPECT_EQ(forward_spmd_info.first.size(), static_cast<size_t>(1));
+  EXPECT_EQ(forward_spmd_info.second.size(), static_cast<size_t>(1));
+  check_dim_mapping(forward_spmd_info.first[0], {0, -1, 1});
+  check_dim_mapping(forward_spmd_info.second[0], {-1, 1, 0});
+  check_partial_dims(forward_spmd_info.second[0], {});
+  // test backward
+  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
+      phi::make_ddim({8, 10, 6}),
+      PADDLE_GET_CONST(TensorDistAttr, forward_spmd_info.second[0]));
+  phi::distributed::SpmdInfo backward_spmd_info =
+      TransposeGradInferSpmd(out_grad, perm);
+  EXPECT_EQ(backward_spmd_info.first.size(), static_cast<size_t>(1));
+  EXPECT_EQ(backward_spmd_info.second.size(), static_cast<size_t>(1));
+  check_dim_mapping(backward_spmd_info.first[0], {-1, 1, 0});
+  check_dim_mapping(backward_spmd_info.second[0], {0, -1, 1});
+  check_partial_dims(backward_spmd_info.second[0], {});
+}
+
 }  // namespace auto_parallel
 }  // namespace distributed
 }  // namespace paddle
