[auto parallel] add softmax backward spmd rule (#59039)

cxxly · web-flow · commit d86f6866464f · 2023-11-30T10:27:36.000+08:00
* [auto parallel] add softmax backward spmd rule

* update test to new eager parallel api
diff --git a/paddle/phi/api/yaml/legacy_backward.yaml b/paddle/phi/api/yaml/legacy_backward.yaml
@@ -660,6 +660,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [out]
+    spmd_rule : SoftmaxGradInferSpmd
   kernel :
     func : softmax_grad
   composite : softmax_grad(out, out_grad, axis, x_grad)
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -1003,6 +1003,7 @@
   output : Tensor(out)
   infer_meta :
     func : SoftmaxInferMeta
+    spmd_rule : SoftmaxInferSpmd
   kernel :
     func : softmax
   inplace : (x -> out)
diff --git a/paddle/phi/infermeta/spmd_rules/reduction.cc b/paddle/phi/infermeta/spmd_rules/reduction.cc
@@ -112,7 +112,6 @@ SpmdInfo ReductionInferSpmdBase(const DistMetaTensor& x,
       ResoluteOutputPartialDimension(axis_to_dim_map, out_axes);
   out_dist_attr.set_partial_status(partial_on_dims,
                                    static_cast<ReduceType>(reduce_type));
-
   // Step3.2  handle input tensor partial (TODO)
   // If the op is a linear op, i.e. `linearity` is true, it supports
   // the input to be partial. Otherwise, the input cannot be partial
diff --git a/paddle/phi/infermeta/spmd_rules/softmax.cc b/paddle/phi/infermeta/spmd_rules/softmax.cc
@@ -19,7 +19,9 @@ limitations under the License. */
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/utils.h"
+#include "paddle/phi/infermeta/spmd_rules/rules.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace phi {
 namespace distributed {
@@ -150,5 +152,58 @@ SpmdInfo SoftmaxInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr}, {out_dist_attr_dst}};
 }
 
+SpmdInfo SoftmaxGradInferSpmd(const DistMetaTensor& out,
+                              const DistMetaTensor& out_grad,
+                              int axis) {
+  axis = axis < 0 ? out.dims().size() + axis : axis;
+
+  PADDLE_ENFORCE_EQ(out_grad.dims().size(),
+                    out_grad.dist_attr().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.dims().size(),
+                        out_grad.dist_attr().dims_mapping().size()));
+
+  PADDLE_ENFORCE_GE(
+      out_grad.dist_attr().dims_mapping().size(),
+      axis,
+      phi::errors::InvalidArgument("The Tensor out_grad's rank [%d] must be "
+                                   "greater than axis [%d].",
+                                   out_grad.dist_attr().dims_mapping().size(),
+                                   axis));
+
+  // To keeping consistent with forward propagation, sharding on softmax_axis
+  // is not supported now, the axis should be resharded as replicated.
+  auto out_grad_dims_mapping = out_grad.dist_attr().dims_mapping();
+  if (out_grad_dims_mapping[axis] >= 0) {
+    out_grad_dims_mapping[axis] = -1;
+    VLOG(6) << "SoftmaxGradInferSpmd: The out_grad's softmax_axis is reshard "
+               "to be replicated: "
+            << "original dims_mapping["
+            << str_join(out_grad.dist_attr().dims_mapping()) << "], "
+            << "resharded dims_mapping[" << str_join(out_grad_dims_mapping)
+            << "].";
+  }
+  auto out_dims_mapping = out.dist_attr().dims_mapping();
+  if (out_dims_mapping[axis] >= 0) {
+    out_dims_mapping[axis] = -1;
+    VLOG(6) << "SoftmaxGradInferSpmd: The out's softmax_axis is reshard "
+               "to be replicated: "
+            << "original dims_mapping["
+            << str_join(out.dist_attr().dims_mapping()) << "], "
+            << "resharded dims_mapping[" << str_join(out_dims_mapping) << "].";
+  }
+
+  auto out_dist_attr = CopyTensorDistAttrForOutput(out.dist_attr());
+  out_dist_attr.set_dims_mapping(out_dims_mapping);
+  auto out_grad_dist_attr = CopyTensorDistAttrForOutput(out_grad.dist_attr());
+  out_grad_dist_attr.set_dims_mapping(out_grad_dims_mapping);
+
+  return ElementwiseBinaryInferSpmd(
+      DistMetaTensor(out.dims(), out_dist_attr),
+      DistMetaTensor(out_grad.dims(), out_grad_dist_attr));
+}
+
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/softmax.h b/paddle/phi/infermeta/spmd_rules/softmax.h
@@ -31,5 +31,8 @@ SpmdInfo SoftmaxInferSpmdReverse(const DistMetaTensor& x,
                                  const DistMetaTensor& out,
                                  int axis);
 
+SpmdInfo SoftmaxGradInferSpmd(const DistMetaTensor& out,
+                              const DistMetaTensor& out_grad,
+                              int axis);
 }  // namespace distributed
 }  // namespace phi
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -126,6 +126,16 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
                   test_semi_auto_parallel_basic)
   set_tests_properties(test_semi_auto_parallel_basic
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 300)
+
+  py_test_modules(test_semi_auto_parallel_softmax_basic MODULES
+                  test_semi_auto_parallel_softmax_basic)
+  set_tests_properties(test_semi_auto_parallel_softmax_basic
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+
+  py_test_modules(test_semi_auto_parallel_compare_basic MODULES
+                  test_semi_auto_parallel_compare_basic)
+  set_tests_properties(test_semi_auto_parallel_compare_basic
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
   py_test_modules(test_semi_auto_parallel_pylayer MODULES
                   test_semi_auto_parallel_pylayer)
   set_tests_properties(test_semi_auto_parallel_pylayer
diff --git a/test/auto_parallel/semi_auto_parallel_for_softmax.py b/test/auto_parallel/semi_auto_parallel_for_softmax.py
@@ -0,0 +1,119 @@
+# 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 TestSoftmaxApiForSemiAutoParallel:
+    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"])
+
+        self._rtol = 1e-6
+        self._atol = 0
+
+        # The gradient of softmax is tiny, ref test_softmax_op.py, use atol
+        # to check the backward precision.
+        self._bwd_rtol = 0
+        self._bwd_atol = 1e-6
+
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+
+    def check_tensor_eq(self, a, b, rtol=1e-6, atol=0):
+        np1 = a.numpy()
+        np2 = b.numpy()
+        np.testing.assert_allclose(np1, np2, rtol=rtol, atol=atol, verbose=True)
+
+    def test_body(self, x_shape, out_shape, x_placements, func):
+        x = paddle.rand(x_shape, dtype=self._dtype)
+        x.stop_gradient = False
+
+        dist_x = dist.shard_tensor(x, self._mesh, x_placements)
+        dist_x.stop_gradient = False
+
+        dist_out = func(dist_x)
+        out = func(x)
+        self.check_tensor_eq(out, dist_out, self._rtol, self._atol)
+
+        dist_out.sum().backward()
+        out.sum().backward()
+        self.check_tensor_eq(
+            x.grad, dist_x.grad, self._bwd_rtol, self._bwd_atol
+        )
+
+    def test_softmax_shard(self):
+        self.test_body(
+            x_shape=[20, 30],
+            out_shape=[4, 4],
+            x_placements=[dist.Shard(0)],
+            func=lambda x: paddle.nn.functional.softmax(x, axis=1),
+        )
+
+    def test_softmax_shard_along_axis(self):
+        self.test_body(
+            x_shape=[20, 30],
+            out_shape=[20, 30],
+            x_placements=[dist.Shard(1)],
+            func=lambda x: paddle.nn.functional.softmax(x, axis=1),
+        )
+
+    def test_multi_axes(self):
+        self.test_body(
+            x_shape=[2, 4, 6, 10],
+            out_shape=[2, 4, 6, 10],
+            x_placements=[dist.Shard(0)],
+            func=lambda x: paddle.nn.functional.softmax(x, axis=1),
+        )
+
+    def test_multi_axes_along_axis(self):
+        self.test_body(
+            x_shape=[2, 4, 6, 10],
+            out_shape=[2, 4, 6, 10],
+            x_placements=[dist.Shard(0)],
+            func=lambda x: paddle.nn.functional.softmax(x, axis=0),
+        )
+
+    def test_negative_axis(self):
+        self.test_body(
+            x_shape=[2, 4, 6, 10],
+            out_shape=[2, 4, 6, 10],
+            x_placements=[dist.Shard(0)],
+            func=lambda x: paddle.nn.functional.softmax(x, axis=-4),
+        )
+
+    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_softmax_shard()
+        self.test_softmax_shard_along_axis()
+        self.test_multi_axes()
+        self.test_multi_axes_along_axis()
+        self.test_negative_axis()
+
+
+if __name__ == '__main__':
+    TestSoftmaxApiForSemiAutoParallel().run_test_case()
diff --git a/test/auto_parallel/test_semi_auto_parallel_compare_basic.py b/test/auto_parallel/test_semi_auto_parallel_compare_basic.py
@@ -0,0 +1,40 @@
+# 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 TestSemiAutoParallelBasic(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=120,
+        )
+        self._default_envs = {"dtype": "float32", "seed": "2023"}
+        self._changeable_envs = {"backend": ["cpu", "gpu"]}
+
+    def test_compare_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_compare.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/auto_parallel/test_semi_auto_parallel_softmax_basic.py b/test/auto_parallel/test_semi_auto_parallel_softmax_basic.py
@@ -0,0 +1,40 @@
+# 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 TestSemiAutoParallelBasic(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=120,
+        )
+        self._default_envs = {"dtype": "float32", "seed": "2023"}
+        self._changeable_envs = {"backend": ["cpu", "gpu"]}
+
+    def test_softmax_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_softmax.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/cpp/auto_parallel/CMakeLists.txt b/test/cpp/auto_parallel/CMakeLists.txt
@@ -8,11 +8,27 @@ cc_test(
   DEPS proto_desc)
 
 if(WITH_DISTRIBUTE)
+  cc_library(spmd_rule_test_util SRCS spmd_rule_test_util.cc)
   cc_test(
     dist_tensor_test
     SRCS dist_tensor_test.cc
     DEPS phi)
-  cc_test_old(spmd_rule_test SRCS spmd_rule_test.cc DEPS spmd_rules phi)
+  paddle_test(
+    spmd_rule_test
+    SRCS
+    spmd_rule_test.cc
+    DEPS
+    spmd_rule_test_util
+    spmd_rules
+    phi)
+  paddle_test(
+    softmax_grad_spmd_rule_test
+    SRCS
+    softmax_grad_spmd_rule_test.cc
+    DEPS
+    spmd_rule_test_util
+    spmd_rules
+    phi)
 endif()
 
 cc_test_old(dist_mapper_test SRCS dist_mapper_test.cc DEPS phi)
diff --git a/test/cpp/auto_parallel/softmax_grad_spmd_rule_test.cc b/test/cpp/auto_parallel/softmax_grad_spmd_rule_test.cc
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
diff --git a/test/cpp/auto_parallel/spmd_rule_test_util.cc b/test/cpp/auto_parallel/spmd_rule_test_util.cc
diff --git a/test/cpp/auto_parallel/spmd_rule_test_util.h b/test/cpp/auto_parallel/spmd_rule_test_util.h
