[BUG] TE FP8 cache not refreshed under per-module quant_recipe

[yezhengmao1]

When quantization is configured only through --te-precision-config-file (which sets config.quant_recipe), the per-module is_first_microbatch flag is never set to True, so Transformer Engine modules never refresh their FP8 weight (transpose) cache on the first micro-batch of each step.

MegatronModule.set_is_first_microbatch() only enters its body when config.fp8, config.fp4, or use_kitchen is truthy.

The --te-precision-config-file path (megatron/training/argument_utils.py:362) sets only config.quant_recipe, leaving fp8/fp4 as None and use_kitchen as False, so the guard is skipped even though FP8 is active per-module.

Steps/Code to reproduce bug

1. Configure FP8 quantization via a recipe file only (no --fp8-format, no kitchen):

   python pretrain_gpt.py ... --te-precision-config-file recipe.yaml

   where recipe.yaml enables FP8 on some modules.

2. Inspect any TE linear layer's is_first_microbatch attribute during training — it stays False on every step.

Expected behavior

With a quant_recipe active, set_is_first_microbatch() should set is_first_microbatch = True on the first micro-batch of each step, exactly as it does for the global --fp8 path, so TE refreshes its FP8 parameter cache correctly.

Additional context

Root cause is the guard condition in set_is_first_microbatch() not covering config.quant_recipe. Same class of bug as expert-only FP8.

Fix: add getattr(self.config, 'quant_recipe', None) is not None to the guard. The flag is a no-op for bf16 modules (TE ignores it outside the FP8 path), so it is safe for mixed bf16/FP8 recipes.

@NVIDIA/mcore-oncall

PR: #5642

[Connor-XY]

Thanks for the clear report. The root cause matches the current code: --te-precision-config-file populates config.quant_recipe, while set_is_first_microbatch() currently only checks fp8, fp4, and use_kitchen. Tracking the fix in #5642.