fix: Do not pass is_first_microbatch to TE when it is not maintained

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
 

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:

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():

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)

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)