diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 85aa4878301..b754b832e8c 100644 --- a/megatron/core/model_parallel_config.py +++ b/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( diff --git a/tests/unit_tests/test_model_parallel_config.py b/tests/unit_tests/test_model_parallel_config.py index 6eab23ac2e2..c19a2d53958 100644 --- a/tests/unit_tests/test_model_parallel_config.py +++ b/tests/unit_tests/test_model_parallel_config.py @@ -1,13 +1,19 @@ # 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(): @@ -15,3 +21,25 @@ def test_native_cross_entropy_loss_fusion_is_allowed(): 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)