diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 95150103f27..fa59d27ce2a 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math @@ -957,6 +957,14 @@ class TransformerConfig(ModelParallelConfig): transformed to an empty list in __post_init__. The deprecated values "full_iteration" and "full_iteration_inference" are also accepted and migrated to the new API in __post_init__.""" + create_attention_mask_in_dataloader: bool = True + """Whether training data loaders create and pass an attention_mask tensor. + + This mirrors the training argument of the same name so Transformer Engine CUDA graph capture + can preserve the actual forward signature. When disabled for causal/no-mask attention, + the graph capture path should not allocate a synthetic global-sequence attention mask. + """ + inference_cuda_graph_scope: Optional[InferenceCudaGraphScope] = field( default=None, metadata={ diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ddd1e7d34cd..b8f36741597 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from __future__ import annotations import functools @@ -19,7 +19,12 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.cuda_graphs import is_graph_capturing -from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType +from megatron.core.transformer.enums import ( + AttnMaskType, + CudaGraphModule, + InferenceCudaGraphScope, + LayerType, +) from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.mlp import MLP from megatron.core.transformer.module import GraphableMegatronModule @@ -1049,13 +1054,29 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): not self.config.cuda_graph_modules or CudaGraphModule.attn in self.config.cuda_graph_modules ): - slen_per_cp = seq_length // self.config.context_parallel_size - static_inputs["attention_mask"] = ( - ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool()) - .to(torch.cuda.current_device()) - .reshape(1, 1, slen_per_cp, seq_length) - .tile(micro_batch_size, 1, 1, 1) - ) + if not self.config.create_attention_mask_in_dataloader: + if self.self_attention.attn_mask_type not in ( + AttnMaskType.causal, + AttnMaskType.no_mask, + AttnMaskType.causal_bottom_right, + ): + log_single_rank( + logger, + logging.WARNING, + "TE CUDA graph capture is omitting attention_mask because " + "create_attention_mask_in_dataloader is False, but " + f"attn_mask_type={self.self_attention.attn_mask_type.name} may require " + "an explicit mask. Ensure this is intended for the current workload.", + ) + else: + slen_per_cp = seq_length // self.config.context_parallel_size + static_inputs["attention_mask"] = ( + ~(torch.tril(torch.ones((slen_per_cp, seq_length))).bool()) + .to(torch.cuda.current_device()) + .reshape(1, 1, slen_per_cp, seq_length) + .tile(micro_batch_size, 1, 1, 1) + ) + return static_inputs def _get_submodules_under_cudagraphs(self): @@ -1300,11 +1321,17 @@ def get_zero_attention_mask(slen_per_tpcp, micro_batch_size): 'attention_mask' in cudagraph_kwargs and cudagraph_kwargs['attention_mask'] is None ): # The attention_mask can be None when there is no padding to the input sequence. - # However, an attention_mask Tensor must be passed into cudagraph for replay, so - # we create an equivalent zero Tensor as the attention_mask. - cudagraph_kwargs["attention_mask"] = get_zero_attention_mask( - hidden_states.size(0), hidden_states.size(1) - ) + # If the dataloader never creates masks, the TE CUDA graph was captured without + # this kwarg. Preserve that signature instead of allocating a synthetic + # [local_seq, global_seq] zero mask. + if not self.config.create_attention_mask_in_dataloader: + cudagraph_kwargs.pop("attention_mask") + else: + # The graph was captured with an attention_mask tensor, so replay needs a + # tensor input with the same semantic value as no padding. + cudagraph_kwargs["attention_mask"] = get_zero_attention_mask( + hidden_states.size(0), hidden_states.size(1) + ) except ImportError: raise RuntimeError("CUDAGraph requires TransformerEngine, but not installed") return tuple(cudagraph_args), cudagraph_kwargs diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e7415dc3019..1a0ff0af0df 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron arguments.""" @@ -276,12 +276,12 @@ def _normalize_cuda_graph_modules_args(args): args.cuda_graph_modules ) validate_deprecated_cuda_graph_modules_migration_inputs( - deprecated_scopes, - args.cuda_graph_impl, - args.inference_cuda_graph_scope, + deprecated_scopes, args.cuda_graph_impl, args.inference_cuda_graph_scope ) if used_full_scope: - warn_rank_0('full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.') + warn_rank_0( + 'full scope is deprecated. Use empty cuda_graph_modules to capture the whole layer.' + ) for scope, attr, value in deprecated_scopes: migration = get_deprecated_cuda_graph_modules_migration( @@ -629,12 +629,10 @@ def validate_args(args, defaults={}): del args.external_cuda_graph if getattr(args, 'cuda_graph_scope_deprecated', None) is not None: - assert not args.cuda_graph_modules, ( - "--cuda-graph-scope and --cuda-graph-modules cannot be used together." - ) - warn_rank_0( - '--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.' - ) + assert ( + not args.cuda_graph_modules + ), "--cuda-graph-scope and --cuda-graph-modules cannot be used together." + warn_rank_0('--cuda-graph-scope is deprecated, use --cuda-graph-modules instead.') args.cuda_graph_modules = args.cuda_graph_scope_deprecated del args.cuda_graph_scope_deprecated @@ -642,10 +640,7 @@ def validate_args(args, defaults={}): # all subsequent validation sees fully-typed enum values. _normalize_cuda_graph_modules_args(args) _normalize_inference_cuda_graph_scope_arg(args) - assert ( - args.inference_cuda_graph_scope - in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl] - ), ( + assert args.inference_cuda_graph_scope in ALLOWED_INFERENCE_SCOPES[args.cuda_graph_impl], ( "Invalid inference CUDA graph scope " f"{args.inference_cuda_graph_scope.name!r} for " f"--cuda-graph-impl={args.cuda_graph_impl!r}." @@ -2095,6 +2090,8 @@ def _add_network_size_args(parser): "num_moe_experts", "fp8_param", "fp4_param", + # already generated by data args + "create_attention_mask_in_dataloader", # incompatible defaults in dataclass "gradient_accumulation_fusion", "overlap_p2p_comm", diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 5dce465cfe6..734790b8915 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import hashlib import inspect @@ -68,6 +68,7 @@ "cpu_offloading_weights": False, "cross_entropy_fusion_impl": "native", "cross_entropy_loss_fusion": True, + "create_attention_mask_in_dataloader": True, "cuda_graph_impl": "none", "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index 93650cf13b0..f5809bfca2d 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc @@ -20,7 +20,7 @@ model_parallel_cuda_manual_seed, ) from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord -from megatron.core.transformer.enums import InferenceCudaGraphScope +from megatron.core.transformer.enums import AttnMaskType, CudaGraphModule, InferenceCudaGraphScope from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( TransformerLayer, @@ -73,6 +73,89 @@ def test_gpu_forward(self): assert hidden_states.shape[1] == micro_batch_size assert hidden_states.shape[2] == config.hidden_size + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.10.0")), + reason="TE CUDA graph kwargs require TransformerEngine >= 1.10", + ) + def test_te_cuda_graph_omits_absent_attention_mask(self): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + cuda_graph_modules=[CudaGraphModule.attn], + create_attention_mask_in_dataloader=False, + ) + layer = TransformerLayer(config, get_gpt_layer_with_transformer_engine_submodules()) + + static_inputs = layer.get_layer_static_inputs(seq_length=32, micro_batch_size=1) + assert "attention_mask" not in static_inputs + + hidden_states = torch.ones((32, 1, config.hidden_size), device="cuda") + _, cudagraph_kwargs = layer._get_te_cuda_graph_replay_args( + hidden_states, attention_mask=None + ) + assert "attention_mask" not in cudagraph_kwargs + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.10.0")), + reason="TE CUDA graph kwargs require TransformerEngine >= 1.10", + ) + def test_te_cuda_graph_keeps_configured_attention_mask(self): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + cuda_graph_modules=[CudaGraphModule.attn], + create_attention_mask_in_dataloader=True, + ) + layer = TransformerLayer(config, get_gpt_layer_with_transformer_engine_submodules()) + + static_inputs = layer.get_layer_static_inputs(seq_length=32, micro_batch_size=1) + assert static_inputs["attention_mask"].shape == (1, 1, 32, 32) + + hidden_states = torch.ones((32, 1, config.hidden_size), device="cuda") + _, cudagraph_kwargs = layer._get_te_cuda_graph_replay_args( + hidden_states, attention_mask=None + ) + assert cudagraph_kwargs["attention_mask"].shape == (1, 1, 32, 32) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + @pytest.mark.skipif( + not (HAVE_TE and is_te_min_version("1.10.0")), + reason="TE CUDA graph kwargs require TransformerEngine >= 1.10", + ) + def test_te_cuda_graph_warns_when_omitting_padding_attention_mask(self, monkeypatch): + config = TransformerConfig( + num_layers=2, + hidden_size=12, + num_attention_heads=4, + use_cpu_initialization=True, + cuda_graph_impl="transformer_engine", + cuda_graph_modules=[CudaGraphModule.attn], + create_attention_mask_in_dataloader=False, + ) + layer = TransformerLayer(config, get_gpt_layer_with_transformer_engine_submodules()) + layer.self_attention.attn_mask_type = AttnMaskType.padding + + warning_messages = [] + monkeypatch.setattr( + "megatron.core.transformer.transformer_layer.log_single_rank", + lambda _logger, _level, message: warning_messages.append(message), + ) + + static_inputs = layer.get_layer_static_inputs(seq_length=32, micro_batch_size=1) + + assert "attention_mask" not in static_inputs + assert any( + "attn_mask_type=padding may require an explicit mask" in msg for msg in warning_messages + ) + def test_chunked_mlp(self): with torch.no_grad(): num_layers = 2