partial_allgather remove use_calc_stream param [fluid_ops] (#68954)

co63oc · web-flow · commit 1eabf1ea47cf · 2024-10-31T11:58:36.000+08:00
* Fix

* Fix

* Fix

* Fix
diff --git a/paddle/fluid/operators/collective/partial_allgather_op.cc b/paddle/fluid/operators/collective/partial_allgather_op.cc
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3132,8 +3132,6 @@ void Pad3dInferMeta(const MetaTensor& x,
 void PartialAllgatherInferMeta(const MetaTensor& x,
                                int nranks,
                                int rank,
-                               int ring_id,
-                               bool use_calc_stream,
                                MetaTensor* out) {
   PADDLE_ENFORCE_GE(
       nranks,
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -499,8 +499,6 @@ void Pad3dInferMeta(const MetaTensor& x,
 void PartialAllgatherInferMeta(const MetaTensor& x,
                                int nranks,
                                int rank,
-                               int ring_id,
-                               bool use_calc_stream,
                                MetaTensor* out);
 
 void PartialSendInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/kernels/cpu/partial_allgather_kernel.cc b/paddle/phi/kernels/cpu/partial_allgather_kernel.cc
@@ -21,8 +21,6 @@ void PartialAllGatherOpCPUKernel(const Context &dev_ctx,
                                  const DenseTensor &x,
                                  int nranks,
                                  int rank,
-                                 int ring_id,
-                                 bool use_calc_stream,
                                  DenseTensor *out) {
   PADDLE_THROW(common::errors::Unavailable(
       "Do not support partial_allgather for cpu kernel now."));
diff --git a/paddle/phi/kernels/gpu/partial_allgather_kernel.cu b/paddle/phi/kernels/gpu/partial_allgather_kernel.cu
@@ -27,8 +27,6 @@ void PartialAllGatherOpCUDAKernel(const Context& dev_ctx,
                                   const DenseTensor& x_in,
                                   int nranks,
                                   int rank,
-                                  int ring_id UNUSED,
-                                  bool use_calc_stream UNUSED,
                                   DenseTensor* out) {
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
   auto in = &x_in;
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -665,15 +665,6 @@
   interfaces : paddle::dialect::InferSymbolicShapeInterface
   traits : paddle::dialect::ForwardOnlyTrait
 
-- op : partial_allgather
-  args : (Tensor x, int nranks, int rank, int ring_id = 0, bool use_calc_stream = false)
-  output : Tensor(out)
-  infer_meta :
-    func: PartialAllgatherInferMeta
-  kernel :
-    func : partial_allgather
-  inplace : (x -> out)
-
 - op : partial_recv
   args : (int ring_id = 0, int peer = 0, DataType dtype=DataType::FLOAT32, int[] out_shape= {}, int num = 1, int id = 0)
   output : Tensor(out)
diff --git a/paddle/phi/ops/yaml/op_version.yaml b/paddle/phi/ops/yaml/op_version.yaml
@@ -499,6 +499,14 @@
           comment : Compute as vector when axis is None and input is matrix.
           default : "false"
 
+- op : partial_allgather
+  version :
+    - checkpoint : Upgrade partial_allgather delete 1 attribute[use_calc_stream].
+      action :
+        - delete_attr : use_calc_stream
+          comment : eject CUDA operations to calculation stream.
+          default : false
+
 - op : pixel_shuffle
   version :
     - checkpoint : Compatible upgrade of pixel_shuffle, add a new attribute [data_format]
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3733,6 +3733,17 @@
   backward : pad3d_grad
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
+- op : partial_allgather
+  args : (Tensor x, int nranks, int rank, int ring_id = 0)
+  output : Tensor(out)
+  infer_meta :
+    func: PartialAllgatherInferMeta
+    param: [x, nranks, rank]
+  kernel :
+    func : partial_allgather
+    param: [x, nranks, rank]
+  inplace : (x -> out)
+
 - op : partial_concat
   args : (Tensor[] x, int start_index = 0, int length = -1)
   output : Tensor(out)
