[deepakn94]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Follows up #6965. That PR stopped passing is_first_microbatch to TE when the config does not maintain it, at config granularity. This extends the same predicate to per-module granularity, because a quant_recipe decides quantization per module and in both directions.

Not a bug fix

Stating this up front, since the phrasing in an earlier revision of this description overstated it. A mis-tracked module today still runs once per microbatch, and set_is_first_microbatch still re-arms the flag every step, because it asks the config and the config does name a recipe. So the microbatch-1 True lands on a gradient buffer that was just zeroed, and nothing is lost.

This removes a trap rather than repairing an observed failure. The staleness that motivated #6965, where nothing re-arms the flag at all, does not apply here.

The gap

is_first_microbatch_tracked reads the config alone:

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)

It answers True for every module as soon as quant_recipe is set, including a module the recipe forces to high precision while the config enables fp8. TEQuantizationRecipe leaves both fp8_quantization_recipe and fp4_quantization_recipe unset for that case, and _get_should_context_be_quantized_recipe returns False for it.

Such a module has no quantized weight cache to refresh, so the flag carries no information for it. It is not inert, though: TE selects the weight-gradient accumulation mode from the same argument,

accumulate = fuse_wgrad_accumulation and not is_first_microbatch

so the module takes the overwrite path for nothing. That is survivable only because the buffer is freshly zeroed and the module is invoked once. It becomes real gradient loss for any module invoked more than once per microbatch, which is exactly the failure #6965 fixed for repeated MTP layers.

Testing te_quant_params is not None would not be enough either: a force-high-precision entry is still a recipe match, so the params are present while the layer runs in high precision. The question has to be answered from the recipe contents.

The change

Delegate to TE's own answer. will_execute_quantized reads the recipe contents rather than its presence, honours override_quantized_autocast, and distinguishes the training and evaluation recipes. It is the same hook fp8_utils.py:965 and the MoE experts already use, so no new mechanism is introduced.

• _resolve_is_first_microbatch passes the module and the ambient autocast state, read before entering the module's own quantization context, since that is the state a recipe overrides.
• Callers with no module in hand, meaning MegatronModule.set_is_first_microbatch, keep the model-wide answer unchanged. Verified across all sixteen combinations of fp8, fp4, use_kitchen and quant_recipe.
• use_kitchen returns early: Kitchen drives quantization outside TE's autocast state, so there is no per-module answer to consult.

Behaviour is unchanged for any config without a quant_recipe.

Tests

Seven cases in tests/unit_tests/transformer/test_module.py, nine once parametrization expands, none needing TE or a GPU: an unquantized config; the model-wide answers with no module passed, which pin what set_is_first_microbatch relies on; Kitchen staying tracked even when the module says otherwise; a recipe quantizing an otherwise-unquantized layer; a recipe forcing high precision under fp8 being untracked; a module with no recipe following the ambient autocast in both directions; and a module lacking the hook falling back to the config.

tests/unit_tests/determinism/kernels/test_te_wrappers.py gains two cases for _resolve_is_first_microbatch on a real TE linear, since megatron/core/extensions/transformer_engine.py is a registered kernel source.

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter on my PR

black, isort, pylint (10.00/10) and ruff clean. The unit tests were not executed locally: the megatron.core import chain does not load on macOS, and the determinism cases need a GPU with TE.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[zhongbozhu]

Do we still need this PR if we have this #6965 merged already?