add kv cache quantization for mcore using bmm_quantizers

Signed-off-by: Kai Xu <[email protected]>

Author: kaix-nv

modelopt/torch/quantization/plugins/__init__.py

@@ -73,4 +73,4 @@
     from .trl import *
 
 with import_plugin("mcore"):
-    from .mcore import *
+    from .mcore import *

modelopt/torch/quantization/plugins/megatron.py

@@ -22,23 +22,23 @@
 import megatron.core.tensor_parallel.layers as megatron_parallel
 import megatron.core.transformer.mlp as megatron_mlp
 import torch
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
 from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
 from megatron.core.transformer import MegatronModule
 from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
 from megatron.core.utils import get_tensor_model_parallel_group_if_none
-from megatron.core.extensions.transformer_engine import TEDotProductAttention
 
 from modelopt.torch.opt.plugins.megatron import (
     _MegatronMLP,
     register_modelopt_extra_state_callbacks,
 )
 from modelopt.torch.utils.distributed import ParallelState
 
+from ..model_calib import max_calibrate
 from ..nn import QuantModule, QuantModuleRegistry, TensorQuantizer
 from ..nn.modules.quant_linear import RealQuantLinear
 from ..qtensor import QTensorWrapper
 from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear
-from ..model_calib import max_calibrate
 
 __all__ = []
 
@@ -468,7 +468,13 @@ def forward(self, input, *args, **kwargs):
 
 @QuantModuleRegistry.register({TEDotProductAttention: "TEDotProductAttention"})
 class _QuantTEDotProductAttention(QuantModule):
-    """Quantized version of TEDotProductAttention for Megatron models with KV cache quantization."""
+    """Quantized version of TEDotProductAttention for Megatron models with KV cache quantization.
+
+    This class adds KV cache quantization support to Transformer Engine's TEDotProductAttention
+    module used in Megatron-Core models. It introduces three quantizers (q_bmm_quantizer,
+    k_bmm_quantizer, v_bmm_quantizer) that quantize the query, key, and value tensors after
+    RoPE has been applied.
+    """
 
     def _setup(self):
         """Initialize quantizers for Q, K, V tensors."""
@@ -478,25 +484,35 @@ def _setup(self):
 
     def _calibrate_quantizers(self):
         """Calibrate quantizers with minimal dummy tensors."""
-        # Get device from parent module parameters
-        device = next(self.parameters()).device if self.parameters() else torch.device('cuda')
-        
+        # Get device and dtype from the parent module's parameters
+        param = next(iter(self.parameters()), None)
+        device = param.device if param is not None else torch.device("cuda")
+        dtype = param.dtype if param is not None else torch.float16
+
         # TEDotProductAttention expects format 'sbhd' or 'bshd' depending on rope_fusion
         batch_size = 1
         seq_len = 1
-        
+
         # Get dimensions from config
         num_heads = self.config.num_attention_heads
-        head_dim = self.config.kv_channels if hasattr(self.config, 'kv_channels') else self.config.hidden_size // num_heads
-        
+        head_dim = (
+            self.config.kv_channels
+            if hasattr(self.config, "kv_channels")
+            else self.config.hidden_size // num_heads
+        )
+
         # Determine tensor format (default to sbhd if not specified)
-        apply_rope_fusion = getattr(self.config, 'apply_rope_fusion', False)
+        apply_rope_fusion = getattr(self.config, "apply_rope_fusion", False)
         qkv_format = "bshd" if apply_rope_fusion else "sbhd"
 
         if qkv_format == "sbhd":
-            dummy_tensor = torch.randn(seq_len, batch_size, num_heads, head_dim, device=device, dtype=torch.float16)
+            dummy_tensor = torch.randn(
+                seq_len, batch_size, num_heads, head_dim, device=device, dtype=dtype
+            )
         else:
-            dummy_tensor = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=torch.float16)
+            dummy_tensor = torch.randn(
+                batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype
+            )
 
         # Calibrate each quantizer
         quantizers = [
@@ -506,14 +522,14 @@ def _calibrate_quantizers(self):
         ]
 
         for _, quantizer in quantizers:
-            if quantizer is not None and quantizer.is_enabled:
+            if quantizer is not None and quantizer.is_enabled():
                 if not hasattr(quantizer, "_amax") or quantizer._amax is None:
                     quantizer.reset_amax()
                     max_calibrate(quantizer, lambda q: q(dummy_tensor), distributed_sync=False)
 
     def forward(self, query, key, value, *args, **kwargs):
         """Apply post-RoPE quantization to KV cache.
-        
+
         TEDotProductAttention receives Q, K, V after RoPE is applied,
         so we quantize them directly for KV cache quantization.
         """
@@ -527,7 +543,7 @@ def forward(self, query, key, value, *args, **kwargs):
     def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
         """Create a sharded state dictionary for distributed checkpointing."""
         sharded_state_dict = {}
-        
+
         # First add non-quantizer parameters
         for k, v in self.state_dict(prefix="", keep_vars=True).items():
             if isinstance(v, torch.Tensor) and v is not None and "_quantizer" not in k:
@@ -544,10 +560,11 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
                 sharded_state_dict[amax_key] = quantizer._amax
 
         # Process other quantizer parameters in bmm_quantizers
-        quantizer_state_dict = {}
-        for k, v in self.state_dict(prefix="", keep_vars=True).items():
-            if isinstance(v, torch.Tensor) and "_quantizer" in k and "_amax" not in k:
-                quantizer_state_dict[k] = v
+        quantizer_state_dict = {
+            k: v
+            for k, v in self.state_dict(prefix="", keep_vars=True).items()
+            if isinstance(v, torch.Tensor) and "_quantizer" in k and "_amax" not in k
+        }
 
         if quantizer_state_dict:
             sharded_state_dict.update(
@@ -583,10 +600,11 @@ def modelopt_post_restore(self, name=""):
         super().modelopt_post_restore(name)
 
         def _check_unsupported_states(quantizer):
+            """Check for unsupported quantizer states and warn if found."""
             if not hasattr(quantizer, "state_dict"):
                 return
 
-            for k in quantizer.state_dict().keys():
+            for k in quantizer.state_dict():
                 if k not in ["_amax", "_pre_quant_scale"]:
                     warnings.warn(
                         f"Restore of {k} for {name} is not supported. The restore of this layer might be "
@@ -600,7 +618,7 @@ def _check_unsupported_states(quantizer):
             ("k_bmm_quantizer", self.k_bmm_quantizer),
             ("v_bmm_quantizer", self.v_bmm_quantizer),
         ]:
-            if not hasattr(self, quantizer_name) or not quantizer.is_enabled:
+            if not hasattr(self, quantizer_name) or not quantizer.is_enabled():
                 continue
 
             _check_unsupported_states(quantizer)

tests/gpu/torch/quantization/plugins/test_megatron.py

@@ -17,7 +17,6 @@
 
 import pytest
 import torch
-import torch.nn as nn
 from _test_utils.import_helper import skip_if_no_megatron
 from _test_utils.torch_dist.dist_utils import spawn_multiprocess_job
 from _test_utils.torch_dist.plugins.megatron_common import (
@@ -368,8 +367,10 @@ def test_fp8_real_quantize():
 
 def _test_kv_cache_quant_helper(config, rank, size):
     """Helper function for testing KV cache quantization with TEDotProductAttention."""
-    initialize_for_megatron(tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED)
-    
+    initialize_for_megatron(
+        tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED
+    )
+
     # Use existing infrastructure to create a minimal GPT model with TEDotProductAttention
     # Note: transformer_impl must be "modelopt" or "transformer_engine" (not "local") to get TEDotProductAttention
     model = get_mcore_gpt_model(
@@ -380,43 +381,45 @@ def _test_kv_cache_quant_helper(config, rank, size):
         vocab_size=32,
         transformer_impl="modelopt",  # This uses TEDotProductAttention via get_gpt_modelopt_spec
     ).cuda()
-    
+
     # Create dummy input for calibration
     prompt_tokens = torch.randint(0, model.vocab_size, (2, model.max_sequence_length)).cuda()
-    
+
     def forward_fn(model):
         return megatron_prefill(model, prompt_tokens)
-    
+
     # Test KV cache quantization with the given config
     quantized_model = mtq.quantize(model, config, forward_fn)
-    
+
     # Find TEDotProductAttention modules and verify they have KV cache quantizers
     te_attention_found = False
     for name, module in quantized_model.named_modules():
         # Check if this is a quantized TEDotProductAttention
-        if hasattr(module, 'q_bmm_quantizer') and hasattr(module, 'k_bmm_quantizer'):
+        if hasattr(module, "q_bmm_quantizer") and hasattr(module, "k_bmm_quantizer"):
             te_attention_found = True
             # Verify all expected quantizers exist
-            assert hasattr(module, 'v_bmm_quantizer'), f"Missing v_bmm_quantizer in {name}"
-            
+            assert hasattr(module, "v_bmm_quantizer"), f"Missing v_bmm_quantizer in {name}"
+
             # Verify K and V quantizers are enabled (main purpose of KV cache configs)
             assert module.k_bmm_quantizer.is_enabled, f"K quantizer not enabled in {name}"
             assert module.v_bmm_quantizer.is_enabled, f"V quantizer not enabled in {name}"
-    
+
     assert te_attention_found, "No TEDotProductAttention with KV cache quantizers found in model"
-    
+
     # Quick smoke test that forward still works
     output = forward_fn(quantized_model)
     assert output is not None, "Forward pass failed"
-    
+
 
 def _test_kv_cache_sharded_state_dict_helper(tmp_path, config, rank, size):
     """Helper for testing KV cache quantization with sharded state dict save/load."""
     # Disable output_layer quantization (same as other sharded state dict tests)
     config["quant_cfg"]["*output_layer*"] = {"enable": False}
-    
-    initialize_for_megatron(tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED)
-    
+
+    initialize_for_megatron(
+        tensor_model_parallel_size=size, pipeline_model_parallel_size=1, seed=SEED
+    )
+
     # Create GPT models with TEDotProductAttention (transformer_impl="modelopt")
     model_ref = get_mcore_gpt_model(
         tensor_model_parallel_size=size,
@@ -426,7 +429,7 @@ def _test_kv_cache_sharded_state_dict_helper(tmp_path, config, rank, size):
         vocab_size=64,
         transformer_impl="modelopt",  # CRITICAL: Use TEDotProductAttention
     ).cuda()
-    
+
     model_test = get_mcore_gpt_model(
         tensor_model_parallel_size=size,
         num_layers=2,
@@ -435,29 +438,31 @@ def _test_kv_cache_sharded_state_dict_helper(tmp_path, config, rank, size):
         vocab_size=64,
         transformer_impl="modelopt",
     ).cuda()
-    
-    prompt_tokens = torch.randint(0, model_ref.vocab_size, (2, model_ref.max_sequence_length)).cuda()
-    
+
+    prompt_tokens = torch.randint(
+        0, model_ref.vocab_size, (2, model_ref.max_sequence_length)
+    ).cuda()
+
     def forward_fn(model):
         return megatron_prefill(model, prompt_tokens)
-    
+
     # Quantize the reference model
     model_ref = mtq.quantize(model_ref, config, forward_fn)
-    
+
     # CRITICAL: model_test must also be quantized with the same config
     # Otherwise it won't have the KV cache quantizer keys when loading state dict
     model_test = mtq.quantize(model_test, config, forward_fn)
-    
+
     # Verify KV cache quantizers were created
     kv_quantizers_found = False
     for name, module in model_ref.named_modules():
-        if hasattr(module, 'k_bmm_quantizer') and hasattr(module, 'v_bmm_quantizer'):
+        if hasattr(module, "k_bmm_quantizer") and hasattr(module, "v_bmm_quantizer"):
             kv_quantizers_found = True
             assert module.k_bmm_quantizer.is_enabled, f"K quantizer not enabled in {name}"
             assert module.v_bmm_quantizer.is_enabled, f"V quantizer not enabled in {name}"
-    
+
     assert kv_quantizers_found, "No KV cache quantizers found in quantized model"
-    
+
     # Test sharded state dict save/load
     sharded_state_dict_test_helper(
         tmp_path,
@@ -467,32 +472,38 @@ def forward_fn(model):
         meta_device=False,
         version=None,
     )
-    
+
     # Verify KV cache quantizers are restored correctly in model_test
     for (name_ref, module_ref), (name_test, module_test) in zip(
         model_ref.named_modules(), model_test.named_modules()
     ):
-        if hasattr(module_ref, 'k_bmm_quantizer'):
-            assert hasattr(module_test, 'k_bmm_quantizer'), f"K quantizer missing after restore in {name_test}"
-            assert hasattr(module_test, 'v_bmm_quantizer'), f"V quantizer missing after restore in {name_test}"
-            
+        if hasattr(module_ref, "k_bmm_quantizer"):
+            assert hasattr(module_test, "k_bmm_quantizer"), (
+                f"K quantizer missing after restore in {name_test}"
+            )
+            assert hasattr(module_test, "v_bmm_quantizer"), (
+                f"V quantizer missing after restore in {name_test}"
+            )
+
             # Check that quantizer states match
-            if hasattr(module_ref.k_bmm_quantizer, '_amax'):
-                assert hasattr(module_test.k_bmm_quantizer, '_amax'), f"K quantizer _amax missing in {name_test}"
+            if hasattr(module_ref.k_bmm_quantizer, "_amax"):
+                assert hasattr(module_test.k_bmm_quantizer, "_amax"), (
+                    f"K quantizer _amax missing in {name_test}"
+                )
                 if module_ref.k_bmm_quantizer._amax is not None:
                     assert torch.allclose(
-                        module_ref.k_bmm_quantizer._amax, 
-                        module_test.k_bmm_quantizer._amax
+                        module_ref.k_bmm_quantizer._amax, module_test.k_bmm_quantizer._amax
                     ), f"K quantizer _amax mismatch in {name_test}"
-            
-            if hasattr(module_ref.v_bmm_quantizer, '_amax'):
-                assert hasattr(module_test.v_bmm_quantizer, '_amax'), f"V quantizer _amax missing in {name_test}"
+
+            if hasattr(module_ref.v_bmm_quantizer, "_amax"):
+                assert hasattr(module_test.v_bmm_quantizer, "_amax"), (
+                    f"V quantizer _amax missing in {name_test}"
+                )
                 if module_ref.v_bmm_quantizer._amax is not None:
                     assert torch.allclose(
-                        module_ref.v_bmm_quantizer._amax,
-                        module_test.v_bmm_quantizer._amax
+                        module_ref.v_bmm_quantizer._amax, module_test.v_bmm_quantizer._amax
                     ), f"V quantizer _amax mismatch in {name_test}"
-            
+
 
 @pytest.mark.parametrize(
     "config",
@@ -503,16 +514,14 @@ def forward_fn(model):
 )
 def test_kv_cache_quant(config):
     """Verify KV cache quantization works correctly with TEDotProductAttention.
-    
-    This test ensures TEDotProductAttention is properly registered and gets the 
+
+    This test ensures TEDotProductAttention is properly registered and gets the
     expected q/k/v_bmm_quantizers when using KV cache configs.
-    
+
     Note: This test requires Transformer Engine to be installed since TEDotProductAttention
     is only available with transformer_impl="modelopt" or "transformer_engine" (not "local").
     """
-    spawn_multiprocess_job(
-        size=1, job=partial(_test_kv_cache_quant_helper, config), backend="nccl"
-    )
+    spawn_multiprocess_job(size=1, job=partial(_test_kv_cache_quant_helper, config), backend="nccl")
 
 
 @pytest.mark.parametrize(
@@ -524,7 +533,7 @@ def test_kv_cache_quant(config):
 )
 def test_kv_cache_sharded_state_dict(tmp_path, config):
     """Test KV cache quantization with sharded state dict save/load.
-    
+
     This test verifies the complete workflow of saving and loading KV cache quantized
     models with distributed checkpointing, ensuring quantizer states are properly
     preserved across the save/load cycle.
@@ -533,5 +542,5 @@ def test_kv_cache_sharded_state_dict(tmp_path, config):
     spawn_multiprocess_job(
         size=size,
         job=partial(_test_kv_cache_sharded_state_dict_helper, tmp_path, config),
-        backend="nccl"
+        backend="nccl",
     )