cp: Move TE cross entropy guard to training args (5162) into core_r0.18.0

megatron/core/model_parallel_config.py

@@ -428,10 +428,14 @@ 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.cross_entropy_loss_fusion and self.cross_entropy_fusion_impl == 'te':
+            warnings.warn(
+                "Transformer Engine cross entropy loss fusion has known stability issues. "
+                "Megatron-LM training args validation rejects this combination by default. "
+                "Use cross_entropy_fusion_impl='native', or disable cross_entropy_loss_fusion.",
+                UserWarning,
+                stacklevel=2,
+            )
 
         if self.defer_embedding_wgrad_compute and self.pipeline_model_parallel_size == 1:
             raise ValueError(

tests/unit_tests/test_model_parallel_config.py

@@ -1,17 +1,45 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import sys
+
 import pytest
 
 from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.training.arguments import parse_args, validate_args
 
 
-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_te_cross_entropy_loss_fusion_warns_in_model_parallel_config():
+    with pytest.warns(UserWarning, match="known stability issues"):
+        config = ModelParallelConfig(cross_entropy_loss_fusion=True, cross_entropy_fusion_impl='te')
+
+    assert config.cross_entropy_loss_fusion
+    assert config.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'
+
+
+def test_te_cross_entropy_loss_fusion_is_disabled_by_training_args(monkeypatch):
+    monkeypatch.setattr(sys, 'argv', ['test_model_parallel_config.py'])
+    args = parse_args()
+    args.num_layers = 2
+    args.hidden_size = 128
+    args.num_attention_heads = 4
+    args.max_position_embeddings = 1024
+    args.seq_length = 1024
+    args.micro_batch_size = 1
+    # Let validate_args derive a global batch size that is valid for the
+    # active data-parallel size in distributed unit-test jobs.
+    args.train_iters = 1
+    args.lr = 1e-4
+    args.tokenizer_type = 'NullTokenizer'
+    args.vocab_size = 1024
+    args.cross_entropy_loss_fusion = True
+    args.cross_entropy_fusion_impl = 'te'
+
+    with pytest.raises(AssertionError, match="Transformer Engine cross entropy loss fusion"):
+        validate_args(args)