From ff6414418075816eac874cdc342092c9c2f6c1b2 Mon Sep 17 00:00:00 2001 From: Mike Chrzanowski Date: Wed, 3 Jun 2026 17:29:31 -0400 Subject: [PATCH 1/2] Disable TE cross entropy loss fusion (#5115) Co-authored-by: Mike Chrzanowski Signed-off-by: NeMo Bot --- megatron/core/model_parallel_config.py | 5 +++++ megatron/core/transformer/moe/README.md | 2 +- megatron/training/arguments.py | 7 +++++++ .../gpt3_15b_8t_release_gb200/model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../models/test_mimo_1f1b_schedule.py | 2 +- tests/unit_tests/test_model_parallel_config.py | 17 +++++++++++++++++ 11 files changed, 37 insertions(+), 8 deletions(-) create mode 100644 tests/unit_tests/test_model_parallel_config.py diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index dabe0d0aced..85aa4878301 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -428,6 +428,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" diff --git a/megatron/core/transformer/moe/README.md b/megatron/core/transformer/moe/README.md index 6a268ce43e9..2f731fcb6f0 100644 --- a/megatron/core/transformer/moe/README.md +++ b/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 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 18cb5ef8a15..9bdc0ee7a21 100644 --- a/megatron/training/arguments.py +++ b/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." diff --git a/tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_gb200/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_gb200/model_config.yaml index 0a6f3403fa9..aa945c0691a 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_gb200/model_config.yaml +++ b/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 diff --git a/tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_sm_gb200/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_sm_gb200/model_config.yaml index ed04acc8671..2a0bb0ddbef 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_15b_8t_release_sm_gb200/model_config.yaml +++ b/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 diff --git a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml b/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml index b0f383f739f..6cba2ff3abf 100644 --- a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release/model_config.yaml +++ b/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 diff --git a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release_sm/model_config.yaml b/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release_sm/model_config.yaml index a85351277f7..c4189ea2de9 100644 --- a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp1pp4emp16etp1cp1_release_sm/model_config.yaml +++ b/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 diff --git a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release/model_config.yaml b/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release/model_config.yaml index aaac6f94848..6aaed5f7b4f 100644 --- a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release/model_config.yaml +++ b/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 diff --git a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release_sm/model_config.yaml b/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release_sm/model_config.yaml index c8492799abd..c2a4f7df6c5 100644 --- a/tests/functional_tests/test_cases/mixtral/deepseekv3_proxy_flex_tp2pp2emp16etp1cp1_gb_200_release_sm/model_config.yaml +++ b/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 diff --git a/tests/unit_tests/models/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/test_mimo_1f1b_schedule.py index 836382b21cc..e86fd85c42b 100644 --- a/tests/unit_tests/models/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/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, ) diff --git a/tests/unit_tests/test_model_parallel_config.py b/tests/unit_tests/test_model_parallel_config.py new file mode 100644 index 00000000000..6eab23ac2e2 --- /dev/null +++ b/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' From 3f9ed1e083a1b9a641534bf1c5ba51cc75b0c0f7 Mon Sep 17 00:00:00 2001 From: Yu Yao <54727607+yaoyu-33@users.noreply.github.com> Date: Fri, 5 Jun 2026 05:55:00 -0700 Subject: [PATCH 2/2] cp: `Move TE cross entropy guard to training args (5162)` into `core_r0.18.0` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: yaoyu-33 (cherry picked from commit b57449928aafaaec7a36e8325d83aacf05e91bd4) Signed-off-by: oliver könig --- megatron/core/model_parallel_config.py | 12 ++++--- .../unit_tests/test_model_parallel_config.py | 34 +++++++++++++++++-- 2 files changed, 39 insertions(+), 7 deletions(-) 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)