[fluid_ops] Modify _legacy_C_ops.c_split (#70441)

co63oc · web-flow · commit 6368f90a30d8 · 2024-12-25T11:25:29.000+08:00
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -150,16 +150,6 @@
     func : c_reduce_prod
   inplace : (x -> out)
 
-- op : c_split
-  args : (Tensor x, int rank = 0, int nranks = 1, int ring_id = 0, bool use_model_parallel = true)
-  output : Tensor(out)
-  infer_meta :
-    func : CSplitInferMeta
-    param : [x, nranks]
-  kernel :
-    func : c_split
-    param: [x, rank, nranks, use_model_parallel]
-
 - op : coalesce_tensor_
   args : (Tensor[] input, DataType dtype, bool copy_data = false, bool set_constant = false, bool persist_output = false, float constant = 0.0, bool use_align = true, int align_size = -1, int size_of_dtype = -1, int64_t[] concated_shapes = {}, int64_t[] concated_ranks = {})
   output : Tensor[](output){input.size()}, Tensor(fused_output)
diff --git a/paddle/phi/ops/yaml/legacy/ops_exclude.yaml b/paddle/phi/ops/yaml/legacy/ops_exclude.yaml
@@ -19,6 +19,7 @@
 - c_concat
 - c_identity
 - c_reduce_sum
+- c_split
 - cast
 - conv2d_transpose
 - conv2d_transpose_bias
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -849,6 +849,16 @@
     func : c_scatter
   traits : paddle::dialect::ForwardOnlyTrait
 
+- op : c_split
+  args : (Tensor x, int rank = 0, int nranks = 1, int ring_id = 0, bool use_model_parallel = true)
+  output : Tensor(out)
+  infer_meta :
+    func : CSplitInferMeta
+    param : [x, nranks]
+  kernel :
+    func : c_split
+    param: [x, rank, nranks, use_model_parallel]
+
 - op : calc_reduced_attn_scores
   args : (Tensor q, Tensor k, Tensor softmax_lse)
   output : Tensor(reduced_scores)
diff --git a/python/paddle/distributed/fleet/layers/mpu/mp_ops.py b/python/paddle/distributed/fleet/layers/mpu/mp_ops.py
@@ -58,19 +58,7 @@ class c_split_eager(PyLayer):
     def forward(ctx, tensor, group, rank, nranks):
         ctx.group = group
         ctx.nranks = nranks
-        return _legacy_C_ops.c_split(
-            tensor,
-            'use_calc_stream',
-            True,
-            'ring_id',
-            group.id,
-            'rank',
-            rank,
-            'nranks',
-            nranks,
-            'use_model_parallel',
-            True,
-        )
+        return _C_ops.c_split(tensor, rank, nranks, group.id, True)
 
     @staticmethod
     def backward(ctx, dy):
