PaddlePaddle
diff --git a/‎paddle/fluid/operators/collective/c_softmax_with_cross_entropy_op.cu
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/operators/collective/c_softmax_with_cross_entropy_op.cu
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/completion.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/static/completion.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/operators/__init__.py
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/static/operators/__init__.py
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
@@ -170,12 +170,12 @@ struct CSoftmaxWithCrossEntropyFunctor<phi::GPUContext, T> {
                             "NCCLCommContext is nullptr, collective op should "
                             "has ring_id attr."));
 
-      stream = comm_ctx->GetStream();
+      stream = dev_ctx.stream();
       VLOG(3) << "new comm_context_manager has ring_id " << rid;
     } else {  // old comm_context
       comm = platform::NCCLCommContext::Instance().Get(rid, place);
 
-      stream = comm->stream();
+      stream = dev_ctx.stream();
       VLOG(3) << "old NCCLCommContext has ring_id " << rid;
     }
 
 
@@ -150,6 +150,9 @@ def _can_apply_infer_spmd_rule(dist_op):
         "split",
         "unsqueeze2",
     ]
+    parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")
+    if parallel_ce == "true":
+        __adapted_ops__.append("softmax_with_cross_entropy")
     op_type = dist_op.serial_op.type
     return enable and contains_spmd_rule(op_type) and op_type in __adapted_ops__
 
 
@@ -11,6 +11,7 @@
 # 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
+import os
 
 from .common import DistributedOperatorImplContainer
 from .common import DistributedOperatorImpl
@@ -42,3 +43,7 @@
 from . import dist_dropout
 from . import dist_flash_attn
 from . import dist_layer_norm
+
+parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")
+if parallel_ce == "true":
+    from . import dist_cross_entropy
Original file line number	Diff line number	Diff line change
`@@ -150,6 +150,9 @@ def _can_apply_infer_spmd_rule(dist_op):`
`150`	`150`	`"split",`
`151`	`151`	`"unsqueeze2",`
`152`	`152`	`]`
	`153`	`+ parallel_ce = os.getenv("PARALLEL_CROSS_ENTROPY")`
	`154`	`+ if parallel_ce == "true":`
	`155`	`+ __adapted_ops__.append("softmax_with_cross_entropy")`
`153`	`156`	`op_type = dist_op.serial_op.type`
`154`	`157`	`return enable and contains_spmd_rule(op_type) and op_type in __adapted_ops__`
`155`	`158`