Modify _legacy_C_ops.c_softmax_with_cross_entropy (#70211)

co63oc · web-flow · commit 390b0f387efc · 2025-02-14T14:16:24.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -90,7 +90,6 @@
 {{"sparse_{name}", (PyCFunction)(void (*)(void))sparse_{name}, METH_VARARGS | METH_KEYWORDS, "C++ interface function for sparse_{name}."}},"""
 
 NEED_GEN_STATIC_ONLY_APIS = [
-    'c_softmax_with_cross_entropy',
     'c_softmax_with_multi_label_cross_entropy',
     'distributed_fused_lamb_init',
     'distributed_fused_lamb_init_',
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -337,6 +337,19 @@
   output : Tensor(x_grad)
   invoke: c_split(out_grad, rank, nranks, ring_id, use_model_parallel)
 
+- backward_op : c_softmax_with_cross_entropy_grad
+  forward: c_softmax_with_cross_entropy (Tensor logits, Tensor label,  int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0) -> Tensor(softmax), Tensor(loss)
+  args: (Tensor softmax, Tensor label, Tensor loss_grad,int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0)
+  output: Tensor(logits_grad)
+  infer_meta :
+    func: CSoftmaxWithCrossEntropyGradInferMeta
+    spmd_rule : CSoftmaxWithCrossEntropyGradSpmd
+    param: [softmax, label, loss_grad, ignore_index, rank, nranks]
+  kernel:
+    func: c_softmax_with_cross_entropy_grad
+    data_type: loss_grad
+    param: [softmax, label, loss_grad, ignore_index, rank, nranks]
+
 - backward_op : cast_grad
   forward : cast (Tensor x, DataType dtype) -> Tensor(out)
   args : (Tensor x, Tensor out_grad)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_backward.yaml b/paddle/phi/ops/yaml/inconsistent/static_backward.yaml
@@ -107,19 +107,6 @@
     func : c_embedding_grad
   no_need_buffer : weight
 
-- backward_op : c_softmax_with_cross_entropy_grad
-  forward: c_softmax_with_cross_entropy (Tensor logits, Tensor label,  int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0) -> Tensor(softmax), Tensor(loss)
-  args: (Tensor softmax, Tensor label, Tensor loss_grad,int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0)
-  output: Tensor(logits_grad)
-  infer_meta :
-    func: CSoftmaxWithCrossEntropyGradInferMeta
-    spmd_rule : CSoftmaxWithCrossEntropyGradSpmd
-    param: [softmax, label, loss_grad, ignore_index, rank, nranks]
-  kernel:
-    func: c_softmax_with_cross_entropy_grad
-    data_type: loss_grad
-    param: [softmax, label, loss_grad, ignore_index, rank, nranks]
-
 - backward_op : c_softmax_with_multi_label_cross_entropy_grad
   forward: c_softmax_with_multi_label_cross_entropy (Tensor logits, Tensor label, Tensor smooth_weight, int64_t ignore_index=-100, bool sum_multi_label_loss=true, int ring_id=0, int rank=0, int nranks=0) -> Tensor(softmax), Tensor(loss)
   args: (Tensor softmax, Tensor label, Tensor smooth_weight, Tensor loss_grad, int64_t ignore_index=-100, bool sum_multi_label_loss=true, int ring_id=0, int rank=0, int nranks=0)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -939,19 +939,6 @@
   args : (Tensor i, Tensor x)
   output : Tensor[](out)
 
-- op: c_softmax_with_cross_entropy
-  args: (Tensor logits, Tensor label,  int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0)
-  output: Tensor(softmax), Tensor(loss)
-  infer_meta:
-    func : CSoftmaxWithCrossEntropyInferMeta
-    spmd_rule : CSoftmaxWithCrossEntropyInferSpmd
-    param: [logits, label, ignore_index, rank, nranks]
-  kernel:
-    func: c_softmax_with_cross_entropy
-    data_type : logits
-    param: [logits, label, ignore_index, rank, nranks]
-  backward: c_softmax_with_cross_entropy_grad
-
 - op: c_softmax_with_multi_label_cross_entropy
   args: (Tensor logits, Tensor label, Tensor smooth_weight, int64_t ignore_index=-100, bool sum_multi_label_loss=true, int ring_id=0, int rank=0, int nranks=0)
   output: Tensor(softmax), Tensor(loss)
diff --git a/paddle/phi/ops/yaml/legacy/backward_exclude.yaml b/paddle/phi/ops/yaml/legacy/backward_exclude.yaml
@@ -4,6 +4,7 @@
 
 - amax_grad
 - amin_grad
+- c_softmax_with_cross_entropy_grad
 - cast_grad
 - conv2d_transpose_double_grad
 - conv2d_transpose_grad
diff --git a/paddle/phi/ops/yaml/legacy/ops_exclude.yaml b/paddle/phi/ops/yaml/legacy/ops_exclude.yaml
@@ -13,6 +13,7 @@
 - c_allreduce_sum
 - c_concat
 - c_identity
+- c_softmax_with_cross_entropy
 - c_split
 - cast
 - conv2d_transpose
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -796,6 +796,19 @@
     func : c_scatter
   traits : paddle::dialect::ForwardOnlyTrait
 
+- op : c_softmax_with_cross_entropy
+  args: (Tensor logits, Tensor label,  int64_t ignore_index=-100, int ring_id=0, int rank=0, int nranks=0)
+  output: Tensor(softmax), Tensor(loss)
+  infer_meta:
+    func : CSoftmaxWithCrossEntropyInferMeta
+    spmd_rule : CSoftmaxWithCrossEntropyInferSpmd
+    param: [logits, label, ignore_index, rank, nranks]
+  kernel:
+    func: c_softmax_with_cross_entropy
+    data_type : logits
+    param: [logits, label, ignore_index, rank, nranks]
+  backward: c_softmax_with_cross_entropy_grad
+
 - op : c_split
   args : (Tensor x, int rank = 0, int nranks = 1, int ring_id = 0, bool use_model_parallel = true)
   output : Tensor(out)
diff --git a/python/paddle/distributed/fleet/layers/mpu/mp_ops.py b/python/paddle/distributed/fleet/layers/mpu/mp_ops.py
@@ -419,17 +419,8 @@ def _c_softmax_with_cross_entropy(
         )
 
     if in_dynamic_mode():
-        softmax, loss = _legacy_C_ops.c_softmax_with_cross_entropy(
-            logits,
-            label,
-            'ring_id',
-            ring_id,
-            'rank',
-            rank,
-            'nranks',
-            nranks,
-            'ignore_index',
-            ignore_index,
+        softmax, loss = _C_ops.c_softmax_with_cross_entropy(
+            logits, label, ignore_index, ring_id, rank, nranks
         )
         if not return_softmax:
             return loss
