[auto parallal] add spmd rule for cast (#59050)

FeixLiu · web-flow · commit df81ec6b6b09 · 2023-11-17T16:55:32.000+08:00
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -228,6 +228,7 @@
   output : Tensor(out)
   infer_meta :
     func : CastInferMeta
+    spmd_rule : CastInferSpmd
   kernel :
     func : cast
     param : [x, dtype]
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h
@@ -76,7 +76,7 @@ CommContext* CreateOrGetCommContext(const DeviceContext& dev_ctx,
   do {                                                                \
     if (phi::CPUContext::classof(dev_ctx)) {                          \
       VLOG(4) << "Call `" << #fn_name << "` in Resharding on CPU.";   \
-      PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES(                  \
+      PD_VISIT_FLOATING_AND_INTEGRAL_TYPES(                           \
           dtype, #fn_name, ([&] {                                     \
             fn_name<data_t>(static_cast<const CPUContext&>(*dev_ctx), \
                             __VA_ARGS__);                             \
diff --git a/paddle/phi/core/visit_type.h b/paddle/phi/core/visit_type.h
@@ -150,6 +150,7 @@ namespace phi {
 
 ///////// BOOL and Floating and Integral Dispatch Marco ///////////
 
+#if NCCL_VERSION_CODE >= 21000
 #define PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES(TYPE, NAME, ...)        \
   [&] {                                                                       \
     const auto& __dtype__ = TYPE;                                             \
@@ -168,12 +169,44 @@ namespace phi {
           NAME, ::paddle::DataType::UINT8, uint8_t, __VA_ARGS__)              \
       PD_PRIVATE_CASE_TYPE(                                                   \
           NAME, ::paddle::DataType::INT16, int16_t, __VA_ARGS__)              \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::FLOAT16, float16, __VA_ARGS__)            \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::BFLOAT16, bfloat16, __VA_ARGS__)          \
       default:                                                                \
         PD_THROW("function " #NAME " is not implemented for data type `",     \
                  __dtype__,                                                   \
                  "`");                                                        \
     }                                                                         \
   }()
+#else
+#define PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES(TYPE, NAME, ...)        \
+  [&] {                                                                       \
+    const auto& __dtype__ = TYPE;                                             \
+    switch (__dtype__) {                                                      \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::BOOL, bool, __VA_ARGS__)    \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::FLOAT32, float, __VA_ARGS__)              \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::FLOAT64, double, __VA_ARGS__)             \
+      PD_PRIVATE_CASE_TYPE(NAME, ::paddle::DataType::INT32, int, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::INT64, int64_t, __VA_ARGS__)              \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::INT8, int8_t, __VA_ARGS__)                \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::UINT8, uint8_t, __VA_ARGS__)              \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::INT16, int16_t, __VA_ARGS__)              \
+      PD_PRIVATE_CASE_TYPE(                                                   \
+          NAME, ::paddle::DataType::FLOAT16, float16, __VA_ARGS__)            \
+      default:                                                                \
+        PD_THROW("function " #NAME " is not implemented for data type `",     \
+                 __dtype__,                                                   \
+                 "`");                                                        \
+    }                                                                         \
+  }()
+#endif
 
 ///////// Floating and Complex Dispatch Marco ///////////
 
diff --git a/paddle/phi/infermeta/spmd_rules/cast.cc b/paddle/phi/infermeta/spmd_rules/cast.cc
@@ -0,0 +1,21 @@
+/* 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/infermeta/spmd_rules/cast.h"
+#include "paddle/phi/infermeta/spmd_rules/elementwise.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo CastInferSpmd(const DistMetaTensor& x, phi::DataType dtype) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/cast.h b/paddle/phi/infermeta/spmd_rules/cast.h
@@ -0,0 +1,26 @@
+// 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 <vector>
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo CastInferSpmd(const DistMetaTensor& x, phi::DataType dtype);
+}
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.cc b/paddle/phi/infermeta/spmd_rules/elementwise.cc
@@ -113,15 +113,15 @@ SpmdInfo ElementwiseUnaryInferSpmd(const DistMetaTensor& x) {
   // input dist_attr.
   TensorDistAttr out_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
   out_dist_attr.set_dims_mapping(out_dims_mapping);
+  TensorDistAttr x_dst_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  x_dst_dist_attr.set_dims_mapping(out_dims_mapping);
 
-  // Step3: Handle partial
-  // Handle input tensor partial (TODO)
   VLOG(4) << "ElementwiseSPMDRule InferForward:";
   VLOG(4) << "Input0 shape: [" << str_join(x_shape) << "] "
           << "src_dims_mapping: [" << str_join(x_dims_mapping) << "] ";
   VLOG(4) << "Output dims_mapping: [" + str_join(out_dims_mapping) + "]\n\n";
 
-  return {{x_dist_attr_src}, {out_dist_attr}};
+  return {{x_dst_dist_attr}, {out_dist_attr}};
 }
 
 SpmdInfo ElementwiseUnaryInferSpmdReverse(const DistMetaTensor& x,
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -16,6 +16,7 @@ limitations under the License. */
 
 #include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
 
+#include "paddle/phi/infermeta/spmd_rules/cast.h"
 #include "paddle/phi/infermeta/spmd_rules/concat.h"
 #include "paddle/phi/infermeta/spmd_rules/default_data_parallel.h"
 #include "paddle/phi/infermeta/spmd_rules/elementwise.h"
diff --git a/test/auto_parallel/semi_auto_parallel_for_cast.py b/test/auto_parallel/semi_auto_parallel_for_cast.py
@@ -0,0 +1,55 @@
+# 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.
+
+
+from semi_auto_parallel_util import SemiAutoParallelTestBase
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestCastApiForSemiAutoParallel(SemiAutoParallelTestBase):
+    def __init__(self):
+        super().__init__()
+
+    def check_specs_unchanged(self, input, output):
+        input_dist_attr = input.dist_attr
+        output_dist_attr = output.dist_attr
+        assert input_dist_attr.dims_mapping == output_dist_attr.dims_mapping
+
+    def test_cast_shard(self):
+        x_shape = [16, 32]
+        x_specs = ['x', None]
+        inputs, outputs = self.runfunc_and_check(
+            x_shape,
+            x_specs,
+            op_func=paddle.cast,
+            with_backward=True,
+            dtype="float64",
+        )
+        self.check_specs_unchanged(inputs, outputs)
+
+    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_cast_shard()
+
+
+if __name__ == '__main__':
+    TestCastApiForSemiAutoParallel().run_test_case()
diff --git a/test/auto_parallel/semi_auto_parallel_simple_net_amp.py b/test/auto_parallel/semi_auto_parallel_simple_net_amp.py
@@ -150,6 +150,8 @@ def test_mp_demo_net(self):
             self.check_tensor_eq(param.grad, param_base.grad)
 
     def run_test_case(self):
+        if self._dtype == "bfloat16" and not paddle.amp.is_bfloat16_supported():
+            return
         self.test_dp_demo_net()
         self.test_mp_demo_net()
 
diff --git a/test/auto_parallel/test_semi_auto_parallel_basic.py b/test/auto_parallel/test_semi_auto_parallel_basic.py
@@ -27,8 +27,10 @@ def setUp(self):
         self._changeable_envs = {"backend": ["cpu", "gpu"]}
 
     def test_matmul_api(self):
+        default_envs = self._default_envs
+        default_envs["NVIDIA_TF32_OVERRIDE"] = "0"
         envs_list = test_base.gen_product_envs_list(
-            self._default_envs, self._changeable_envs
+            default_envs, self._changeable_envs
         )
         for envs in envs_list:
             self.run_test_case(
@@ -96,6 +98,16 @@ def test_add_n_api(self):
                 user_defined_envs=envs,
             )
 
+    def test_cast_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_cast.py",
+                user_defined_envs=envs,
+            )
+
     def test_custom_relu_api(self):
         envs_list = test_base.gen_product_envs_list(
             self._default_envs, self._changeable_envs
