From 97e442d87cf2b1cf3b0d524be4d0412c885c86ec Mon Sep 17 00:00:00 2001 From: yezhengmao Date: Fri, 3 Jul 2026 16:16:25 +0800 Subject: [PATCH 1/2] Set is_first_microbatch when quant_recipe is configured Signed-off-by: yezhengmao --- megatron/core/transformer/module.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index bae3c70cf9c..6a55bfbc348 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -106,11 +106,13 @@ def set_is_first_microbatch(self): """Sets the is_first_microbatch flag if it exists and config.fp8==True. When this flag is set, TE modules will update their fp8 parameter cache. 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 not hasattr(self, "modules_with_is_first_microbatch"): self.modules_with_is_first_microbatch = [] From c25948d010900704b612c66c82128b93af7c67e8 Mon Sep 17 00:00:00 2001 From: yezhengmao Date: Mon, 6 Jul 2026 15:33:34 +0800 Subject: [PATCH 2/2] Add regression tests for is_first_microbatch with quant_recipe Signed-off-by: yezhengmao --- tests/unit_tests/transformer/test_module.py | 59 +++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/unit_tests/transformer/test_module.py b/tests/unit_tests/transformer/test_module.py index 73b0235f474..92f15b2f46d 100644 --- a/tests/unit_tests/transformer/test_module.py +++ b/tests/unit_tests/transformer/test_module.py @@ -57,6 +57,65 @@ def test_megatron_module(self): # failed_module.bf16 = True +class _FirstMicrobatchModule(torch.nn.Module): + """Stand-in for a TE module that exposes the is_first_microbatch flag.""" + + def __init__(self): + super().__init__() + self.is_first_microbatch = False + + +class DummyQuantModule(MegatronModule): + def __init__(self, config: TransformerConfig): + super().__init__(config) + self.child = _FirstMicrobatchModule() + + def forward(self, x): + return x + + +class TestSetIsFirstMicrobatch: + + def setup_method(self, method): + Utils.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(123) + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + def _build_module(self, **overrides): + config = TransformerConfig( + num_layers=2, hidden_size=12, num_attention_heads=4, use_cpu_initialization=True + ) + for key, value in overrides.items(): + setattr(config, key, value) + return DummyQuantModule(config=config) + + def test_quant_recipe_sets_flag(self): + # quant_recipe alone must enable the flag, even with fp8/fp4/kitchen off. + module = self._build_module(quant_recipe=object()) + assert module.config.fp8 is None + assert module.config.fp4 is None + assert getattr(module.config, 'use_kitchen', False) is False + assert module.config.quant_recipe is not None + assert module.child.is_first_microbatch is False + + module.set_is_first_microbatch() + assert module.child.is_first_microbatch is True + + def test_no_quant_leaves_flag_untouched(self): + # With no quantization mode configured the flag must not be touched. + module = self._build_module() + assert module.config.fp8 is None + assert module.config.fp4 is None + assert getattr(module.config, 'use_kitchen', False) is False + assert module.config.quant_recipe is None + assert module.child.is_first_microbatch is False + + module.set_is_first_microbatch() + assert module.child.is_first_microbatch is False + + class TestFloat16Module: def setup_method(self, method):