Set is_first_microbatch when quant_recipe is configured

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 = []

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