diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 55abafa78aa..176f039f923 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- The functional quantized-weight export API now exports the KV cache: ``capture_kv_cache_export_state`` returns an attention module's ``k_scale`` / ``v_scale``, and ``build_hf_quantization_config`` takes the KV-cache formats from ``get_kv_cache_export_spec``. - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections. - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time. @@ -44,6 +45,8 @@ Changelog **Bug Fixes** +- Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges. +- ``export_mcore_gpt_to_hf`` and ``capture_kv_cache_export_state`` now write NVFP4 KV-cache ``k_scale`` / ``v_scale`` as amax / 448, like the HF exporter and as vLLM and TensorRT-LLM expect. They wrote amax / (6 * 448), which overflows in the attention kernels. Re-export Megatron checkpoints with an NVFP4 KV cache. - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. - ``mtq.quantize`` and ``mtq.calibrate`` now give every module back the training mode it had. They used to leave every module in the mode of the model passed in, so a module in eval mode inside a model in training mode came back in training mode. - Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file. diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 2e73258f9cd..5fb6f32ddea 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -315,10 +315,13 @@ def get_kv_cache_scaling_factor( ): return [None, None] - scaling_factors = [ - get_scaling_factor(getattr(self_attention_module, quantizer)) - for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer") - ] + # Both FP8 and NVFP4 cache scales are amax / 448, as the HF exporter writes them and + # vLLM / TensorRT-LLM read them. + scaling_factors = [] + for quantizer_name in ("k_bmm_quantizer", "v_bmm_quantizer"): + quantizer = getattr(self_attention_module, quantizer_name) + amax = quantizer.export_amax() if quantizer.is_enabled else None + scaling_factors.append(None if amax is None else amax.float() / 448.0) # For FP8, we recommend default KV-cache scaling factor to be 1. The # asymmetric format applies this only to K; V remains NVFP4. diff --git a/modelopt/torch/export/quantized_weight_export.py b/modelopt/torch/export/quantized_weight_export.py index 79a29e7c400..c3e70624d90 100644 --- a/modelopt/torch/export/quantized_weight_export.py +++ b/modelopt/torch/export/quantized_weight_export.py @@ -37,6 +37,9 @@ from ..quantization.nn import SequentialQuantizer, TensorQuantizer from .convert_hf_config import convert_hf_quant_config_format from .quant_format import ( + KV_CACHE_FP8, + KV_CACHE_FP8_K_NVFP4_V, + KV_CACHE_NVFP4, QUANTIZATION_FP8, QUANTIZATION_FP8_PB_WO, QUANTIZATION_FP8_PC_PT, @@ -51,6 +54,8 @@ ) from .quant_utils import ( _get_quantization_from_quantizers, + get_kv_cache_dtype, + get_kv_cache_scaling_factor, get_scaling_factor, process_layer_quant_config, to_quantized_weight, @@ -58,8 +63,10 @@ __all__ = [ "build_hf_quantization_config", + "capture_kv_cache_export_state", "capture_quantized_weight_export_state", "export_quantized_weight_tensors", + "get_kv_cache_export_spec", "get_quantized_weight_export_spec", "merge_quantized_weight_export_states", "restore_quantized_weight_export_state", @@ -83,6 +90,7 @@ QUANTIZATION_MXFP8, } ) +_KV_CACHE_EXPORT_FORMATS = {KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4} _STATIC_INPUT_SCALE_FORMATS = { QUANTIZATION_FP8, QUANTIZATION_NVFP4, @@ -126,6 +134,15 @@ class _QuantizedWeightExportSpec: input_quantizer_enabled: bool +@dataclass(frozen=True) +class _KVCacheExportState: + """Detached K and V cache scales of one attention module.""" + + quantization_format: str + k_scale: torch.Tensor + v_scale: torch.Tensor + + @dataclass(frozen=True) class _ExportStateTensorMetadata: name: str @@ -497,6 +514,56 @@ def get_quantized_weight_export_spec( ) +def get_kv_cache_export_spec(attention: nn.Module) -> str | None: + """Return the KV-cache format of one attention module, or ``None`` when it is not quantized.""" + quantizers = [getattr(attention, name, None) for name in ("k_bmm_quantizer", "v_bmm_quantizer")] + enabled = [quantizer is not None and quantizer.is_enabled for quantizer in quantizers] + if not any(enabled): + return None + if not all(enabled): + raise _UnsupportedQuantizedWeightExportFormatError( + "Functional export requires both or neither of the K and V cache quantizers" + ) + # The captured state carries no K/V biases. + if any( + getattr(quantizer, attr, None) is not None + for quantizer in quantizers + for attr in ("bias", "bias_value") + ): + raise _UnsupportedQuantizedWeightExportFormatError( + "Functional export does not support affine KV caches" + ) + quantization_format = get_kv_cache_dtype(attention) + if quantization_format not in _KV_CACHE_EXPORT_FORMATS: + raise _UnsupportedQuantizedWeightExportFormatError( + f"Functional export does not support the {quantization_format!r} KV cache" + ) + return quantization_format + + +def capture_kv_cache_export_state( + attention: nn.Module, + *, + clamp_fp8_scales: bool = True, + cpu: bool = True, +) -> _KVCacheExportState | None: + """Capture the ``k_scale`` / ``v_scale`` of one attention module, or ``None`` when not quantized. + + ``clamp_fp8_scales`` raises FP8 cache scales to at least 1.0, as the other exporters do. + """ + quantization_format = get_kv_cache_export_spec(attention) + if quantization_format is None: + return None + k_scale, v_scale = get_kv_cache_scaling_factor(attention, clamp_fp8_scales=clamp_fp8_scales) + if k_scale is None or v_scale is None: + raise RuntimeError("Missing calibrated KV-cache scales") + if cpu: + k_scale, v_scale = k_scale.cpu(), v_scale.cpu() + return _KVCacheExportState( + quantization_format, k_scale.detach().clone(), v_scale.detach().clone() + ) + + def split_quantized_weight_export_state( state: _QuantizedWeightExportState, ) -> tuple[object, tuple[torch.Tensor, ...]]: @@ -844,8 +911,19 @@ def _quantized_layer_name(weight_name: str) -> str: def build_hf_quantization_config( named_states: Mapping[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None] | Iterable[tuple[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None]], + kv_cache_formats: Iterable[str | None] = (), ) -> dict[str, Any]: - """Build canonical HF configuration, rejecting unrepresentable activation settings.""" + """Build canonical HF configuration, rejecting unrepresentable activation settings. + + ``kv_cache_formats`` holds :func:`get_kv_cache_export_spec` of the attention modules; they must + agree, and ``None`` (an unquantized cache) is allowed alongside one format. + """ + kv_cache_format_set = {fmt for fmt in kv_cache_formats if fmt is not None} + if len(kv_cache_format_set) > 1: + raise NotImplementedError( + "Cannot build canonical HF configuration for mixed KV-cache formats " + f"{sorted(kv_cache_format_set)}" + ) layers: dict[str, tuple[str, int] | None] = {} for weight_name, state in dict(named_states).items(): if state is not None: @@ -875,4 +953,6 @@ def build_hf_quantization_config( "producer": {"name": "modelopt", "version": __version__}, "quantization": process_layer_quant_config(layer_config), } + if kv_cache_format_set: + config["quantization"]["kv_cache_quant_algo"] = kv_cache_format_set.pop() return convert_hf_quant_config_format(config) diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 78a6e03a16b..d82f916b248 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -45,7 +45,7 @@ TensorQuantizer, ) from .utils import is_quantized, is_quantized_linear -from .utils.shared_input import SharedWeightGlobalAmaxState +from .utils.shared_input import MLA_DOWN_PROJ_PATTERN, SharedWeightGlobalAmaxState __all__ = [ "preserve_quantizer_attributes_context", @@ -115,11 +115,17 @@ def _restore_shared_quant_state_aliases( """Rebuild shared-state ties before checkpoint tensor values are loaded.""" if not metadata.get("shared_quant_states"): return - # max / mse / local_hessian all carry ``shared_states`` (via _SharedStatesConfig) and use - # the same grouping; resolve the patterns that were in effect at save and rebuild the ties. - patterns = SharedWeightGlobalAmaxState.resolve_patterns( - shared_states=getattr(config, "shared_states", None) - ) + patterns = metadata.get("shared_quant_patterns") + if patterns is None: + # Saved before the patterns were recorded, when the MLA down projections were not grouped: + # tying them now would load both members' ranges into one buffer. + patterns = [ + pattern + for pattern in SharedWeightGlobalAmaxState.resolve_patterns( + shared_states=getattr(config, "shared_states", None) + ) + if pattern != MLA_DOWN_PROJ_PATTERN + ] SharedWeightGlobalAmaxState.restore(model, patterns=patterns) @@ -206,8 +212,12 @@ def update_quantize_metadata( metadata["quantizer_state"] = quantizer_state(model) if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model): metadata["shared_quant_states"] = shared_state_metadata + # Record the patterns calibration grouped with, so restore ties the same groups; ``config`` + # need not be the calibration algorithm's (e.g. LoRA and AutoQuantize pass a default one). + metadata["shared_quant_patterns"] = SharedWeightGlobalAmaxState.grouping_patterns(model) else: metadata.pop("shared_quant_states", None) + metadata.pop("shared_quant_patterns", None) def _linear_attention_modules(model): diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index cf14a5b8c3c..f02732a9334 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -100,9 +100,10 @@ def _is_calibrated_nvfp4_static(q) -> bool: def _collect_grouped_linears(model: nn.Module) -> list[list[nn.Module]]: - """Collect name-based sibling groups (Q/K/V, gate/up, w1/w3) of calibrated NVFP4-static linears.""" + """Collect name-based sibling groups (Q/K/V, MLA down, gate/up, w1/w3) of calibrated NVFP4-static linears.""" patterns: tuple[tuple[str, ...], ...] = ( ("q_proj", "k_proj", "v_proj"), + ("linear_q_down_proj", "linear_kv_down_proj"), *list_all_possible("gate_up_pairs"), ) groups: list[list[nn.Module]] = [] @@ -146,6 +147,41 @@ def _check_grouped_weight_global_amax_synced(model: nn.Module) -> None: ) +@torch.no_grad() +def _sync_mla_fused_down_proj_amax(model: torch.nn.Module) -> None: + """Give Megatron-Core MLA's q_down / kv_down projections one per-tensor input and weight range each. + + vLLM and TRT-LLM fuse them into one GEMM (``fused_qkv_a_proj``) with one scale per tensor. + Static NVFP4 weights share their global range through ``SHARED_PATTERNS`` instead. + """ + for module in model.modules(): + members = [ + getattr(module, name, None) for name in ("linear_q_down_proj", "linear_kv_down_proj") + ] + if any(member is None for member in members): + continue + for quantizer_name in ("input_quantizer", "weight_quantizer"): + quantizers = [] + for member in members: + quantizer = getattr(member, quantizer_name, None) + # W4A8 keeps its per-tensor FP8 range in the last quantizer. + if isinstance(quantizer, SequentialQuantizer): + quantizer = quantizer[-1] + amax = getattr(quantizer, "_amax", None) + if ( + quantizer is None + or not quantizer.is_enabled + or amax is None + or amax.numel() != 1 + ): + break + quantizers.append((quantizer, amax)) + else: + shared = torch.stack([amax.float().reshape(()) for _, amax in quantizers]).max() + for quantizer, amax in quantizers: + quantizer.amax = shared.to(amax).reshape(amax.shape) + + def _finalize_with_shared_state(model: nn.Module, weight_patterns: list[str]) -> None: """Finalize calibrated static quantizers and attached shared state. @@ -155,6 +191,7 @@ def _finalize_with_shared_state(model: nn.Module, weight_patterns: list[str]) -> the default patterns, verifies the name groups were actually synced. Call once ``_amax`` is final: single-process, or after the distributed amax sync. """ + _sync_mla_fused_down_proj_amax(model) SharedWeightGlobalAmaxState.populate(model) promote_static_block_weight_quantizers(model) # Under the default patterns, verify the fusible name groups were actually synced. @@ -833,6 +870,8 @@ def _mse_calibrate_weights( pbar.update(1) pbar.close() + # The search refines each weight's range on its own; share the fused MLA pair's again. + _sync_mla_fused_down_proj_amax(model) class _LocalHessianAccumulator: diff --git a/modelopt/torch/quantization/utils/shared_input.py b/modelopt/torch/quantization/utils/shared_input.py index d80de2fcf9d..1d7e1c2a139 100644 --- a/modelopt/torch/quantization/utils/shared_input.py +++ b/modelopt/torch/quantization/utils/shared_input.py @@ -16,9 +16,10 @@ """Shared quantization state for fusible sibling modules. Weight ``global_amax`` must be unified across modules that get **fused** at export -(q/k/v -> qkv, gate/up -> gate_up) so they quantize with one per-tensor scale. -:func:`find_shared_input_groups` discovers these groups by regex over module FQNs; -``SHARED_PATTERNS`` covers the standard q/k/v, gate/up and per-expert w1/w3 names. +(q/k/v -> qkv, gate/up -> gate_up, MLA q_down/kv_down -> fused_qkv_a_proj) so they quantize +with one per-tensor scale. :func:`find_shared_input_groups` discovers these groups by regex over +module FQNs; ``SHARED_PATTERNS`` covers the standard q/k/v, gate/up, per-expert w1/w3 and MCore +MLA down-projection names. Discovery is name/pattern-based (not input-hook-based) on purpose: "shares an input tensor" is broader than "gets fused" — e.g. a ``shared_expert_gate`` reads the same @@ -52,11 +53,15 @@ "iter_shared_quant_states", ] -# Groups export fuses: q/k/v -> qkv, gate/up (incl. Mixtral w1/w3) -> gate_up. +# Megatron-Core MLA's q_down / kv_down projections, which deployment fuses into fused_qkv_a_proj. +MLA_DOWN_PROJ_PATTERN = r"(?:(.*)\.)?(?:linear_q_down_proj|linear_kv_down_proj)" + +# Groups export fuses: q/k/v -> qkv, MLA q_down/kv_down, gate/up (incl. Mixtral w1/w3) -> gate_up. # Regexes are ``re.fullmatch``-ed against module FQNs; ``(?:(.*)\.)?`` captures the # immediate parent so MoE ``w1``/``w3`` groups are per-expert. SHARED_PATTERNS = ( r"(?:(.*)\.)?(?:q_proj|k_proj|v_proj)", + MLA_DOWN_PROJ_PATTERN, r"(?:(.*)\.)?(?:gate_proj|up_proj)", r"(?:(.*)\.)?(?:w1|w3)", ) @@ -65,6 +70,7 @@ # same ``experts``/``local_experts`` container. Kept unused until a state opts into it. SHARED_PATTERNS_ACROSS_EXPERTS = ( r"(?:(.*)\.)?(?:q_proj|k_proj|v_proj)", + MLA_DOWN_PROJ_PATTERN, r"(?:(.*)\.)?(?:gate_proj|up_proj)", r"(?:(.*)\.)?(?:experts|local_experts)\.\d+\.(?:w1|w3)", ) @@ -103,6 +109,8 @@ def __init__(self) -> None: object.__setattr__(self, "_parent", None) object.__setattr__(self, "_members", ()) object.__setattr__(self, "_hook_handles", []) + # Patterns that formed this group, recorded so restore groups modules the same way. + object.__setattr__(self, "patterns", ()) def set_members(self, parent: nn.Module, members: Sequence[nn.Module]) -> None: """Set the owning parent and linked member modules.""" @@ -235,6 +243,7 @@ def attach(cls, model: nn.Module, patterns: Sequence[str] | None = None) -> int: target_quantizer_kind=cls.target_quantizer_kind, ): state = cls() + object.__setattr__(state, "patterns", tuple(patterns or ())) state.set_members(parent, members) created = _register_parent_shared_state(parent, state) n_created += int(created) @@ -272,6 +281,14 @@ def restore(cls, model: nn.Module, patterns: Sequence[str] | None = None) -> Non for state in iter_shared_quant_states(model, cls): state.restore_from_members() + @classmethod + def grouping_patterns(cls, model: nn.Module) -> list[str]: + """Return the patterns that formed this state's groups in ``model``.""" + patterns: list[str] = [] + for state in iter_shared_quant_states(model, cls): + patterns += [pattern for pattern in state.patterns if pattern not in patterns] + return patterns + @classmethod def metadata(cls, model: nn.Module) -> dict[str, bool]: """Return restore metadata for this state when present in ``model``.""" diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 043970ef42a..c70e5bb432f 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -14,6 +14,7 @@ # limitations under the License. import json +import re from copy import deepcopy from functools import partial from pathlib import Path @@ -26,6 +27,7 @@ from _test_utils.torch.export.unified_checkpoint import ( assert_exported_checkpoint_matches, assert_per_expert_experts_complete, + load_safetensors_dir, ) from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model from _test_utils.torch.megatron.utils import get_forward @@ -49,6 +51,11 @@ from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping from modelopt.torch.export.quant_format import GGML_FORMATS +from modelopt.torch.export.quantized_weight_export import ( + build_hf_quantization_config, + capture_kv_cache_export_state, + get_kv_cache_export_spec, +) from modelopt.torch.export.unified_export_megatron import GPTModelExporter from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import TensorQuantizer @@ -501,6 +508,27 @@ def _verify_exported_metadata(export_dir: Path, model_type, quant_config, extra_ assert generation_config["top_p"] == 0.5 +def _verify_functional_kv_cache_export(model, export_dir: Path): + """The functional KV-cache export writes the same scales and config as this exporter.""" + attentions = [m for m in model.modules() if hasattr(m, "k_bmm_quantizer")] + states = [capture_kv_cache_export_state(attention) for attention in attentions] + assert states and None not in states + exported = load_safetensors_dir(export_dir) + for kind in ("k_scale", "v_scale"): + names = sorted( + (name for name in exported if name.endswith(kind)), + key=lambda name: int(re.search(r"layers\.(\d+)\.", name).group(1)), + ) + assert len(names) == len(states) + for name, state in zip(names, states): + torch.testing.assert_close(exported[name], getattr(state, kind).to(exported[name])) + config = build_hf_quantization_config({}, [get_kv_cache_export_spec(a) for a in attentions]) + quantization_config = json.loads((export_dir / "config.json").read_text())[ + "quantization_config" + ] + assert config["kv_cache_scheme"] == quantization_config["kv_cache_scheme"] + + def _test_unified_export_megatron( tmp_path, model_type, @@ -661,6 +689,8 @@ def _test_unified_export_megatron( if quant_config: _verify_model_quant_config(tmp_export_dir, quant_config, kv_cache_quant_cfg) + if kv_cache_quant_cfg and rank == 0: + _verify_functional_kv_cache_export(model, tmp_export_dir) if rank == 0: _verify_exported_metadata(tmp_export_dir, model_type, quant_config, extra_module) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index c53a87276e6..24bbe809b44 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -2648,3 +2648,38 @@ def unrecognised(self, prefix="", sharded_offsets=(), metadata=None): ) def test_extra_state_has_data(self, entry, expected): assert _output_layer_extra_state_has_data(entry) is expected + + +@pytest.mark.parametrize( + "quant_cfg", ["FP8_DEFAULT_CFG", "NVFP4_DEFAULT_CFG", "NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG"] +) +def test_mla_fused_down_projections_share_ranges(distributed_setup_size_1, quant_cfg): + """vLLM and TRT-LLM fuse MLA's q_down and kv_down into one GEMM with one scale per tensor.""" + initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED) + try: + model = get_mcore_gpt_model( + hidden_size=128, + num_attention_heads=2, + vocab_size=128, + max_sequence_length=32, + transformer_impl="transformer_engine", + multi_latent_attention=True, + ).cuda() + attentions = [layer.self_attention for layer in model.decoder.layers] + with torch.no_grad(): # give the two weights different ranges + for attention in attentions: + attention.linear_kv_down_proj.weight.mul_(4) + + model = mtq.quantize(model, copy.deepcopy(getattr(mtq, quant_cfg)), get_forward(model)) + + for attention in attentions: + q_down, kv_down = attention.linear_q_down_proj, attention.linear_kv_down_proj + assert torch.equal(q_down.input_quantizer.amax, kv_down.input_quantizer.amax) + if q_down.weight_quantizer.is_nvfp4_static: + assert torch.equal( + q_down.weight_quantizer.global_amax, kv_down.weight_quantizer.global_amax + ) + else: + assert torch.equal(q_down.weight_quantizer.amax, kv_down.weight_quantizer.amax) + finally: + destroy_model_parallel() # not a spawned process diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py index ef0734aff96..73e92734954 100644 --- a/tests/unit/torch/export/test_quantized_weight_state.py +++ b/tests/unit/torch/export/test_quantized_weight_state.py @@ -24,8 +24,10 @@ from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.quantized_weight_export import ( build_hf_quantization_config, + capture_kv_cache_export_state, capture_quantized_weight_export_state, export_quantized_weight_tensors, + get_kv_cache_export_spec, get_quantized_weight_export_spec, merge_quantized_weight_export_states, restore_quantized_weight_export_state, @@ -713,3 +715,77 @@ def test_mixed_config_groups_moe_experts_by_projection_family(): "model.layers.0.mlp.experts.down_proj", } assert "lm_head" in config["ignore"] + + +_FP8_KV = {"num_bits": (4, 3)} +_NVFP4_KV = {"num_bits": (2, 1), "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}} + + +def _kv_attention(k_cfg, v_cfg, k_amax=896.0, v_amax=224.0) -> nn.Module: + attention = nn.Module() + for name, cfg, amax in (("k_bmm_quantizer", k_cfg, k_amax), ("v_bmm_quantizer", v_cfg, v_amax)): + quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) + quantizer.amax = torch.tensor(amax) + setattr(attention, name, quantizer) + return attention + + +@pytest.mark.parametrize( + ("k_cfg", "v_cfg", "quantization_format", "k_scale", "v_scale"), + [ + # Scales are amax / 448; FP8 ones are raised to at least 1.0 like the other exporters do. + (_FP8_KV, _FP8_KV, "FP8", 2.0, 1.0), + (_NVFP4_KV, _NVFP4_KV, "NVFP4", 2.0, 0.5), + (_FP8_KV, _NVFP4_KV, "FP8_K_NVFP4_V", 2.0, 0.5), + ], +) +def test_kv_cache_capture(k_cfg, v_cfg, quantization_format, k_scale, v_scale): + attention = _kv_attention(k_cfg, v_cfg) + + state = capture_kv_cache_export_state(attention) + + assert get_kv_cache_export_spec(attention) == state.quantization_format == quantization_format + torch.testing.assert_close(state.k_scale.squeeze(), torch.tensor(k_scale)) + torch.testing.assert_close(state.v_scale.squeeze(), torch.tensor(v_scale)) + assert not state.k_scale.requires_grad and state.k_scale.device.type == "cpu" + if quantization_format == "FP8": + unclamped = capture_kv_cache_export_state(attention, clamp_fp8_scales=False) + assert unclamped.v_scale.item() == 0.5 + + +def test_kv_cache_spec_rejects_what_it_cannot_export(): + attention = _kv_attention(_FP8_KV, _FP8_KV) + attention.k_bmm_quantizer.disable() + attention.v_bmm_quantizer.disable() + assert get_kv_cache_export_spec(attention) is None + assert capture_kv_cache_export_state(attention) is None + assert get_kv_cache_export_spec(nn.Linear(2, 2)) is None + + attention.k_bmm_quantizer.enable() + with pytest.raises(NotImplementedError, match="both or neither"): + get_kv_cache_export_spec(attention) + + attention = _kv_attention({"num_bits": 8}, {"num_bits": 8}) + with pytest.raises(NotImplementedError, match="'INT8' KV cache"): + get_kv_cache_export_spec(attention) + + affine_fp8 = {**_FP8_KV, "bias": {-2: None, -4: None, "type": "static"}} + attention = _kv_attention(affine_fp8, affine_fp8) + with pytest.raises(NotImplementedError, match="affine KV caches"): + get_kv_cache_export_spec(attention) + + +@pytest.mark.parametrize( + "weight_spec", [None, weight_export._QuantizedWeightExportSpec("fp8", 0, True)] +) +def test_hf_config_carries_the_kv_cache_format(weight_spec): + named_specs = {"model.layers.0.proj.weight": weight_spec} + + config = build_hf_quantization_config(named_specs, kv_cache_formats=["FP8", None, "FP8"]) + + assert config["kv_cache_scheme"] == {"dynamic": False, "num_bits": 8, "type": "float"} + assert "kv_cache_scheme" not in build_hf_quantization_config(named_specs) + config = build_hf_quantization_config(named_specs, kv_cache_formats=["FP8_K_NVFP4_V"]) + assert config["kv_cache_quant_algo"] == "FP8_K_NVFP4_V" + with pytest.raises(NotImplementedError, match="mixed KV-cache formats"): + build_hf_quantization_config(named_specs, kv_cache_formats=["FP8", "NVFP4"]) diff --git a/tests/unit/torch/quantization/test_shared_input.py b/tests/unit/torch/quantization/test_shared_input.py index 63933a819ff..9ea01e64dbe 100644 --- a/tests/unit/torch/quantization/test_shared_input.py +++ b/tests/unit/torch/quantization/test_shared_input.py @@ -26,9 +26,18 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -from modelopt.torch.quantization.config import MaxCalibConfig, QuantizerAttributeConfig -from modelopt.torch.quantization.model_calib import max_calibrate -from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer +from modelopt.torch.quantization.config import ( + MaxCalibConfig, + QuantizeConfig, + QuantizerAttributeConfig, +) +from modelopt.torch.quantization.conversion import update_quantize_metadata +from modelopt.torch.quantization.model_calib import _sync_mla_fused_down_proj_amax, max_calibrate +from modelopt.torch.quantization.nn import ( + NVFP4StaticQuantizer, + SequentialQuantizer, + TensorQuantizer, +) from modelopt.torch.quantization.utils import ( SHARED_PATTERNS, SharedWeightGlobalAmaxState, @@ -507,3 +516,119 @@ def test_config_rejects_invalid_shared_states(self): MaxCalibConfig(shared_states={"weight_global_amax": {"pattern": [r".*"]}}) with pytest.raises(ValueError, match="invalid regex"): MaxCalibConfig(shared_states={"weight_global_amax": {"patterns": ["("]}}) + + +class _DummyMLA(nn.Module): + """Megatron-Core MLA's q_down / kv_down projections, which deployment fuses into one GEMM.""" + + def __init__(self) -> None: + super().__init__() + self.linear_q_down_proj = nn.Linear(32, 32, bias=False) + self.linear_kv_down_proj = nn.Linear(32, 32, bias=False) + with torch.no_grad(): + self.linear_kv_down_proj.weight.mul_(4) # give the two weights different ranges + + def forward(self, x): + return self.linear_q_down_proj(x) + self.linear_kv_down_proj(x) + + +_PER_TENSOR_INT8 = {"num_bits": 8, "axis": None} +_STATIC_NVFP4 = { + "num_bits": (2, 1), + "block_sizes": {-1: NVFP4_BLOCK, "type": "static", "scale_bits": (4, 3)}, +} + + +def _quantize_mla(weight_cfg, algorithm, input_cfg=None): + quant_cfg = [ + {"enable": False, "quantizer_name": "*"}, + {"cfg": weight_cfg, "quantizer_name": "*weight_quantizer"}, + ] + if input_cfg is not None: + quant_cfg.append({"cfg": input_cfg, "quantizer_name": "*input_quantizer"}) + mla = _DummyMLA() + mtq.quantize( + mla, {"quant_cfg": quant_cfg, "algorithm": algorithm}, lambda m: m(torch.randn(4, 32)) + ) + return mla + + +@pytest.mark.parametrize("algorithm", ["max", {"method": "mse"}]) +def test_mla_down_projections_share_per_tensor_ranges(algorithm): + """The MSE weight search refines each range on its own, so the pair is shared again after it.""" + mla = _quantize_mla(_PER_TENSOR_INT8, algorithm, input_cfg=_PER_TENSOR_INT8) + q_down, kv_down = mla.linear_q_down_proj, mla.linear_kv_down_proj + + assert torch.equal(q_down.input_quantizer.amax, kv_down.input_quantizer.amax) + assert torch.equal(q_down.weight_quantizer.amax, kv_down.weight_quantizer.amax) + assert kv_down.weight_quantizer.amax > q_down.weight.abs().max() + + +def test_mla_down_projections_share_the_fp8_stage_of_sequential_weight_quantizers(): + mla = _DummyMLA() + mtq.replace_quant_module(mla) + first_stages = [] + for scale, proj in enumerate((mla.linear_q_down_proj, mla.linear_kv_down_proj), start=1): + proj.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer()) + proj.weight_quantizer[0].amax = torch.full((32, 1), float(scale)) + proj.weight_quantizer[1].amax = torch.tensor(float(scale)) + proj.input_quantizer.amax = torch.tensor(float(scale)) + first_stages.append(proj.weight_quantizer[0].amax.clone()) + mla.linear_kv_down_proj.input_quantizer.disable() + + _sync_mla_fused_down_proj_amax(mla) + + q_down, kv_down = mla.linear_q_down_proj, mla.linear_kv_down_proj + assert q_down.weight_quantizer[1].amax == kv_down.weight_quantizer[1].amax == 2.0 + for proj, first_stage in zip((q_down, kv_down), first_stages): + assert torch.equal(proj.weight_quantizer[0].amax, first_stage) # per-channel: left alone + assert q_down.input_quantizer.amax == 1.0 # the other side is not quantized + + +def test_mla_down_projections_of_older_checkpoints_restore_ungrouped(tmp_path): + """A checkpoint saved before the MLA pair was grouped keeps each member's own range.""" + mla = _quantize_mla(_STATIC_NVFP4, "max") + shared = mla.linear_q_down_proj.weight_quantizer.global_amax.clone() + assert torch.equal(shared, mla.linear_kv_down_proj.weight_quantizer.global_amax) + path = tmp_path / "model.pth" + mto.save(mla, path) + + restored = _DummyMLA() + mto.restore(restored, path) + _single_global_amax_state(restored) + + # Older checkpoints did not record the patterns; restore must not tie the pair. + checkpoint = torch.load(path, weights_only=False) + for _, mode_state in checkpoint["modelopt_state"]["modelopt_state_dict"]: + mode_state["metadata"].pop("shared_quant_patterns", None) + torch.save(checkpoint, path) + restored = _DummyMLA() + mto.restore(restored, path) + assert not hasattr(restored, "_shared_quant_states") + + +def test_custom_grouping_patterns_survive_save_and_restore(tmp_path): + """Restore groups as calibration did, even where the metadata is written with a default config.""" + q_and_k = r"(?:(.*)\.)?(?:q_proj|k_proj)" + cfg = { + "quant_cfg": [ + {"enable": False, "quantizer_name": "*"}, + {"cfg": _STATIC_NVFP4, "quantizer_name": "*weight_quantizer"}, + ], + "algorithm": { + "method": "max", + "shared_states": {"weight_global_amax": {"patterns": [q_and_k]}}, + }, + } + attn = _DummyAttention() + mtq.quantize(attn, cfg, lambda m: m(torch.randn(2, 32))) + assert _single_global_amax_state(attn).members == (attn.q_proj, attn.k_proj) + metadata = {} + update_quantize_metadata(attn, QuantizeConfig(), metadata) # as LoRA and AutoQuantize do + assert metadata["shared_quant_patterns"] == [q_and_k] + + path = tmp_path / "model.pth" + mto.save(attn, path) + restored = _DummyAttention() + mto.restore(restored, path) + assert _single_global_amax_state(restored).members == (restored.q_proj, restored.k_proj)