diff --git a/docs/user-guide/deterministic-training.md b/docs/user-guide/deterministic-training.md index 9242a3e42fe..e21a10d1f04 100644 --- a/docs/user-guide/deterministic-training.md +++ b/docs/user-guide/deterministic-training.md @@ -45,7 +45,7 @@ Checked against the parsed `args` Namespace in `apply_determinism_to_args`. Inco | `--tp-comm-overlap` | Must be off — asserted (the overlap path is not bit-exact); drop the flag yourself | | `torch.use_deterministic_algorithms` | Set to `True` | -Flash attention is permitted: Transformer Engine's flash-attention backend is deterministic when `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` (see the [Transformer Engine docs](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html)). +Flash attention is permitted: Transformer Engine's flash-attention backend is deterministic when `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` (see the [Transformer Engine docs](https://docs.nvidia.com/deeplearning/transformer-engine/api/pytorch.html)). ## Verifying determinism diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 5e75c34b1b2..8ee7dbe5f97 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -49,6 +49,7 @@ from megatron.core.tensor_parallel.utils import divide from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.module import is_first_microbatch_tracked from megatron.core.transformer.torch_norm import LayerNormInterface from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import ( @@ -400,6 +401,21 @@ def _get_should_context_be_quantized_params( ) +def _resolve_is_first_microbatch(module) -> Optional[bool]: + """The value to pass TE, or ``None`` meaning "no opinion, just accumulate". + + A ``True`` tells TE the gradient is fresh, so backward writes over ``main_grad`` instead of + adding into it. Pass the flag on only when it can be trusted. + """ + if ( + module.disable_parameter_transpose_cache + or not is_first_microbatch_tracked(module.config) + or getattr(module, 'is_repeated_layer', False) + ): + return None + return module.is_first_microbatch + + def _get_extra_te_kwargs(config: TransformerConfig): extra_transformer_engine_kwargs = {"params_dtype": config.params_dtype} @@ -1374,9 +1390,7 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: """Forward.""" - _is_first_microbatch = ( - None if self.disable_parameter_transpose_cache else self.is_first_microbatch - ) + _is_first_microbatch = _resolve_is_first_microbatch(self) quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) with quant_context: @@ -1623,9 +1637,7 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: def forward(self, x): """Forward.""" - _is_first_microbatch = ( - None if self.disable_parameter_transpose_cache else self.is_first_microbatch - ) + _is_first_microbatch = _resolve_is_first_microbatch(self) quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) # FP32 residual connections pass the FP32 residual stream into this fused module, but @@ -2775,9 +2787,7 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: def forward(self, x, m_splits): """Forward.""" - _is_first_microbatch = ( - None if self.disable_parameter_transpose_cache else self.is_first_microbatch - ) + _is_first_microbatch = _resolve_is_first_microbatch(self) quant_context = _get_fp8_autocast_for_quant_params(self.te_quant_params, self.training) with quant_context: diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index a1524f10c00..2cd43bea15c 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -27,6 +27,24 @@ def param_is_not_shared(param): # pylint: disable=missing-function-docstring return not hasattr(param, 'shared') or not param.shared +def is_first_microbatch_tracked(config) -> bool: + """True if ``is_first_microbatch`` is still being kept up to date. + + A training step runs N microbatches. The flag marks microbatch 1 -- the one that + re-quantizes the weights and starts a fresh main_grad, while 2..N reuse and accumulate:: + + layer is built -> flag = True + every forward -> flag = False (microbatch 1 is over) + start of each step -> flag = True (only quantized configs) + """ + return ( + config.fp8 is not None + or config.fp4 is not None + or getattr(config, 'use_kitchen', False) + or getattr(config, 'quant_recipe', None) is not None + ) + + class MegatronModule(torch.nn.Module): """Base Megatron module inhertied by all Models. @@ -117,12 +135,7 @@ def set_is_first_microbatch(self): If kitchen is being used, kitchen controls quantization level. A quant_recipe (e.g. from --te-precision-config-file) also enables the flag. """ - if ( - self.config.fp8 is not None - or self.config.fp4 is not None - or getattr(self.config, 'use_kitchen', False) - or getattr(self.config, 'quant_recipe', None) is not None - ): + if is_first_microbatch_tracked(self.config): if not hasattr(self, "modules_with_is_first_microbatch"): self.modules_with_is_first_microbatch = [] for m in self.modules(): diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index eb4c565ec50..4c1d1d35f03 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -2182,6 +2182,13 @@ def __init__( self._build_layers(pg_collection) assert len(self.layers) > 0, "MultiTokenPredictionBlock must have at least one layer." + + if self.mtp_use_repeated_layer: + # One layer object, called once per MTP depth, every call adding into the same + # main_grad. A True would make one of those calls overwrite instead of add. + for m in self.layers.modules(): + if hasattr(m, 'is_first_microbatch'): + m.is_repeated_layer = True self.cp_group = pg_collection.cp self.tp_group = pg_collection.tp self.tp_cp_group = getattr(pg_collection, 'tp_cp', None) diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 7b7eea29fc3..c4698258558 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -10,7 +10,7 @@ from megatron.core.context_parallel import ContextParallelBatch from megatron.core.enums import ModelType -from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.extensions.transformer_engine import HAVE_TE, _resolve_is_first_microbatch from megatron.core.inference.utils import InferenceMode from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_spec, @@ -103,7 +103,7 @@ def teardown_method(self, method): destroy_global_vars() destroy_num_microbatches_calculator() - def _create_config_and_mtp_block_spec(self, tp, cp, use_te=False): + def _create_config_and_mtp_block_spec(self, tp, cp, use_te=False, use_repeated_layer=False): Utils.initialize_model_parallel(tensor_model_parallel_size=tp, context_parallel_size=cp) config = TransformerConfig( mtp_num_layers=2, @@ -114,6 +114,7 @@ def _create_config_and_mtp_block_spec(self, tp, cp, use_te=False): tensor_model_parallel_size=tp, sequence_parallel=True if tp > 1 else False, context_parallel_size=cp, # Enable CP for MTP testing + mtp_use_repeated_layer=use_repeated_layer, ) if use_te: transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec() @@ -626,6 +627,29 @@ def test_constructor_ues_te(self, tp, cp): elif tp == 4: assert num_weights == 15216 * config.mtp_num_layers + @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") + @pytest.mark.parametrize('repeated', [False, True]) + def test_repeated_layer_opts_out_of_is_first_microbatch(self, repeated): + """A shared MTP layer must hand TE None, because its first forward is its last backward.""" + torch.manual_seed(_SEED) + config, mtp_block_spec = self._create_config_and_mtp_block_spec( + tp=1, cp=1, use_te=True, use_repeated_layer=repeated + ) + mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec) + + # Quantization is what would otherwise make TE act on the flag, so ask under it. + config.fp8 = "hybrid" + mtp.set_is_first_microbatch() + te_modules = [m for m in mtp.modules() if hasattr(m, 'is_first_microbatch')] + assert te_modules, "expected the TE spec to produce modules carrying is_first_microbatch" + + if repeated: + assert len(mtp.layers) == 1 + assert all(_resolve_is_first_microbatch(m) is None for m in te_modules) + else: + assert len(mtp.layers) == config.mtp_num_layers + assert all(_resolve_is_first_microbatch(m) is True for m in te_modules) + def test_get_embeddings_rolls_padding_mask(self): """Test that _get_embeddings rolls padding_mask alongside input ids.""" torch.manual_seed(_SEED)