Disable TE cross entropy loss fusion

megatron/core/model_parallel_config.py

@@ -434,6 +434,11 @@ def __post_init__(self):
         if self.autocast_dtype is None:
             self.autocast_dtype = self.params_dtype
 
+        assert not (self.cross_entropy_loss_fusion and self.cross_entropy_fusion_impl == 'te'), (
+            "Transformer Engine cross entropy loss fusion is disabled due to stability issues. "
+            "Use cross_entropy_fusion_impl='native', or disable cross_entropy_loss_fusion."
+        )
+
         if self.defer_embedding_wgrad_compute and self.pipeline_model_parallel_size == 1:
             raise ValueError(
                 "Cannot defer embedding wgrad compute when pipeline model parallel is not used"

megatron/core/transformer/moe/README.md

@@ -122,7 +122,7 @@ The following flags are general performance flags that can help to achieve highe
 --moe-router-fusion
 --moe-permute-fusion
 --cross-entropy-loss-fusion
---cross-entropy-fusion-impl te
+--cross-entropy-fusion-impl native
 
 ## Communication optimization
 --use-distributed-optimizer

megatron/training/arguments.py

@@ -1525,6 +1525,13 @@ def validate_args(args, defaults={}):
         assert args.fim_spm_rate, "--fim-spm-rate should be specified."
         assert all(token is not None for token in extra_tokens), "FIM extra tokens should be specified."
 
+    assert not (
+        args.cross_entropy_loss_fusion and args.cross_entropy_fusion_impl == 'te'
+    ), (
+        "Transformer Engine cross entropy loss fusion is disabled due to stability issues. "
+        "Use --cross-entropy-fusion-impl native, or omit --cross-entropy-loss-fusion."
+    )
+
     # Deterministic mode
     if args.deterministic_mode:
         assert not args.use_flash_attn, "Flash attention can not be used in deterministic mode."

tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_gb200/model_config.yaml

@@ -29,7 +29,7 @@ MODEL_ARGS:
   --train-samples: 19531250
   --manual-gc: true
   --cross-entropy-loss-fusion: true
-  --cross-entropy-fusion-impl: te
+  --cross-entropy-fusion-impl: native
   # Transformer Engine args
   --transformer-impl: transformer_engine
   # Data args

tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_sm_gb200/model_config.yaml

@@ -29,7 +29,7 @@ MODEL_ARGS:
   --train-samples: 19531250
   --manual-gc: true
   --cross-entropy-loss-fusion: true
-  --cross-entropy-fusion-impl: te
+  --cross-entropy-fusion-impl: native
   # Transformer Engine args
   --transformer-impl: transformer_engine
   # Data args

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml

@@ -39,7 +39,7 @@ MODEL_ARGS:
   --exit-duration-in-mins: 220
   --no-check-for-nan-in-loss-and-grad: true
   --cross-entropy-loss-fusion: true
-  --cross-entropy-fusion-impl: te
+  --cross-entropy-fusion-impl: native
   --manual-gc: true
   --manual-gc-interval: 10
 

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release_sm/model_config.yaml

@@ -40,7 +40,7 @@ MODEL_ARGS:
   --exit-duration-in-mins: 220
   --no-check-for-nan-in-loss-and-grad: true
   --cross-entropy-loss-fusion: true
-  --cross-entropy-fusion-impl: te
+  --cross-entropy-fusion-impl: native
   --manual-gc: true
   --manual-gc-interval: 10
 

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release/model_config.yaml

@@ -41,7 +41,7 @@ MODEL_ARGS:
   --exit-duration-in-mins: 220
   --no-check-for-nan-in-loss-and-grad: true
   --cross-entropy-loss-fusion: true
-  --cross-entropy-fusion-impl: te
+  --cross-entropy-fusion-impl: native
   --manual-gc: true
   --manual-gc-interval: 10
 

tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release_sm/model_config.yaml

@@ -41,7 +41,7 @@ MODEL_ARGS:
   --exit-duration-in-mins: 220
   --no-check-for-nan-in-loss-and-grad: true
   --cross-entropy-loss-fusion: true
-  --cross-entropy-fusion-impl: te
+  --cross-entropy-fusion-impl: native
   --manual-gc: true
   --manual-gc-interval: 10
 

tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py

@@ -250,7 +250,7 @@ def get_language_model_spec(
         pipeline_dtype=pipeline_dtype,
         bf16=bf16,
         cross_entropy_loss_fusion=True,
-        cross_entropy_fusion_impl='te',
+        cross_entropy_fusion_impl='native',
         calculate_per_token_loss=per_token_loss,
         **extra_kwargs,
     )

tests/unit_tests/test_model_parallel_config.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import pytest
+
+from megatron.core.model_parallel_config import ModelParallelConfig
+
+
+def test_te_cross_entropy_loss_fusion_is_disabled():
+    with pytest.raises(AssertionError, match="Transformer Engine cross entropy loss fusion"):
+        ModelParallelConfig(cross_entropy_loss_fusion=True, cross_entropy_fusion_impl='te')
+
+
+def test_native_cross_entropy_loss_fusion_is_allowed():
+    config = ModelParallelConfig(cross_entropy_loss_fusion=True, cross_entropy_fusion_impl='native')
+
+    assert config.cross_entropy_loss_fusion
+    assert config.cross_entropy_fusion_impl == 'native'