Set is_first_microbatch when quant_recipe is configured

[yezhengmao1]

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

What does this PR do?

Enable the is_first_microbatch flag when quantization is driven by a
quant_recipe (e.g. via --te-precision-config-file), not only for the
fp8 / fp4 / kitchen paths.

Issue tracking

Linked issue: #5641

Problem

set_is_first_microbatch() only propagates the flag when config.fp8,
config.fp4, or use_kitchen is set. When quantization is configured
solely through --te-precision-config-file (which sets config.quant_recipe
while leaving fp8/fp4 as None and use_kitchen as False), the guard
never enters, so per-module is_first_microbatch is never set to True.

As a result TE modules never refresh their fp8 weight (transpose) cache on
the first micro-batch of each step under per-module quantization — the same
root cause as expert-only fp8.

Fix

Add quant_recipe is not None to the guard. The flag is a no-op for modules
that run in bf16 (TE ignores it outside the fp8 path), so this is safe for
mixed bf16/fp8 recipes and mirrors the existing global fp8 behavior.

Contribution process

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.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Connor-XY]

Thanks—the fix direction looks correct and is appropriately scoped. Before approval, please add a focused regression test covering fp8=None, fp4=None, use_kitchen=False, and a non-None quant_recipe, asserting that a child module exposing is_first_microbatch is updated. Please also retain coverage showing that the flag is not touched when no quantization mode is configured. Once that is added, we can start CI.

[yezhengmao1]

Thanks for the review! I've pushed a focused regression test (tests/unit_tests/transformer/test_module.py, class TestSetIsFirstMicrobatch) covering both cases you asked for. @Connor-XY