【Hackathon 7th Fundable Projects 2 No.59-61】 [fluid_ops] mp_allreduce_sum (#68921)

co63oc · web-flow · commit 8886ac17b46d · 2024-10-30T10:20:48.000+08:00
* Fix

* Fix

* Fix

* Fix

* Fix
diff --git a/paddle/fluid/operators/collective/mp_allreduce_sum_op.cc b/paddle/fluid/operators/collective/mp_allreduce_sum_op.cc
@@ -40,10 +40,6 @@ class MpAllReduceSumOpMaker : public framework::OpProtoAndCheckerMaker {
     AddAttr<int>("ring_id", "(int default 0) communication ring id.")
         .SetDefault(0);
 
-    AddAttr<bool>(
-        "use_calc_stream",
-        "(bool default false) eject CUDA operations to calculation stream.")
-        .SetDefault(false);
     AddComment(string::Sprintf(R"DOC(
 MpAllReduceSum Operator
 
@@ -82,13 +78,3 @@ REGISTER_OPERATOR(mp_allreduce_sum,
                   ops::MpAllReduceSumOpGradMaker<paddle::imperative::OpBase>,
                   ops::MpAllReduceSumOpMaker,
                   ops::MpAllReduceSumInplaceInferer);
-
-PD_REGISTER_STRUCT_KERNEL(mp_allreduce_sum,
-                          CPU,
-                          ALL_LAYOUT,
-                          ops::MpAllReduceSumCPUKernel,
-                          float,
-                          double,
-                          int,
-                          int64_t,
-                          phi::dtype::float16) {}
diff --git a/paddle/fluid/operators/collective/mp_allreduce_sum_op.cu.cc b/paddle/fluid/operators/collective/mp_allreduce_sum_op.cu.cc
diff --git a/paddle/fluid/operators/ops_signature/mp_allreduce_sum_sig.cc b/paddle/fluid/operators/ops_signature/mp_allreduce_sum_sig.cc
@@ -0,0 +1,27 @@
+/* Copyright (c) 2024 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/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature MpAllreduceSumOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  return KernelSignature("mp_allreduce_sum", {"X"}, {}, {"Out"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(mp_allreduce_sum,
+                           phi::MpAllreduceSumOpArgumentMapping);
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -209,6 +209,8 @@
     'lod_reset',
     'lod_reset_',
     'max_pool2d_v2',
+    'mp_allreduce_sum',
+    'mp_allreduce_sum_',
     'partial_sum',
     'pull_gpups_sparse',
     'pull_gpups_sparse_',
diff --git a/paddle/phi/kernels/cpu/mp_allreduce_sum_kernel.cc b/paddle/phi/kernels/cpu/mp_allreduce_sum_kernel.cc
@@ -0,0 +1,35 @@
+// Copyright (c) 2024 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/kernel_registry.h"
+#include "paddle/phi/kernels/all_reduce_kernel.h"
+
+namespace phi {
+template <typename T, typename Context>
+void MpAllReduceSumKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          DenseTensor* out) {
+  AllReduceKernel<T, Context>(
+      dev_ctx, x, static_cast<int>(ReduceType::kRedSum), out);
+}
+}  // namespace phi
+PD_REGISTER_KERNEL(mp_allreduce_sum,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::MpAllReduceSumKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float16) {}
diff --git a/paddle/phi/kernels/gpu/mp_allreduce_sum_kernel.cu b/paddle/phi/kernels/gpu/mp_allreduce_sum_kernel.cu
@@ -0,0 +1,50 @@
+// Copyright (c) 2024 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/kernel_registry.h"
+#include "paddle/phi/kernels/all_reduce_kernel.h"
+
+namespace phi {
+template <typename T, typename Context>
+void MpAllReduceSumKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          DenseTensor* out) {
+  AllReduceKernel<T, Context>(
+      dev_ctx, x, static_cast<int>(ReduceType::kRedSum), out);
+}
+}  // namespace phi
+
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(mp_allreduce_sum,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::MpAllReduceSumKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::bfloat16,
+                   phi::dtype::float16) {}
+#else
+PD_REGISTER_KERNEL(mp_allreduce_sum,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::MpAllReduceSumKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::float16) {}
+#endif
diff --git a/paddle/phi/kernels/xpu/mp_allreduce_sum_kernel.cc b/paddle/phi/kernels/xpu/mp_allreduce_sum_kernel.cc
@@ -0,0 +1,34 @@
+// Copyright (c) 2024 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/kernel_registry.h"
+#include "paddle/phi/kernels/all_reduce_kernel.h"
+
+namespace phi {
+template <typename T, typename Context>
+void MpAllReduceSumKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          DenseTensor* out) {
+  AllReduceKernel<T, Context>(
+      dev_ctx, x, static_cast<int>(ReduceType::kRedSum), out);
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(mp_allreduce_sum,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::MpAllReduceSumKernel,
+                   float,
+                   int,
+                   phi::dtype::float16) {}
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2089,6 +2089,12 @@
   kernel :
     func : mode_grad
 
+- backward_op : mp_allreduce_sum_grad
+  forward : mp_allreduce_sum(Tensor x, int ring_id = 0) -> Tensor(out)
+  args : (Tensor out_grad, int ring_id = 0)
+  output : Tensor(x_grad)
+  invoke : c_identity(out_grad, ring_id, false, false)
+
 - backward_op : multi_dot_grad
   forward : multi_dot (Tensor[] x) -> Tensor(out)
   args : (Tensor[] x, Tensor out_grad)
diff --git a/paddle/phi/ops/yaml/legacy/backward_exclude.yaml b/paddle/phi/ops/yaml/legacy/backward_exclude.yaml
@@ -25,6 +25,7 @@
 - mean_double_grad
 - mean_grad
 - mish_grad
+- mp_allreduce_sum_grad
 - norm_grad
 - pad_double_grad
 - pad_grad
diff --git a/paddle/phi/ops/yaml/legacy/ops_exclude.yaml b/paddle/phi/ops/yaml/legacy/ops_exclude.yaml
@@ -62,6 +62,7 @@
 - memcpy_d2h
 - memcpy_h2d
 - mish
+- mp_allreduce_sum
 - norm
 - one_hot
 - ones
diff --git a/paddle/phi/ops/yaml/op_compat.yaml b/paddle/phi/ops/yaml/op_compat.yaml
@@ -4482,6 +4482,12 @@
   outputs :
     out : Out
 
+- op: mp_allreduce_sum
+  inputs :
+    x : X
+  outputs :
+    out: Out
+
 - op: nce
   backward: nce_grad
   inputs:
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3482,6 +3482,18 @@
   inplace : (param -> param_out), (velocity -> velocity_out), (master_param -> master_param_out)
   traits : pir::SideEffectTrait, paddle::dialect::ForwardOnlyTrait
 
+- op : mp_allreduce_sum
+  args : (Tensor x, int ring_id = 0)
+  output : Tensor(out)
+  infer_meta :
+    func : AllReduceInferMeta
+    param: [x]
+  kernel :
+    func : mp_allreduce_sum
+    param: [x]
+  backward: mp_allreduce_sum_grad
+  inplace: (x -> out)
+
 - op : multi_dot
   args : (Tensor[] x)
   output : Tensor
