diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4580a045f39..a1a2c8a4c41 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,7 @@ Changelog *Quantization* - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize_kv_cache`` and ``constraints.kv_effective_bits``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, and writes a JSON sensitivity report alongside the checkpoint. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included. *Megatron Framework (M-LM / M-Bridge)* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index fea69221825..d77dbd520f4 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -422,8 +422,37 @@ leaving the original recipe unchanged. For models without backprop support (e.g. Llama-4), use the `kl_div` scoring method — see the shipped `general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits` recipe. -KV cache is applied as a uniform post-step, not part of the per-layer search. An AutoQuantize recipe -falls back to `--kv_cache_qformat` (default `fp8_cast`) unless it sets an explicit `kv_cache` field. +Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall back to +`--kv_cache_qformat` (default `fp8_cast`) unless they set an explicit `kv_cache` field. + +KV-cache AutoQuantize recipes instead set `constraints.kv_effective_bits`. Their +`candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes +packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are +preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice. +The shipped canary recipe searches calibrated FP8 K/V (8.0 bits/scalar) and packed NVFP4 K/V +(4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the +companion vLLM implementation does not support that asymmetric per-layer format: + +```bash +python hf_ptq.py \ + --pyt_ckpt_path Qwen/Qwen3-1.7B \ + --recipe general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits \ + --auto_quantize_checkpoint /path/to/kv_autoquant.pth \ + --export_path /path/to/qwen3-1.7b-mixed-kv +``` + +Each candidate uses max calibration so its persistent K/V scales are present in the unified HF +checkpoint. Unified export records the selected formats in `kv_cache_quantized_layers` and writes +the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`; +`--auto_quantize_checkpoint` stores the resumable raw search state. + +> [!NOTE] +> Layer-wise KV checkpoints require the companion +> [vLLM mixed-KV metadata consumer](https://github.com/vllm-project/vllm/pull/52813) or a later +> vLLM release containing it. The repository's currently pinned vLLM 0.26.0 does not consume +> `kv_cache_quantized_layers`, so these checkpoints are export-only in that stock environment. +> Do not deploy them with the pinned runtime. Full FP8 K/V and full NVFP4 K/V use existing vLLM +> kernels once the layer-wise metadata consumer is available. The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an interrupted search (skips re-scoring): @@ -460,6 +489,8 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. +> *AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Use a PTQ recipe with FSDP2.* + ### Usage #### Slurm (recommended) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a56a62b54b5..8a52a78040b 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -96,6 +96,23 @@ from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 +_FSDP2_AUTOQUANT_ERROR = ( + "AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, " + "selection, and checkpoint writes are synchronized across ranks." +) + + +def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor: + """Return logits only for token positions selected by ``attention_mask``.""" + attention_mask = batch.get("attention_mask") + if attention_mask is None: + return logits + if logits.shape[:-1] != attention_mask.shape: + raise ValueError( + "AutoQuant KL logits and attention_mask must have matching token dimensions; " + f"got {tuple(logits.shape[:-1])} and {tuple(attention_mask.shape)}." + ) + return logits[attention_mask.bool()] def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool: @@ -338,6 +355,24 @@ def _mtq_candidate_formats(formats) -> list[dict]: return quantization_formats +def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]: + """Translate format-agnostic KV candidates while preserving useful preset names.""" + candidates = [] + for idx, fmt in enumerate(formats): + quant_cfg = fmt.model_dump(exclude_none=True) + candidate_quantizers = quant_cfg.get("quant_cfg", []) + name = None + for preset_name, preset in KV_QUANT_CFG_CHOICES.items(): + normalized_preset_quantizers = ( + type(fmt)(**preset).model_dump(exclude_none=True).get("quant_cfg", []) + ) + if normalized_preset_quantizers == candidate_quantizers: + name = preset_name + break + candidates.append((quant_cfg, name or f"KV_CACHE_FORMAT_{idx}")) + return candidates + + def _mtq_inputs_from_auto_quantize_config( aq_config, args: argparse.Namespace, fixed_quantize_config=None ) -> dict: @@ -349,6 +384,16 @@ def _mtq_inputs_from_auto_quantize_config( to ``--kv_cache_qformat`` when the recipe omits it. """ constraints = aq_config.constraints.model_dump(exclude_none=True) + is_kv_search = aq_config.constraints.kv_effective_bits is not None + if is_kv_search: + return { + "search_domain": "kv_cache", + "constraints": {"kv_effective_bits": constraints["kv_effective_bits"]}, + "quantization_formats": _mtq_kv_candidate_formats(aq_config.candidate_formats), + "disabled_layers": aq_config.disabled_layers, + "method": aq_config.auto_quantize_method, + "score_size": aq_config.score_size, + } # cost_excluded_layers (sibling of disabled_layers) maps to the mtq cost key: these layers are # kept out of the bit-budget denominator (cost_weight 0) — e.g. VL vision towers — distinct from # disabled_layers, which removes them from the search. @@ -380,6 +425,7 @@ def _mtq_inputs_from_auto_quantize_config( for search_space in aq_config.module_search_spaces ] return { + "search_domain": "weight", "constraints": constraints, "quantization_formats": quantization_formats, "fixed_quantization_config": fixed_quantization_config, @@ -414,11 +460,7 @@ def auto_quantize( ) if args.use_fsdp2: - warnings.warn( - "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet " - "(distributed calibration, sensitivity scoring, and recipe/checkpoint " - "synchronization across ranks); use at your own risk." - ) + raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR) inputs = _mtq_inputs_from_auto_quantize_config( aq_config, args, fixed_quantize_config=fixed_quantize_config @@ -461,14 +503,33 @@ def forward_step(model, batch): output = model(**inputs_) if is_base_model: assert full_model is not None - return full_model.lm_head(output.last_hidden_state) - return output.logits + logits = full_model.lm_head(output.last_hidden_state) + else: + logits = output.logits + return _select_unpadded_logits(logits, batch) else: raise ValueError( f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'" ) + if inputs["search_domain"] == "kv_cache": + language_model, _ = mtq.auto_quantize_kv_cache( + language_model, + constraints=inputs["constraints"], + data_loader=calib_dataloader, + forward_step=forward_step, + quantization_formats=inputs["quantization_formats"], + num_calib_steps=len(calib_dataloader), + num_score_steps=min( + len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1) + ), + verbose=True, + disabled_layers=inputs["disabled_layers"], + checkpoint=args.auto_quantize_checkpoint, + ) + return language_model + language_model, _ = mtq.auto_quantize( language_model, constraints=inputs["constraints"], @@ -512,6 +573,8 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool: def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False + if args.use_fsdp2 and _recipe_is_auto_quantize(args.recipe): + raise NotImplementedError(_FSDP2_AUTOQUANT_ERROR) if args.use_fsdp2: hf_config = AutoConfig.from_pretrained( args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index 2ca8f4f462b..664953a137c 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -149,10 +149,22 @@ def _validate_active_moe_expert_ratio(cls, v: float | None) -> float | None: class AutoQuantizeConstraints(ModeloptBaseConfig): """LP search constraints + cost model; matches the ``mtq.auto_quantize`` constraints dict.""" - effective_bits: float = ModeloptField( - default=4.8, + effective_bits: float | None = ModeloptField( + default=None, title="Effective bits per weight", - description="Average weight-storage bits target for the LP, in (0, 16].", + description=( + "Average weight-storage bits target for the LP, in (0, 16]. Defaults to 4.8 " + "when neither bit constraint is specified." + ), + ) + kv_effective_bits: float | None = ModeloptField( + default=None, + title="Effective bits per KV-cache scalar", + description=( + "Average KV-cache storage bits target across eligible layers for layer-wise KV " + "AutoQuant, in (0, 16]. Exactly one of effective_bits and kv_effective_bits may " + "be set." + ), ) cost_model: Literal["weight", "active_moe"] = ModeloptField( default="weight", @@ -165,13 +177,36 @@ class AutoQuantizeConstraints(ModeloptBaseConfig): description="Extra cost-model parameters; omit for the 'weight' cost model.", ) - @field_validator("effective_bits") + @model_validator(mode="before") @classmethod - def _validate_effective_bits(cls, v: float) -> float: - if not (0 < v <= 16): + def _default_weight_constraint(cls, data): + if isinstance(data, dict): + data = dict(data) + if "effective_bits" not in data and "kv_effective_bits" not in data: + data["effective_bits"] = 4.8 + return data + + @field_validator("effective_bits", "kv_effective_bits") + @classmethod + def _validate_effective_bits(cls, v: float | None) -> float | None: + if v is not None and not (0 < v <= 16): raise ValueError(f"effective_bits must be in (0, 16], got {v}") return v + @model_validator(mode="after") + def _exactly_one_bit_constraint(self): + if (self.effective_bits is None) == (self.kv_effective_bits is None): + raise ValueError( + "Exactly one of effective_bits and kv_effective_bits must be specified." + ) + if self.kv_effective_bits is not None and ( + self.cost_model != "weight" or self.cost is not None + ): + raise ValueError( + "KV-cache AutoQuant does not support weight or active-MoE cost settings." + ) + return self + class AutoQuantizeModuleSearchSpace(ModeloptBaseConfig): """Candidate formats selectable for modules matching one or more name patterns.""" @@ -272,6 +307,20 @@ def _has_search_space(self): "auto_quantize requires candidate_formats or at least one module_search_spaces " "entry. For uniform quantization, use a PTQ recipe instead." ) + if self.constraints.kv_effective_bits is not None: + if self.auto_quantize_method != "kl_div": + raise ValueError( + "KV-cache AutoQuant currently requires auto_quantize_method=kl_div." + ) + if self.module_search_spaces: + raise ValueError( + "KV-cache AutoQuant uses one candidate space for all eligible attention " + "layers; module_search_spaces is not supported." + ) + if self.kv_cache is not None: + raise ValueError( + "KV-cache AutoQuant candidate_formats replace the uniform kv_cache post-step." + ) return self diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 45fa0c30f3b..534aefe3ba8 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -269,6 +269,14 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An if kv_cache_quant_algo: if kv_cache_quant_algo == "FP8": new_config["kv_cache_scheme"] = {"dynamic": False, "num_bits": 8, "type": "float"} + elif kv_cache_quant_algo == "MIXED_PRECISION": + new_config["kv_cache_quant_algo"] = kv_cache_quant_algo + new_config["kv_cache_quantized_layers"] = original_quantization_details.get( + "kv_cache_quantized_layers", {} + ) + new_config["kv_cache_schema_version"] = original_quantization_details.get( + "kv_cache_schema_version", 1 + ) else: # TODO: Handle other kv cache quantization algorithms new_config["kv_cache_scheme"] = kv_cache_quant_algo diff --git a/modelopt/torch/export/model_config.py b/modelopt/torch/export/model_config.py index 5f92cc2e5dc..cf778bf937d 100755 --- a/modelopt/torch/export/model_config.py +++ b/modelopt/torch/export/model_config.py @@ -45,6 +45,7 @@ QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" KV_CACHE_FP8 = "FP8" +KV_CACHE_FP8_K_NVFP4_V = "FP8_K_NVFP4_V" KV_CACHE_INT8 = "INT8" KV_CACHE_NVFP4 = "NVFP4" KV_CACHE_NVFP4_AFFINE = "NVFP4_AFFINE" diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 2e32f123869..fa594f104c4 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -278,7 +278,7 @@ def _map(name: str) -> str: def revert_quant_config_names(quantization: dict, mapper) -> None: - """Revert ``exclude_modules`` / ``quantized_layers`` keys to hub names, in place. + """Revert layer-reference keys to hub names, in place. ``mapper`` is the callable from :func:`build_reverse_name_mapper` (a no-op when ``None``). Applies to the ModelOpt ``{"quantization": {...}}`` sub-dict before it is @@ -293,6 +293,11 @@ def revert_quant_config_names(quantization: dict, mapper) -> None: quantized_layers = quantization.get("quantized_layers") if isinstance(quantized_layers, dict) and quantized_layers: quantization["quantized_layers"] = {mapper(k): v for k, v in quantized_layers.items()} + kv_cache_quantized_layers = quantization.get("kv_cache_quantized_layers") + if isinstance(kv_cache_quantized_layers, dict) and kv_cache_quantized_layers: + quantization["kv_cache_quantized_layers"] = { + mapper(k): v for k, v in kv_cache_quantized_layers.items() + } def _assert_experts_pre_expanded( diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index c86af3aa9f5..b91ffc4079a 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -50,6 +50,7 @@ from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer from .model_config import ( KV_CACHE_FP8, + KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_INT8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE, @@ -389,27 +390,36 @@ def get_kv_cache_scaling_factor(self_attention_module: nn.Module) -> list[torch. for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer") ] - # For FP8, we recommend default kv cache scaling factor to be 1. - if get_kv_cache_dtype(self_attention_module) == KV_CACHE_FP8: - for i, factor in enumerate(scaling_factors): - if factor is None: - continue - if factor.item() > 0.5: - warn( - f"Warning: Large KV activation detected: {factor.item()}, " - "Quantized KV cache may lead to higher accuracy drop." - ) - scaling_factors[i] = torch.max( - factor, torch.tensor([1.0], dtype=torch.float, device=factor.device) + # For FP8, we recommend default KV-cache scaling factor to be 1. The + # asymmetric format applies this only to K; V remains NVFP4. + kv_cache_dtype = get_kv_cache_dtype(self_attention_module) + if kv_cache_dtype == KV_CACHE_FP8: + fp8_indices = range(len(scaling_factors)) + elif kv_cache_dtype == KV_CACHE_FP8_K_NVFP4_V: + fp8_indices = (0,) + else: + fp8_indices = () + for i in fp8_indices: + factor = scaling_factors[i] + if factor is None: + continue + if factor.item() > 0.5: + warn( + f"Warning: Large KV activation detected: {factor.item()}, " + "Quantized KV cache may lead to higher accuracy drop." ) + scaling_factors[i] = torch.max( + factor, torch.tensor([1.0], dtype=torch.float, device=factor.device) + ) return scaling_factors def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None: """Returns the kv_cache dtype. - If num_bits of output_quantizer is (4, 3) then returns FP8; if it is 8, returns int8, - otherwise returns None. + K/V quantizers are inspected as a pair so FP8 K with NVFP4 V remains + distinguishable from uniform FP8 or NVFP4. The output quantizer is retained + as a fallback for the unified Megatron export path. Args: modules: The module or list of modules to inspect. @@ -424,6 +434,29 @@ def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None: modules = [modules] for module in modules: + k_quantizer = getattr(module, "k_bmm_quantizer", None) + v_quantizer = getattr(module, "v_bmm_quantizer", None) + if ( + k_quantizer is not None + and v_quantizer is not None + and k_quantizer.is_enabled + and v_quantizer.is_enabled + ): + k_dtype = _compute_kv_cache_dtype( + [k_quantizer.num_bits], hasattr(k_quantizer, "_bias_value") + ) + v_dtype = _compute_kv_cache_dtype( + [v_quantizer.num_bits], hasattr(v_quantizer, "_bias_value") + ) + if k_dtype == KV_CACHE_FP8 and v_dtype == KV_CACHE_NVFP4: + return KV_CACHE_FP8_K_NVFP4_V + if k_dtype == v_dtype: + return k_dtype + raise NotImplementedError( + "Unsupported mixed K/V cache quantization pair: " + f"K uses {k_dtype}, while V uses {v_dtype}." + ) + # Case where the module has both k_bmm_quantizer and v_bmm_quantizer # Still check for output quantizer for the unified_megatron_export path for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer", "output_quantizer"): @@ -1001,7 +1034,7 @@ def _postprocess_single_tensor( key: str, value: torch.Tensor, kv_cache_max_bound: float, - kv_cache_format: str | None, + kv_cache_format: str | dict[str, dict[str, str]] | None, is_modelopt_qlora: bool = False, ) -> tuple[str | None, torch.Tensor | None]: """Per-tensor subset of :func:`postprocess_state_dict`, for streaming export. @@ -1035,12 +1068,15 @@ def _postprocess_single_tensor( if key.endswith(old_suffix): prefix = key[: -len(old_suffix)] if "_amax" in key: - assert kv_cache_format in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], ( - "Invalid KV cache quantization format." - ) + layer_quantization = _resolve_kv_cache_format_for_key(key, kv_cache_format) + assert layer_quantization in [ + KV_CACHE_FP8, + KV_CACHE_NVFP4, + KV_CACHE_NVFP4_AFFINE, + ], "Invalid KV cache quantization format." assert kv_cache_max_bound > 0, "Maxbound must be greater than zero." value = value.float() / kv_cache_max_bound - if kv_cache_format == KV_CACHE_FP8 and value.item() > 0.5: + if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5: logger.warning( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) @@ -1051,10 +1087,40 @@ def _postprocess_single_tensor( return None, None +def _resolve_kv_cache_format_for_key( + key: str, quantization: str | dict[str, dict[str, str]] | None +) -> str | None: + """Resolve uniform or per-layer metadata to the format for this K or V tensor.""" + if isinstance(quantization, dict): + matches = [ + (layer_name, layer_config.get("quant_algo")) + for layer_name, layer_config in quantization.items() + if key == layer_name or key.startswith(layer_name + ".") + ] + quantization = max(matches, key=lambda item: len(item[0]))[1] if matches else None + if quantization == KV_CACHE_FP8_K_NVFP4_V: + if key.endswith("k_bmm_quantizer._amax"): + return KV_CACHE_FP8 + if key.endswith("v_bmm_quantizer._amax"): + return KV_CACHE_NVFP4 + return None + return quantization + + +def _get_kv_cache_postprocess_config( + quantization_details: dict[str, Any], +) -> str | dict[str, dict[str, str]] | None: + """Return the uniform format or layer map consumed by both HF exporters.""" + kv_cache_format = quantization_details.get("kv_cache_quant_algo") + if kv_cache_format == "MIXED_PRECISION": + return quantization_details.get("kv_cache_quantized_layers", {}) + return kv_cache_format + + def postprocess_state_dict( state_dict: dict, maxbound: float, - quantization: str | None, + quantization: str | dict[str, dict[str, str]] | None, is_modelopt_qlora: bool = False, tied_map: "TiedWeightMap | None" = None, ) -> dict: @@ -1063,7 +1129,8 @@ def postprocess_state_dict( Args: state_dict: The full model state_dict. maxbound: The maximum bound value for the output quantizer. - quantization: The KV cache quantization format. + quantization: The uniform KV cache quantization format, or a per-attention-layer + ``{layer_name: {"quant_algo": ...}}`` mapping for mixed precision. is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. tied_map: Optional :class:`TiedWeightMap`. When provided, tied-weight dedup is authoritative and name-based: a declared alias key whose canonical @@ -1101,15 +1168,18 @@ def _export_key(key: str) -> str: prefix = key[: -len(old_suffix)] if "_amax" in key: - assert quantization in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], ( - "Invalid KV cache quantization format." - ) + layer_quantization = _resolve_kv_cache_format_for_key(key, quantization) + assert layer_quantization in [ + KV_CACHE_FP8, + KV_CACHE_NVFP4, + KV_CACHE_NVFP4_AFFINE, + ], "Invalid KV cache quantization format." assert maxbound > 0, "Maxbound must be greater than zero." value = value.float() / maxbound # Warn if scale exceeds threshold - if quantization == KV_CACHE_FP8 and value.item() > 0.5: + if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5: logger.warning( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) @@ -1601,7 +1671,7 @@ def get_quant_config( block_size = None # Create base config - quant_config = { + quant_config: dict[str, Any] = { "producer": { "name": "modelopt", "version": __version__, @@ -1619,7 +1689,9 @@ def get_quant_config( # It also holds awq_block_size information for applicable layers. layer_config_dict = {} - kv_cache_format = QUANTIZATION_NONE + kv_cache_formats: set[str] = set() + kv_cache_quantized_layers: dict[str, dict[str, str]] = {} + kv_cache_eligible_layers = 0 for name, module in dict(model.named_modules()).items(): # Check for standard quantizers or any quantizers from weight attributes weight_names = list(weight_attr_names(module)) @@ -1675,18 +1747,22 @@ def get_quant_config( not_enabled = SimpleNamespace(is_enabled=False) # Find kv cache quant format + has_kv_quantizers = all( + hasattr(module, quantizer_name) + for quantizer_name in ("k_bmm_quantizer", "v_bmm_quantizer") + ) + if has_kv_quantizers: + kv_cache_eligible_layers += 1 + if ( getattr(module, "k_bmm_quantizer", not_enabled).is_enabled or getattr(module, "v_bmm_quantizer", not_enabled).is_enabled or getattr(module, "output_quantizer", not_enabled).is_enabled ): module_kv_quant = get_kv_cache_dtype(module) - if kv_cache_format == QUANTIZATION_NONE: - kv_cache_format = module_kv_quant - else: - assert kv_cache_format == module_kv_quant, ( - "Do not support mixed precision kv cache quantization" - ) + if module_kv_quant != QUANTIZATION_NONE: + kv_cache_formats.add(module_kv_quant) + kv_cache_quantized_layers[name] = {"quant_algo": module_kv_quant} # MoE routers/gates are intentionally kept in original precision. On transformers>=5.0 they # are not nn.Linear modules (e.g. TopKRouter), never receive a quantizer, and would otherwise @@ -1699,8 +1775,23 @@ def get_quant_config( # Process per layer quantization config dict quant_config["quantization"].update(process_layer_quant_config(layer_config_dict)) - if kv_cache_format is not None: - quant_config["quantization"]["kv_cache_quant_algo"] = kv_cache_format + all_kv_layers_quantized = ( + kv_cache_eligible_layers > 0 and len(kv_cache_quantized_layers) == kv_cache_eligible_layers + ) + if len(kv_cache_formats) == 1 and all_kv_layers_quantized: + quant_config["quantization"]["kv_cache_quant_algo"] = next(iter(kv_cache_formats)) + elif kv_cache_quantized_layers: + weight_quant_algo = quant_config["quantization"].get("quant_algo") + if weight_quant_algo not in (None, "MIXED_PRECISION"): + raise NotImplementedError( + "Mixed-precision KV-cache export with a uniform quantized-weight format is " + "not supported yet. Use BF16 weights or a mixed-weight AutoQuant recipe." + ) + quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION" + quant_config["quantization"].setdefault("quantized_layers", {}) + quant_config["quantization"]["kv_cache_quant_algo"] = "MIXED_PRECISION" + quant_config["quantization"]["kv_cache_quantized_layers"] = kv_cache_quantized_layers + quant_config["quantization"]["kv_cache_schema_version"] = 1 return quant_config diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 77429b1cfaf..dd964098c6b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,6 +15,7 @@ """Code that export quantized Hugging Face models for deployment.""" +import copy import json import re import tempfile @@ -100,6 +101,7 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( + _get_kv_cache_postprocess_config, fuse_prequant_layernorm, fuse_prequant_to_linear, get_activation_scaling_factor, @@ -1010,12 +1012,13 @@ def _export_transformers_checkpoint( # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. kv_cache_max_bound = 448 - kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] + quantization_details = quant_config["quantization"] + kv_cache_postprocess_config = _get_kv_cache_postprocess_config(quantization_details) quantized_state_dict = postprocess_state_dict( quantized_state_dict, kv_cache_max_bound, - kv_cache_format, + kv_cache_postprocess_config, is_modelopt_qlora, tied_map=tied_map, ) @@ -1461,6 +1464,7 @@ def _write_hf_export_config( model: nn.Module, hf_quant_config: dict | None, export_dir: Path, + name_mapper: Callable[[str], str] | None = None, ) -> None: """Write hf_quant_config.json (if quantized) and embed quantization_config into config.json.""" quantization_details = (hf_quant_config or {}).get("quantization", {}) @@ -1474,6 +1478,27 @@ def _write_hf_export_config( json.dump(hf_quant_config, file, indent=4) quantization_config = convert_hf_quant_config_format(hf_quant_config) + kv_autoquant_report = next( + ( + getattr(module, "_modelopt_kv_cache_auto_quantize_state") + for module in model.modules() + if hasattr(module, "_modelopt_kv_cache_auto_quantize_state") + ), + None, + ) + if kv_autoquant_report is not None: + kv_autoquant_report = copy.deepcopy(kv_autoquant_report) + if name_mapper is not None: + kv_autoquant_report["layers"] = { + name_mapper(name): value + for name, value in kv_autoquant_report["layers"].items() + } + signature_layers = kv_autoquant_report.get("search_signature", {}).get("layers", []) + for layer in signature_layers: + layer["name"] = name_mapper(layer["name"]) + with open(f"{export_dir}/kv_cache_auto_quantize_report.json", "w") as file: + json.dump(kv_autoquant_report, file, indent=4) + original_config = f"{export_dir}/config.json" with open(original_config) as file: config_data = json.load(file) @@ -1571,6 +1596,7 @@ def export_hf_checkpoint( ) if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None + name_mapper = None try: name_mapper = build_reverse_name_mapper(model) if name_mapper is not None and hf_quant_config: @@ -1580,7 +1606,7 @@ def export_hf_checkpoint( f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor " "names may not match the original HF hub checkpoint." ) - _write_hf_export_config(model, hf_quant_config, export_dir) + _write_hf_export_config(model, hf_quant_config, export_dir, name_mapper) return post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) @@ -1602,6 +1628,7 @@ def export_hf_checkpoint( # and fails). Best-effort and atomic: any failure (an op we cannot reverse yet, # transformers API drift, unexpected shapes) falls back to the in-memory names for BOTH # weights and config so they stay mutually consistent. + name_mapper = None try: name_mapper = build_reverse_name_mapper(model) export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict) @@ -1636,7 +1663,7 @@ def export_hf_checkpoint( finally: _unpatch_revert_weight_conversion(_patches) - _write_hf_export_config(model, hf_quant_config, export_dir) + _write_hf_export_config(model, hf_quant_config, export_dir, name_mapper) except Exception as e: warnings.warn( diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 50b5d797c6a..73ca33e2a26 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -34,7 +34,11 @@ from safetensors.torch import save_file from .quant_aware_conversion import build_reverse_name_mapper -from .quant_utils import _postprocess_single_tensor, get_quant_config +from .quant_utils import ( + _get_kv_cache_postprocess_config, + _postprocess_single_tensor, + get_quant_config, +) from .registry import ExportContext from .unified_export_hf import ( _add_mtp_exclusions, @@ -253,7 +257,7 @@ def _export_transformers_checkpoint_streaming( # --- Per-tensor constants --- kv_cache_max_bound = 448 - kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] + kv_cache_format = _get_kv_cache_postprocess_config(quant_config["quantization"]) # --- Tied alias keys to skip --- # data_ptr() is unreliable for disk-offloaded weights, so we use _tied_weights_keys. diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py new file mode 100644 index 00000000000..3e0700abda3 --- /dev/null +++ b/modelopt/torch/quantization/kv_cache_auto_quant.py @@ -0,0 +1,673 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Layer-wise KV-cache AutoQuant using isolated forward KL sensitivity.""" + +from __future__ import annotations + +import fnmatch +import math +import os +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +from modelopt.torch.opt.searcher import LPS +from modelopt.torch.utils import print_rank_0, safe_load, safe_save + +from .config import QuantizeConfig +from .conversion import set_quantizer_by_cfg +from .nn import TensorQuantizer + +__all__ = ["auto_quantize_kv_cache"] + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + +_KV_QUANTIZER_ATTRS = ("k_bmm_quantizer", "v_bmm_quantizer") +_KV_AUTOQUANT_SCHEMA_VERSION = 1 +_KV_CANDIDATE_HOLDER_NAME = "layer" +_KV_CANDIDATE_NAMES = {f"{_KV_CANDIDATE_HOLDER_NAME}.{attr}" for attr in _KV_QUANTIZER_ATTRS} +_NON_KV_PROBE_NAMES = { + f"{_KV_CANDIDATE_HOLDER_NAME}.{name}" + for name in ( + "q_bmm_quantizer", + "p_bmm_quantizer", + "input_quantizer", + "output_quantizer", + "q_proj.input_quantizer", + "q_proj.weight_quantizer", + "q_proj.output_quantizer", + ) +} + + +def _disabled_quantizer() -> TensorQuantizer: + quantizer = TensorQuantizer() + quantizer.disable() + return quantizer + + +def _candidate_quantizers(config: QuantizeConfig) -> dict[str, TensorQuantizer]: + """Build a candidate with the same qualified K/V names used by a full model.""" + _validate_candidate_patterns(config) + root = nn.Module() + holder = nn.Module() + root.add_module(_KV_CANDIDATE_HOLDER_NAME, holder) + for attr in _KV_QUANTIZER_ATTRS: + setattr(holder, attr, _disabled_quantizer()) + set_quantizer_by_cfg(root, config.quant_cfg) + quantizers = {attr: getattr(holder, attr) for attr in _KV_QUANTIZER_ATTRS} + for attr, quantizer in quantizers.items(): + if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled: + raise ValueError( + f"KV-cache candidate must enable {attr}; got {type(quantizer).__name__}." + ) + return quantizers + + +def _validate_candidate_patterns(config: QuantizeConfig) -> None: + """Require every ordered config entry to match only a qualified K/V name.""" + matched_names: set[str] = set() + probe_names = _KV_CANDIDATE_NAMES | _NON_KV_PROBE_NAMES + for entry in config.quant_cfg: + if entry.parent_class is not None: + raise ValueError("KV-cache AutoQuant candidates do not support parent_class filters.") + matches = {name for name in probe_names if fnmatch.fnmatch(name, entry.quantizer_name)} + if not matches: + raise ValueError( + "KV-cache AutoQuant candidate pattern " + f"{entry.quantizer_name!r} does not match a supported qualified K/V quantizer." + ) + non_kv_matches = matches - _KV_CANDIDATE_NAMES + if non_kv_matches: + raise ValueError( + "KV-cache AutoQuant candidates may configure only k_bmm_quantizer and " + f"v_bmm_quantizer; pattern {entry.quantizer_name!r} also matches " + f"{sorted(non_kv_matches)}." + ) + matched_names.update(matches) + if matched_names != _KV_CANDIDATE_NAMES: + raise ValueError( + "KV-cache AutoQuant candidates must completely configure both " + "k_bmm_quantizer and v_bmm_quantizer." + ) + + +def _algorithm_method(config: QuantizeConfig) -> str | None: + algorithm = config.algorithm + if algorithm is None or isinstance(algorithm, str): + return algorithm + if isinstance(algorithm, dict): + return algorithm.get("method") + return getattr(algorithm, "method", None) + + +def _deployable_kv_bits(quantizer: TensorQuantizer) -> float: + """Return storage bits for the narrow K/V formats supported by unified export.""" + if quantizer.bias is not None: + raise ValueError("KV-cache AutoQuant does not support affine candidates yet.") + if quantizer.is_fp8: + return 8.0 + if quantizer.is_nvfp4_dynamic and quantizer.block_sizes.get(-1) == 16: + return 4.5 + raise ValueError( + "KV-cache AutoQuant candidates must use unified-export-compatible per-tensor FP8 " + "or block-16 dynamic NVFP4 quantizers." + ) + + +def _validate_deployable_candidate(config: QuantizeConfig) -> None: + quantizers = _candidate_quantizers(config) + k_quantizer = quantizers["k_bmm_quantizer"] + v_quantizer = quantizers["v_bmm_quantizer"] + k_bits = _deployable_kv_bits(k_quantizer) + v_bits = _deployable_kv_bits(v_quantizer) + if k_bits != v_bits and (k_bits, v_bits) != (8.0, 4.5): + raise ValueError( + "Unified export supports only uniform FP8, uniform NVFP4, or FP8-K/NVFP4-V " + "KV-cache AutoQuant candidates." + ) + + algorithm_method = _algorithm_method(config) + for attr, quantizer in quantizers.items(): + will_calibrate = algorithm_method == "max" and not quantizer._use_constant_amax + if not hasattr(quantizer, "_amax") and not will_calibrate: + raise ValueError( + f"KV-cache AutoQuant candidate {attr} has no persistent export scale. " + "Use max calibration or constant_amax; dynamic and use_constant_amax-only " + "candidates cannot be exported." + ) + + assert config.effective_bits is not None + actual_effective_bits = (k_bits + v_bits) / 2.0 + if not math.isclose(config.effective_bits, actual_effective_bits, rel_tol=0.0, abs_tol=1e-12): + raise ValueError( + "KV-cache AutoQuant candidate effective_bits does not match its configured K/V " + f"storage cost: declared {config.effective_bits}, actual {actual_effective_bits}." + ) + + +def _validate_kv_only_config(config: QuantizeConfig) -> None: + if config.effective_bits is None: + raise ValueError( + "Each KV-cache AutoQuant candidate must declare config-level effective_bits." + ) + algorithm_method = _algorithm_method(config) + if algorithm_method != "max": + if algorithm_method is not None: + raise ValueError( + "KV-cache AutoQuant supports only non-structural calibration algorithms " + f"None and 'max'; got {algorithm_method!r}." + ) + _validate_deployable_candidate(config) + + +def _validate_search_inputs( + constraints: dict[str, Any], + quantization_formats: list[tuple[dict[str, Any], str]], + num_calib_steps: int, + num_score_steps: int, +) -> tuple[float, list[tuple[str, QuantizeConfig]]]: + """Validate a KV-cache search before the caller converts the model.""" + if set(constraints) != {"kv_effective_bits"}: + raise ValueError( + "KV-cache AutoQuant constraints must contain only kv_effective_bits; " + f"got {sorted(constraints)}." + ) + target_bits = float(constraints["kv_effective_bits"]) + if not (0 < target_bits <= 16): + raise ValueError(f"kv_effective_bits must be in (0, 16], got {target_bits}.") + if num_calib_steps <= 0: + raise ValueError("num_calib_steps must be positive.") + if num_score_steps <= 0: + raise ValueError("num_score_steps must be positive.") + + candidates = [] + seen_names = set() + for raw_config, name in quantization_formats: + if name in seen_names: + raise ValueError(f"Duplicate KV-cache AutoQuant candidate name: {name!r}.") + config = QuantizeConfig(**raw_config) + _validate_kv_only_config(config) + candidates.append((name, config)) + seen_names.add(name) + if not candidates: + raise ValueError("KV-cache AutoQuant requires at least one candidate format.") + return target_bits, candidates + + +def _projection_width(module: nn.Module, side: str) -> int | None: + projection = getattr(module, f"{side}_proj", None) + out_features = getattr(projection, "out_features", None) + if isinstance(out_features, int) and out_features > 0: + return out_features + + config = getattr(module, "config", None) + num_kv_heads = getattr(config, "num_key_value_heads", None) + head_dim = getattr(module, "head_dim", None) or getattr(config, "head_dim", None) + if not isinstance(head_dim, int): + hidden_size = getattr(config, "hidden_size", None) + num_heads = getattr(config, "num_attention_heads", None) + if isinstance(hidden_size, int) and isinstance(num_heads, int) and num_heads > 0: + head_dim = hidden_size // num_heads + if isinstance(num_kv_heads, int) and isinstance(head_dim, int): + return num_kv_heads * head_dim + return None + + +def _kv_scalar_weight(module: nn.Module, name: str) -> int: + k_width = _projection_width(module, "k") + v_width = _projection_width(module, "v") + if k_width is None or v_width is None: + raise ValueError( + "Cannot determine exact KV width for eligible attention layer " + f"{name!r}. Expected k_proj/v_proj.out_features or config " + "num_key_value_heads plus head_dim." + ) + return k_width + v_width + + +def _eligible_layers( + model: nn.Module, disabled_layers: list[str] | str | None +) -> list[tuple[str, nn.Module, int]]: + patterns = [disabled_layers] if isinstance(disabled_layers, str) else disabled_layers or [] + layers = [] + for name, module in model.named_modules(): + if not all(hasattr(module, attr) for attr in _KV_QUANTIZER_ATTRS): + continue + if any(fnmatch.fnmatch(name, pattern) for pattern in patterns): + continue + layers.append((name, module, _kv_scalar_weight(module, name))) + if not layers: + raise ValueError("KV-cache AutoQuant found no eligible attention layers.") + return layers + + +def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuantizer]) -> None: + for attr, quantizer in quantizers.items(): + setattr(module, attr, quantizer) + + +@contextmanager +def _freeze_existing_quantizers(model: nn.Module, candidate_quantizers: list[TensorQuantizer]): + """Freeze calibration without changing existing quantizers' execution mode.""" + candidate_ids = {id(quantizer) for quantizer in candidate_quantizers} + states = [] + for module in model.modules(): + if ( + not isinstance(module, TensorQuantizer) + or id(module) in candidate_ids + or not module.is_enabled + ): + continue + states.append((module, module._if_calib)) + module.disable_calib() + try: + yield + finally: + for quantizer, if_calib in states: + quantizer._if_calib = if_calib + + +def _get_logits( + forward_step: Callable[[nn.Module, Any], torch.Tensor], model: nn.Module, data: Any +) -> torch.Tensor: + logits = forward_step(model, data) + if not isinstance(logits, torch.Tensor): + raise TypeError("KV-cache AutoQuant forward_step must return a logits tensor.") + if logits.ndim < 2 or logits.shape[-1] == 0: + raise ValueError( + "KV-cache AutoQuant forward_step must return logits with a non-empty vocabulary " + "dimension." + ) + if not torch.isfinite(logits).all(): + raise ValueError("KV-cache AutoQuant encountered NaN or Inf logits.") + return logits + + +def _solve_additive_recipe( + layer_names: list[str], + scalar_weights: list[int], + candidate_names: list[str], + candidate_bits: list[float], + scores: list[list[float]], + target_bits: float, + verbose: bool, +) -> tuple[list[int], str]: + denominator = float(sum(scalar_weights)) + candidate_costs = [ + [weight * bits / 16.0 for bits in candidate_bits] for weight in scalar_weights + ] + max_cost = denominator * target_bits / 16.0 + lps = LPS( + name="KVCacheAutoQuant", + constraints={"kv_cache_size_after_compression": max_cost}, + constraints_to_candidate_costs={"kv_cache_size_after_compression": candidate_costs}, + candidate_scores=scores, + objective_type="minimize", + verbose=verbose, + ) + selections, status = lps() + if status != "Optimal": + minimum_bits = sum(weight * min(candidate_bits) for weight in scalar_weights) / denominator + raise ValueError( + f"KV-cache AutoQuant could not satisfy kv_effective_bits={target_bits}; " + f"minimum achievable value is {minimum_bits:.4f}. Solver status: {status}." + ) + if len(selections) != len(layer_names): + raise RuntimeError( + "KV-cache AutoQuant solver returned an invalid selection count: " + f"{len(selections)} for {len(layer_names)} layers and candidates {candidate_names}." + ) + return selections, status + + +def _search_signature( + candidates: list[tuple[str, QuantizeConfig]], + layers: list[tuple[str, nn.Module, int]], + target_bits: float, + num_calib_steps: int, + num_score_steps: int, +) -> dict[str, Any]: + return { + "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION, + "kv_effective_bits": target_bits, + "num_calib_steps": num_calib_steps, + "num_score_steps": num_score_steps, + "candidates": [ + { + "name": name, + "config": config.model_dump(mode="json", exclude_none=True), + } + for name, config in candidates + ], + "layers": [{"name": name, "kv_scalar_weight": weight} for name, _, weight in layers], + } + + +def _checkpoint_state_is_compatible(state: dict[str, Any], signature: dict[str, Any]) -> bool: + return state.get("search_signature") == signature + + +def _quantizer_state_dict( + candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]], +) -> dict[str, dict[str, dict[str, dict[str, torch.Tensor]]]]: + return { + layer_name: { + candidate_name: { + attr: quantizer.state_dict() for attr, quantizer in layer_quantizers.items() + } + for candidate_name, layer_quantizers in layer_candidates.items() + } + for layer_name, layer_candidates in candidate_quantizers.items() + } + + +def _restore_quantizer_state_dict( + candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]], + state: dict[str, dict[str, dict[str, dict[str, torch.Tensor]]]], +) -> None: + """Restore calibration buffers into config-created candidate quantizers.""" + for layer_name, layer_candidates in candidate_quantizers.items(): + for candidate_name, layer_quantizers in layer_candidates.items(): + for attr, quantizer in layer_quantizers.items(): + quantizer_state = state[layer_name][candidate_name][attr] + for key, value in quantizer_state.items(): + if "." not in key and key not in quantizer._buffers: + quantizer.register_buffer(key, torch.empty_like(value)) + quantizer.load_state_dict(quantizer_state) + + +def _report_state(state: dict[str, Any]) -> dict[str, Any]: + """Return the JSON-safe search report, excluding calibration tensors.""" + return {key: value for key, value in state.items() if key != "quantizer_state"} + + +@torch.inference_mode() +def auto_quantize_kv_cache( + model: nn.Module, + constraints: dict[str, Any], + quantization_formats: list[tuple[dict[str, Any], str]], + data_loader: Iterable, + forward_step: Callable[[nn.Module, Any], torch.Tensor], + *, + num_calib_steps: int, + num_score_steps: int, + disabled_layers: list[str] | str | None = None, + verbose: bool = False, + checkpoint: str | None = None, +) -> tuple[nn.Module, dict[str, Any]]: + """Select one supplied K/V format per attention layer using isolated forward KL. + + Candidate formats are format-agnostic ``QuantizeConfig`` dictionaries. Each must + configure K and V together and declare ``effective_bits`` matching its packed + storage per K-or-V scalar, including scale overhead. Candidate calibration is scoped + to the candidate K/V quantizers, while pre-existing fixed quantizers keep executing + with frozen state. Persistent ``constant_amax`` formats may skip calibration forwards. + """ + target_bits, candidates = _validate_search_inputs( + constraints, quantization_formats, num_calib_steps, num_score_steps + ) + + layers = _eligible_layers(model, disabled_layers) + signature = _search_signature( + candidates, + layers, + target_bits, + num_calib_steps, + num_score_steps, + ) + candidate_names = [name for name, _ in candidates] + candidate_bits = [] + for _, config in candidates: + assert config.effective_bits is not None + candidate_bits.append(config.effective_bits) + + original_quantizers = { + name: {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS} + for name, module, _ in layers + } + disabled_quantizers = { + name: {attr: _disabled_quantizer() for attr in _KV_QUANTIZER_ATTRS} for name, _, _ in layers + } + candidate_quantizers = { + name: { + candidate_name: _candidate_quantizers(config) for candidate_name, config in candidates + } + for name, _, _ in layers + } + + is_training = model.training + model.eval() + try: + for name, module, _ in layers: + _apply_layer_quantizers(module, disabled_quantizers[name]) + + state: dict[str, Any] | None = None + if checkpoint is not None and os.path.exists(checkpoint): + restored = safe_load(checkpoint) + if not isinstance(restored, dict): + raise ValueError( + "KV-cache AutoQuant checkpoint must contain a search-state dictionary." + ) + if _checkpoint_state_is_compatible(restored, signature): + state = restored + if verbose: + print_rank_0(f"KV-cache AutoQuant restored search state from {checkpoint}.") + else: + raise ValueError( + "KV-cache AutoQuant checkpoint does not match the current candidates " + "or eligible layers. Use a different checkpoint path." + ) + + if state is not None and state.get("calibration_complete"): + quantizer_state = state.get("quantizer_state") + if quantizer_state is None: + raise ValueError( + "KV-cache AutoQuant checkpoint is missing calibrated quantizer state. " + "Use a different checkpoint path." + ) + _restore_quantizer_state_dict(candidate_quantizers, quantizer_state) + else: + from .model_quant import calibrate + + for candidate_name, config in candidates: + for layer_name, module, _ in layers: + _apply_layer_quantizers( + module, candidate_quantizers[layer_name][candidate_name] + ) + + if config.algorithm is not None: + + def calibration_loop(calibration_model): + for step, data in enumerate(data_loader): + if step >= num_calib_steps: + break + _get_logits(forward_step, calibration_model, data) + + active_quantizers = [ + quantizer + for layer_name, _, _ in layers + for quantizer in candidate_quantizers[layer_name][candidate_name].values() + ] + calibration_proxy = nn.Module() + calibration_proxy.quantizers = nn.ModuleList(active_quantizers) + with _freeze_existing_quantizers(model, active_quantizers): + calibrate( + calibration_proxy, + algorithm=config.algorithm, + forward_loop=lambda _: calibration_loop(model), + ) + + for layer_name, module, _ in layers: + _apply_layer_quantizers(module, disabled_quantizers[layer_name]) + + state = { + "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION, + "search_signature": signature, + "calibration_complete": True, + "num_calib_steps": num_calib_steps, + "quantizer_state": _quantizer_state_dict(candidate_quantizers), + } + if checkpoint is not None: + checkpoint_dir = os.path.dirname(checkpoint) + if checkpoint_dir: + os.makedirs(checkpoint_dir, exist_ok=True) + safe_save(state, checkpoint) + + assert state is not None + if not state.get("layers"): + score_sums: dict[str, dict[str, torch.Tensor | None]] = { + layer_name: dict.fromkeys(candidate_names) for layer_name, _, _ in layers + } + scored_tokens = 0 + scored_steps = 0 + iterator = tqdm( + data_loader, + total=num_score_steps, + desc="Estimating KV-cache KL sensitivity", + disable=not verbose, + ) + for data in iterator: + if scored_steps >= num_score_steps: + break + logits_ref = _get_logits(forward_step, model, data) + log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1) + scored_tokens += logits_ref.numel() // logits_ref.shape[-1] + + for layer_name, module, _ in layers: + for candidate_name, _ in candidates: + _apply_layer_quantizers( + module, candidate_quantizers[layer_name][candidate_name] + ) + logits_quant = _get_logits(forward_step, model, data) + if logits_quant.shape != logits_ref.shape: + raise ValueError( + "KV-cache AutoQuant forward_step returned different reference and " + f"candidate logits shapes: {tuple(logits_ref.shape)} and " + f"{tuple(logits_quant.shape)}." + ) + score = F.kl_div( + torch.log_softmax(logits_quant.float(), dim=-1), + log_prob_ref, + reduction="sum", + log_target=True, + ) + previous_score = score_sums[layer_name][candidate_name] + score_sums[layer_name][candidate_name] = ( + score if previous_score is None else previous_score + score + ) + _apply_layer_quantizers(module, disabled_quantizers[layer_name]) + scored_steps += 1 + + if scored_steps == 0 or scored_tokens == 0: + raise ValueError("KV-cache AutoQuant data_loader produced no scoring batches.") + scores = [] + for layer_name, _, _ in layers: + layer_scores = [] + for candidate_name in candidate_names: + score_sum = score_sums[layer_name][candidate_name] + if score_sum is None: + raise RuntimeError( + "KV-cache AutoQuant did not collect a score for " + f"{layer_name!r}/{candidate_name!r}." + ) + if not torch.isfinite(score_sum): + raise ValueError( + "KV-cache AutoQuant produced a non-finite KL score for " + f"{layer_name!r}/{candidate_name!r}." + ) + layer_scores.append(float(score_sum.item()) / scored_tokens) + scores.append(layer_scores) + selections, status = _solve_additive_recipe( + [name for name, _, _ in layers], + [weight for _, _, weight in layers], + candidate_names, + candidate_bits, + scores, + target_bits, + verbose, + ) + denominator = float(sum(weight for _, _, weight in layers)) + achieved_bits = ( + sum( + weight * candidate_bits[selected] + for selected, (_, _, weight) in zip(selections, layers) + ) + / denominator + ) + selected_score = sum( + layer_scores[selected] for selected, layer_scores in zip(selections, scores) + ) + state.update( + { + "method": "kl_div", + "score_reduction": "mean_per_scored_token", + "constraints": {"kv_effective_bits": target_bits}, + "num_score_steps": scored_steps, + "num_scored_tokens": scored_tokens, + "candidates": [ + { + "name": name, + "effective_bits": effective_bits, + "config": config.model_dump(mode="json", exclude_none=True), + } + for (name, config), effective_bits in zip(candidates, candidate_bits) + ], + "layers": { + layer_name: { + "kv_scalar_weight": weight, + "scores": dict(zip(candidate_names, layer_scores)), + "selected": candidate_names[selected], + } + for selected, layer_scores, (layer_name, _, weight) in zip( + selections, scores, layers + ) + }, + "best": { + "effective_bits": achieved_bits, + "score": selected_score, + "is_satisfied": achieved_bits <= target_bits + 1e-12, + "solver_status": status, + }, + } + ) + if checkpoint is not None: + checkpoint_dir = os.path.dirname(checkpoint) + if checkpoint_dir: + os.makedirs(checkpoint_dir, exist_ok=True) + safe_save(state, checkpoint) + if verbose: + print_rank_0(f"Saved KV-cache AutoQuant report to {checkpoint}.") + + for layer_name, module, _ in layers: + selected_name = state["layers"][layer_name]["selected"] + _apply_layer_quantizers(module, candidate_quantizers[layer_name][selected_name]) + if verbose: + print_rank_0(f"KV-cache AutoQuant selected {selected_name} for {layer_name}.") + report = _report_state(state) + model._modelopt_kv_cache_auto_quantize_state = report + return model, report + except Exception: + for layer_name, module, _ in layers: + _apply_layer_quantizers(module, original_quantizers[layer_name]) + raise + finally: + model.train(is_training) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 966a3643fe3..8fc9aa37800 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -15,6 +15,7 @@ """User-facing quantization API.""" +import copy import fnmatch import inspect import os @@ -39,12 +40,15 @@ from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe from .algorithms import get_auto_quantize_config as _get_auto_quantize_config from .config import QuantizeAlgoCfgType +from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs +from .kv_cache_auto_quant import auto_quantize_kv_cache as _auto_quantize_kv_cache from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import QuantModule, TensorQuantizer from .utils import is_quantized __all__ = [ "auto_quantize", + "auto_quantize_kv_cache", "calibrate", "compute_quantization_mse", "disable_quantizer", @@ -656,6 +660,125 @@ def _process_quantization_formats(formats, custom_name_prefix): return model, searcher.state_dict() +def auto_quantize_kv_cache( + model: nn.Module, + constraints: dict[str, Any], + quantization_formats: list[dict[str, Any] | tuple[dict[str, Any], str]], + data_loader: Iterable, + forward_step: Callable[[nn.Module, Any], torch.Tensor], + *, + num_calib_steps: int = 512, + num_score_steps: int = 128, + disabled_layers: list[str] | str | None = None, + verbose: bool = False, + checkpoint: str | None = None, +) -> tuple[nn.Module, dict[str, Any]]: + """Search layer-wise KV-cache formats using isolated forward KL sensitivity. + + Unlike weight AutoQuant, BF16/no-quant is only the scoring reference and is + not an implicit solver choice. Each supplied format must configure K and V + together and declare config-level ``effective_bits`` equal to their packed + storage cost per K-or-V scalar, including scale overhead. The budget is weighted + by the K/V projection widths of eligible layers. Formats use their own calibration + algorithm; candidates with persistent ``constant_amax`` skip calibration forwards. + + Args: + model: Model whose attention K/V quantizers will be searched. + constraints: A ``{"kv_effective_bits": target}`` storage constraint across + eligible layers. + quantization_formats: Candidate ``QuantizeConfig`` dictionaries, optionally + paired with display names. + data_loader: Re-iterable calibration and scoring batches. + forward_step: Callable returning full-vocabulary logits for the token positions + to score in one batch. + num_calib_steps: Maximum calibration batches per candidate. + num_score_steps: Maximum batches used for isolated forward-KL scoring. + disabled_layers: Optional layer-name patterns excluded from the search and + bit budget, and preserved in their existing KV-cache format. + verbose: Whether to print progress and selected formats. + checkpoint: Optional path for resumable calibration and sensitivity state. + + Returns: + The converted model with the selected per-layer K/V quantizers and a + JSON-safe sensitivity report. + """ + if ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and torch.distributed.get_world_size() > 1 + ): + raise RuntimeError( + "auto_quantize_kv_cache is single-process only; distributed scoring, selection, " + "and checkpoint writes are not synchronized." + ) + + processed_formats = [] + for idx, candidate in enumerate(quantization_formats): + if isinstance(candidate, tuple): + raw_config, name = candidate + else: + raw_config, name = candidate, f"KV_CACHE_FORMAT_{idx}" + if not isinstance(raw_config, dict): + raise TypeError("KV-cache AutoQuant formats must be config dictionaries.") + if not isinstance(name, str) or not name: + raise ValueError("KV-cache AutoQuant candidate names must be non-empty strings.") + processed_formats.append((raw_config, name)) + + _validate_kv_cache_search_inputs( + constraints, processed_formats, num_calib_steps, num_score_steps + ) + + converted_for_search = not is_quantized(model) + conversion_snapshot = _snapshot_model_structure(model) if converted_for_search else [] + try: + if converted_for_search: + model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry) + set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}]) + return _auto_quantize_kv_cache( + model, + constraints, + processed_formats, + data_loader, + forward_step, + num_calib_steps=num_calib_steps, + num_score_steps=num_score_steps, + disabled_layers=disabled_layers, + verbose=verbose, + checkpoint=checkpoint, + ) + except Exception: + if converted_for_search: + _restore_model_structure(conversion_snapshot) + raise + + +def _snapshot_model_structure( + model: nn.Module, +) -> list[tuple[nn.Module, type[nn.Module], dict[str, Any]]]: + """Capture lightweight module metadata for failure-atomic fresh conversion.""" + return [ + ( + module, + type(module), + { + key: copy.copy(value) if isinstance(value, (dict, list, set)) else value + for key, value in module.__dict__.items() + }, + ) + for module in model.modules() + ] + + +def _restore_model_structure( + snapshot: list[tuple[nn.Module, type[nn.Module], dict[str, Any]]], +) -> None: + """Undo an in-place quantization conversion without copying model tensors.""" + for module, original_type, original_state in reversed(snapshot): + object.__setattr__(module, "__class__", original_type) + module.__dict__.clear() + module.__dict__.update(original_state) + + def get_auto_quantize_config(search_state, constraints=None, verbose=False): """Build a flat quant config from auto_quantize search_state. diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml new file mode 100644 index 00000000000..836273eee53 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Layer-wise KV-cache search over calibrated FP8 and NVFP4 at 5.4 bits/scalar. +# Isolated full-vocabulary forward KL is measured with every other eligible +# attention layer kept in BF16. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers + kv_fp8: configs/ptq/units/kv_fp8 + kv_nvfp4: configs/ptq/units/kv_nvfp4 + +metadata: + recipe_type: auto_quantize + description: Layer-wise calibrated FP8/NVFP4 KV-cache search at 5.4 bits using forward KL. + +auto_quantize: + constraints: + kv_effective_bits: 5.4 + + candidate_formats: + - quant_cfg: + - $import: kv_fp8 + algorithm: max + effective_bits: 8.0 + - quant_cfg: + - $import: kv_nvfp4 + algorithm: max + effective_bits: 4.5 + + auto_quantize_method: kl_div + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + + cost_excluded_layers: + - $import: base_cost_excluded_layers diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 6a2c36e4a7c..5abdfbbf77a 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -20,6 +20,7 @@ from types import SimpleNamespace import pytest +import torch import yaml from modelopt.recipe import load_recipe @@ -83,6 +84,65 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): assert inputs["quantization_formats"][1] == QUANT_CFG_CHOICES["fp8"] +def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch): + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "fp8_cast" + ) + aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) + + assert inputs["search_domain"] == "kv_cache" + assert inputs["constraints"] == {"kv_effective_bits": 5.4} + assert inputs["method"] == "kl_div" + assert [config["effective_bits"] for config, _ in inputs["quantization_formats"]] == [ + 8.0, + 4.5, + ] + assert "kv_cache_quant_cfg" not in inputs + + +def test_kv_autoquant_kl_excludes_padding_positions(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + logits = torch.arange(2 * 4 * 3).reshape(2, 4, 3) + attention_mask = torch.tensor([[1, 1, 0, 0], [0, 1, 1, 0]]) + + selected = hf_ptq._select_unpadded_logits(logits, {"attention_mask": attention_mask}) + + assert torch.equal(selected, logits[attention_mask.bool()]) + + +def test_kv_autoquant_kl_rejects_misaligned_attention_mask(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + + with pytest.raises(ValueError, match="matching token dimensions"): + hf_ptq._select_unpadded_logits(torch.zeros(2, 4, 3), {"attention_mask": torch.ones(2, 3)}) + + +def test_autoquant_rejects_fsdp2(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + args = SimpleNamespace( + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=True, + ) + + with pytest.raises(NotImplementedError, match="does not support --use_fsdp2"): + hf_ptq.auto_quantize(args, torch.nn.Module(), [], SimpleNamespace()) + + +def test_fsdp2_autoquant_rejected_before_model_load(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + monkeypatch.setattr(hf_ptq, "_recipe_is_auto_quantize", lambda _: True) + monkeypatch.setattr( + hf_ptq.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: pytest.fail("The model config must not be loaded."), + ) + + with pytest.raises(NotImplementedError, match="does not support --use_fsdp2"): + hf_ptq.load_model(SimpleNamespace(use_fsdp2=True, recipe="autoquant")) + + def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch): """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict.""" diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 3dfb9906a54..e4dc4b3f422 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -1907,6 +1907,7 @@ def test_load_recipe_autoquantize_fixed_baseline_requires_explicit_search(tmp_pa [ "general/auto_quantize/nvfp4_fp8_at_5p4bits", "general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits", + "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits", "general/auto_quantize/nvfp4_mse_fp8_at_6p0bits", "general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits", "general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe", @@ -1925,6 +1926,22 @@ def test_load_recipe_autoquantize_builtin_general(recipe_path): assert recipe.auto_quantize.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"] +def test_load_recipe_kv_autoquantize_contract(): + recipe = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits") + aq = recipe.auto_quantize + + assert aq.constraints.effective_bits is None + assert aq.constraints.kv_effective_bits == 5.4 + assert aq.auto_quantize_method == "kl_div" + assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5] + for fmt in aq.candidate_formats: + for entry in fmt.quant_cfg: + assert entry.quantizer_name == "*[kv]_bmm_quantizer" + assert not entry.cfg.use_constant_amax + assert entry.cfg.constant_amax is None + assert fmt.algorithm == "max" + + def _all_shipped_ptq_recipe_paths(): """Every shipped PTQ recipe, discovered from disk rather than a hardcoded list.""" root = files("modelopt_recipes") diff --git a/tests/unit/torch/export/test_convert_hf_config.py b/tests/unit/torch/export/test_convert_hf_config.py new file mode 100644 index 00000000000..a0b8295b485 --- /dev/null +++ b/tests/unit/torch/export/test_convert_hf_config.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import torch + +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.unified_export_hf import _write_hf_export_config + + +def test_convert_mixed_kv_cache_config_preserves_layer_map(): + layer_map = { + "model.layers.0.self_attn": {"quant_algo": "FP8"}, + "model.layers.1.self_attn": {"quant_algo": "NVFP4"}, + } + converted = convert_hf_quant_config_format( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": {}, + "kv_cache_quant_algo": "MIXED_PRECISION", + "kv_cache_quantized_layers": layer_map, + "kv_cache_schema_version": 1, + }, + } + ) + + assert converted["quant_method"] == "modelopt" + assert converted["quant_algo"] == "MIXED_PRECISION" + assert converted["config_groups"] == {} + assert converted["kv_cache_quant_algo"] == "MIXED_PRECISION" + assert converted["kv_cache_quantized_layers"] == layer_map + assert converted["kv_cache_schema_version"] == 1 + + +def test_write_hf_export_config_writes_mapped_kv_autoquant_report(tmp_path): + layer_name = "model.layers.0.self_attn" + model = torch.nn.Module() + model._modelopt_kv_cache_auto_quantize_state = { + "layers": {layer_name: {"selected": "fp8"}}, + "search_signature": {"layers": [{"name": layer_name}]}, + } + quant_config = { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": None, + "kv_cache_quant_algo": "MIXED_PRECISION", + "kv_cache_quantized_layers": {layer_name: {"quant_algo": "FP8"}}, + "kv_cache_schema_version": 1, + }, + } + (tmp_path / "config.json").write_text("{}") + + _write_hf_export_config( + model, + quant_config, + tmp_path, + name_mapper=lambda name: f"hub.{name}", + ) + + report = json.loads((tmp_path / "kv_cache_auto_quantize_report.json").read_text()) + assert report["layers"] == {f"hub.{layer_name}": {"selected": "fp8"}} + assert report["search_signature"]["layers"] == [{"name": f"hub.{layer_name}"}] + assert model._modelopt_kv_cache_auto_quantize_state["layers"] == { + layer_name: {"selected": "fp8"} + } + assert (tmp_path / "hf_quant_config.json").is_file() + exported_config = json.loads((tmp_path / "config.json").read_text()) + assert exported_config["quantization_config"]["kv_cache_quant_algo"] == "MIXED_PRECISION" diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index e7ca68d0b69..6c778ec098e 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -27,12 +27,15 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.layer_utils import get_quantization_format from modelopt.torch.export.model_config import ( + KV_CACHE_FP8, + KV_CACHE_FP8_K_NVFP4_V, + KV_CACHE_NVFP4, QUANTIZATION_FP8, QUANTIZATION_NVFP4, QUANTIZATION_W4A8_AWQ, ) -from modelopt.torch.export.quant_utils import get_quant_config -from modelopt.torch.quantization.nn import NVFP4StaticQuantizer +from modelopt.torch.export.quant_utils import get_quant_config, postprocess_state_dict +from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer @pytest.mark.parametrize( @@ -64,6 +67,129 @@ def test_nvfp4_static_quantizer_export(): assert quant_config["quantization"]["group_size"] == 16 +def test_mixed_kv_cache_quantization_exports_per_layer_map(): + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + model = torch.nn.Module() + model.attn0 = FakeAttention() + model.attn1 = FakeAttention() + model.attn2 = FakeAttention() + mtq.set_quantizer_by_cfg( + model.attn0, + [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ], + ) + mtq.set_quantizer_by_cfg( + model.attn1, + [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "use_constant_amax": True, + }, + } + ], + ) + mtq.set_quantizer_by_cfg( + model.attn2, + [ + { + "quantizer_name": "*k_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + }, + { + "quantizer_name": "*v_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "use_constant_amax": True, + }, + }, + ], + ) + + quantization = get_quant_config(model)["quantization"] + assert quantization["quant_algo"] == "MIXED_PRECISION" + assert quantization["kv_cache_quant_algo"] == "MIXED_PRECISION" + assert quantization["quantized_layers"] == {} + assert quantization["kv_cache_quantized_layers"] == { + "attn0": {"quant_algo": "FP8"}, + "attn1": {"quant_algo": "NVFP4"}, + "attn2": {"quant_algo": "FP8_K_NVFP4_V"}, + } + + +def test_unsupported_asymmetric_kv_cache_pair_fails_export(): + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + model = FakeAttention() + mtq.set_quantizer_by_cfg( + model, + [ + { + "quantizer_name": "*k_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "use_constant_amax": True, + }, + }, + { + "quantizer_name": "*v_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + }, + ], + ) + + with pytest.raises(NotImplementedError, match="Unsupported mixed K/V cache"): + get_quant_config(model) + + +def test_mixed_kv_cache_postprocess_uses_each_layers_format(): + state_dict = { + "attn0.k_bmm_quantizer._amax": torch.tensor([448.0]), + "attn0.v_bmm_quantizer._amax": torch.tensor([224.0]), + "attn1.k_bmm_quantizer._amax": torch.tensor([112.0]), + "attn1.v_bmm_quantizer._amax": torch.tensor([56.0]), + } + layer_formats = { + "attn0": {"quant_algo": KV_CACHE_FP8}, + "attn1": {"quant_algo": KV_CACHE_NVFP4}, + "attn2": {"quant_algo": KV_CACHE_FP8_K_NVFP4_V}, + } + state_dict.update( + { + "attn2.k_bmm_quantizer._amax": torch.tensor([448.0]), + "attn2.v_bmm_quantizer._amax": torch.tensor([112.0]), + } + ) + + processed = postprocess_state_dict(state_dict, 448.0, layer_formats) + + assert processed == { + "attn0.k_proj.k_scale": torch.tensor([1.0]), + "attn0.v_proj.v_scale": torch.tensor([0.5]), + "attn1.k_proj.k_scale": torch.tensor([0.25]), + "attn1.v_proj.v_scale": torch.tensor([0.125]), + "attn2.k_proj.k_scale": torch.tensor([1.0]), + "attn2.v_proj.v_scale": torch.tensor([0.25]), + } + + class _FakeTopKRouter(torch.nn.Module): """Mimics a transformers>=5.0 MoE router: owns a ``weight`` but is NOT an ``nn.Linear``. diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 242a64f762b..09551c6bf60 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -36,9 +36,13 @@ ) import modelopt.torch.quantization as mtq -from modelopt.torch.export.model_config import KV_CACHE_FP8 +from modelopt.torch.export.model_config import KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4 from modelopt.torch.export.model_utils import TiedWeightMap -from modelopt.torch.export.quant_utils import _postprocess_single_tensor +from modelopt.torch.export.quant_utils import ( + _get_kv_cache_postprocess_config, + _postprocess_single_tensor, + _resolve_kv_cache_format_for_key, +) from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.export.unified_export_hf_streaming import ( _parse_shard_size, @@ -319,6 +323,53 @@ def test_postprocess_kv_scale_renamed_and_divided(): assert abs(val.item() - 0.5) < 1e-5 +@pytest.mark.parametrize( + ("layer_name", "quant_algo", "side", "resolved_format"), + [ + ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V, "k", KV_CACHE_FP8), + ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V, "v", KV_CACHE_NVFP4), + ("model.layers.1.self_attn", KV_CACHE_NVFP4, "k", KV_CACHE_NVFP4), + ], +) +def test_postprocess_resolves_mixed_kv_format_per_layer_and_side( + layer_name, quant_algo, side, resolved_format +): + quantization = { + "kv_cache_quant_algo": "MIXED_PRECISION", + "kv_cache_quantized_layers": {layer_name: {"quant_algo": quant_algo}}, + } + postprocess_config = _get_kv_cache_postprocess_config(quantization) + original_key = f"{layer_name}.{side}_bmm_quantizer._amax" + + assert _resolve_kv_cache_format_for_key(original_key, postprocess_config) == resolved_format + + key, val = _postprocess_single_tensor( + original_key, + torch.tensor(224.0), + 448.0, + postprocess_config, + ) + + assert key == f"{layer_name}.{side}_proj.{side}_scale" + assert val.item() == pytest.approx(0.5) + + +@pytest.mark.parametrize(("side", "resolved_format"), [("k", KV_CACHE_FP8), ("v", KV_CACHE_NVFP4)]) +def test_postprocess_resolves_uniform_asymmetric_kv_format(side, resolved_format): + original_key = f"model.layers.0.self_attn.{side}_bmm_quantizer._amax" + + assert _resolve_kv_cache_format_for_key(original_key, KV_CACHE_FP8_K_NVFP4_V) == resolved_format + key, val = _postprocess_single_tensor( + original_key, + torch.tensor(224.0), + 448.0, + KV_CACHE_FP8_K_NVFP4_V, + ) + + assert key == f"model.layers.0.self_attn.{side}_proj.{side}_scale" + assert val.item() == pytest.approx(0.5) + + def test_postprocess_scale_squeezed(): """3D scale tensors with shape[0]==1 are squeezed.""" t = torch.ones(1, 4, 4) diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index e3aa22bf3d8..06a1b88f222 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -516,6 +516,7 @@ def test_revert_quant_config_names_mapper(): "lm_head", ], "quantized_layers": {"model.layers.0.mlp.experts.0.w1": {"quant_algo": "NVFP4"}}, + "kv_cache_quantized_layers": {"model.layers.0.mlp.experts.0": {"quant_algo": "FP8"}}, } revert_quant_config_names(quant, mapper) assert quant["exclude_modules"] == [ @@ -524,6 +525,7 @@ def test_revert_quant_config_names_mapper(): "lm_head", ] assert "model.layers.0.block_sparse_moe.experts.0.w1" in quant["quantized_layers"] + assert "model.layers.0.block_sparse_moe.experts.0" in quant["kv_cache_quantized_layers"] # mapper(None) is a no-op q2 = {"exclude_modules": ["x*"]} revert_quant_config_names(q2, None) diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py new file mode 100644 index 00000000000..e415873bc99 --- /dev/null +++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py @@ -0,0 +1,663 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +import torch.nn as nn +from _test_utils.torch.transformers_models import get_tiny_llama + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import get_kv_cache_dtype, get_quant_config +from modelopt.torch.quantization import model_quant, tensor_quant +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.kv_cache_auto_quant import ( + _candidate_quantizers, + _kv_scalar_weight, + _solve_additive_recipe, + _validate_kv_only_config, + auto_quantize_kv_cache, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.fixture +def nvfp4_fake_quant_stub(monkeypatch): + """Keep CPU search tests independent of the CUDA-only NVFP4 fake-quant kernel.""" + + monkeypatch.setattr( + tensor_quant, + "dynamic_block_quantize_op", + lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs), + ) + + +def _quantizer_cfg(bits, *, constant_amax=None): + cfg = {"num_bits": bits} + if bits == (2, 1): + cfg["block_sizes"] = {-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + if constant_amax is not None: + cfg["constant_amax"] = constant_amax + return cfg + + +def _kv_config(bits, effective_bits, *, algorithm="max", constant_amax=None): + return QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": _quantizer_cfg(bits, constant_amax=constant_amax), + } + ], + algorithm=algorithm, + effective_bits=effective_bits, + ) + + +def test_kv_candidate_requires_exact_bits_and_both_sides(): + _validate_kv_only_config(_kv_config((4, 3), 8.0)) + + with pytest.raises(ValueError, match="config-level effective_bits"): + _validate_kv_only_config( + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ] + ) + ) + with pytest.raises(ValueError, match="completely configure both"): + _validate_kv_only_config( + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*k_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ], + effective_bits=8.0, + ) + ) + + +@pytest.mark.parametrize("algorithm", ["svdquant", {"method": "smoothquant"}, {"method": "mse"}]) +def test_kv_candidate_rejects_structural_or_unscoped_algorithms(algorithm): + config = _kv_config((4, 3), 8.0).model_copy(update={"algorithm": algorithm}) + + with pytest.raises(ValueError, match="only non-structural calibration algorithms"): + _validate_kv_only_config(config) + + +def test_kv_additive_solver_spends_fp8_on_more_sensitive_layer(): + selections, status = _solve_additive_recipe( + layer_names=["layer0", "layer1"], + scalar_weights=[256, 256], + candidate_names=["fp8", "nvfp4"], + candidate_bits=[8.0, 4.5], + scores=[[0.0, 10.0], [0.0, 1.0]], + target_bits=6.25, + verbose=False, + ) + + assert status == "Optimal" + assert selections == [0, 1] + + +def test_kv_scalar_weight_counts_k_and_v_widths(): + module = nn.Module() + module.k_proj = nn.Linear(32, 24, bias=False) + module.v_proj = nn.Linear(32, 16, bias=False) + + assert _kv_scalar_weight(module, "attention") == 40 + + +@pytest.mark.parametrize( + ("config", "expected_format"), + [ + (_kv_config((4, 3), 8.0), "FP8"), + (_kv_config((2, 1), 4.5), "NVFP4"), + (_kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0), "FP8"), + (_kv_config((2, 1), 4.5, algorithm=None, constant_amax=1.0), "NVFP4"), + ], +) +def test_kv_candidate_accepts_export_supported_persistent_formats(config, expected_format): + _validate_kv_only_config(config) + quantizers = _candidate_quantizers(config) + module = nn.Module() + for name, quantizer in quantizers.items(): + setattr(module, name, quantizer) + + assert get_kv_cache_dtype(module) == expected_format + + +@pytest.mark.parametrize( + ("config", "match"), + [ + (_kv_config((4, 3), 8.0, algorithm=None), "no persistent export scale"), + ( + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ], + algorithm="max", + effective_bits=8.0, + ), + "no persistent export scale", + ), + (_kv_config(8, 8.0, algorithm=None, constant_amax=1.0), "per-tensor FP8"), + (_kv_config(4, 4.0, algorithm=None, constant_amax=1.0), "per-tensor FP8"), + (_kv_config((4, 3), 6.0), "does not match its configured K/V storage cost"), + ], +) +def test_kv_candidate_rejects_non_exportable_or_incorrect_cost(config, match): + with pytest.raises(ValueError, match=match): + _validate_kv_only_config(config) + + +class _ToyKVAttention(nn.Module): + def __init__(self, width, gain): + super().__init__() + self.k_proj = nn.Linear(width, width, bias=False) + self.v_proj = nn.Linear(width, width, bias=False) + self.k_bmm_quantizer = nn.Identity() + self.v_bmm_quantizer = nn.Identity() + self.gain = gain + + def forward(self, x): + return x + self.gain * (self.k_bmm_quantizer(x) + self.v_bmm_quantizer(x)) + + +class _ToyKVModel(nn.Module): + def __init__(self, width=8): + super().__init__() + self.attn0 = _ToyKVAttention(width, gain=0.25) + self.attn1 = _ToyKVAttention(width, gain=2.0) + self.lm_head = nn.Linear(width, width, bias=False) + + def forward(self, x): + return self.lm_head(self.attn1(self.attn0(x))) + + +def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fake_quant_stub): + torch.manual_seed(123) + model = _ToyKVModel() + data = [torch.randn(2, 3, 8), torch.randn(2, 3, 8)] + candidates = [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + "algorithm": None, + "effective_bits": 8.0, + }, + "fp8", + ), + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + } + ], + "algorithm": None, + "effective_bits": 4.5, + }, + "nvfp4", + ), + ] + + model, state = auto_quantize_kv_cache( + model, + {"kv_effective_bits": 6.25}, + candidates, + data, + lambda model, batch: model(batch), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "kv_search.pth"), + ) + + assert state["best"]["effective_bits"] == pytest.approx(6.25) + assert state["best"]["is_satisfied"] + assert model.training + assert {layer["selected"] for layer in state["layers"].values()} == { + "fp8", + "nvfp4", + } + for layer_name, layer_state in state["layers"].items(): + layer = model.get_submodule(layer_name) + assert layer.k_bmm_quantizer.num_bits == layer.v_bmm_quantizer.num_bits + expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1) + assert layer.k_bmm_quantizer.num_bits == expected_bits + + restored_model = _ToyKVModel().eval() + restored_model, restored_state = auto_quantize_kv_cache( + restored_model, + {"kv_effective_bits": 6.25}, + candidates, + data, + lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "kv_search.pth"), + ) + + assert restored_state == state + assert not restored_model.training + for layer_name, layer_state in restored_state["layers"].items(): + layer = restored_model.get_submodule(layer_name) + expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1) + assert layer.k_bmm_quantizer.num_bits == expected_bits + + +def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quant_stub): + model = _ToyKVModel() + candidate = ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + }, + { + "quantizer_name": "*.k_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + }, + ], + "algorithm": None, + "effective_bits": 6.25, + }, + "fp8_k_nvfp4_v", + ) + + model, state = auto_quantize_kv_cache( + model, + {"kv_effective_bits": 6.25}, + [candidate], + [torch.randn(1, 2, 8)], + lambda search_model, batch: search_model(batch), + num_calib_steps=1, + num_score_steps=1, + ) + + assert state["candidates"][0]["effective_bits"] == pytest.approx(6.25) + assert state["best"]["effective_bits"] == pytest.approx(6.25) + for layer_state in state["layers"].values(): + assert layer_state["selected"] == "fp8_k_nvfp4_v" + for layer in (model.attn0, model.attn1): + assert layer.k_bmm_quantizer.num_bits == (4, 3) + assert layer.v_bmm_quantizer.num_bits == (2, 1) + assert get_quant_config(model)["quantization"]["kv_cache_quant_algo"] == "FP8_K_NVFP4_V" + + +def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state(): + model = _ToyKVModel() + original_quantizers = { + name: (module.k_bmm_quantizer, module.v_bmm_quantizer) + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)) + } + candidates = [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + } + ], + "algorithm": None, + "effective_bits": 4.5, + }, + "nvfp4", + ) + ] + + with pytest.raises(ValueError, match="non-empty vocabulary dimension"): + auto_quantize_kv_cache( + model, + {"kv_effective_bits": 4.5}, + candidates, + [torch.randn(2, 3, 8)], + lambda *_: torch.ones(8), + num_calib_steps=1, + num_score_steps=1, + ) + + assert model.training + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)): + assert (module.k_bmm_quantizer, module.v_bmm_quantizer) == original_quantizers[name] + + +def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_fake_quant_stub): + torch.manual_seed(123) + model = get_tiny_llama(num_hidden_layers=2) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))} for _ in range(2)] + candidates = [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3)}, + }, + ], + "algorithm": "max", + "effective_bits": 8.0, + }, + "fp8", + ), + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + } + ], + "algorithm": None, + "effective_bits": 4.5, + }, + "nvfp4", + ), + ] + + model, state = mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 6.25}, + candidates, + data, + lambda search_model, batch: search_model(**batch).logits, + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "hf_kv_search.pth"), + ) + + assert len(state["layers"]) == model.config.num_hidden_layers + assert state["best"]["effective_bits"] == pytest.approx(6.25) + assert all( + layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled + for layer in model.model.layers + ) + non_kv_quantizers = [ + quantizer + for name, quantizer in model.named_modules() + if isinstance(quantizer, TensorQuantizer) + and not name.endswith(("k_bmm_quantizer", "v_bmm_quantizer")) + ] + assert non_kv_quantizers + assert all(not quantizer.is_enabled for quantizer in non_kv_quantizers) + assert all( + layer.self_attn.q_proj.weight_quantizer.num_bits == 8 for layer in model.model.layers + ) + exported_quantization = get_quant_config(model)["quantization"] + assert exported_quantization["quantized_layers"] == {} + assert exported_quantization["kv_cache_quantized_layers"] + assert set(exported_quantization["kv_cache_quantized_layers"]) <= { + f"model.layers.{idx}.self_attn" for idx in range(model.config.num_hidden_layers) + } + + restored_model = get_tiny_llama(num_hidden_layers=2) + restored_model, restored_state = mtq.auto_quantize_kv_cache( + restored_model, + {"kv_effective_bits": 6.25}, + candidates, + data, + lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "hf_kv_search.pth"), + ) + + assert restored_state == state + assert any( + hasattr(layer.self_attn.k_bmm_quantizer, "_amax") + for layer in restored_model.model.layers + if layer.self_attn.k_bmm_quantizer.num_bits == (4, 3) + ) + + +def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic(): + model = get_tiny_llama(num_hidden_layers=1) + original_types = {name: type(module) for name, module in model.named_modules()} + invalid_candidate = _kv_config((4, 3), 8.0).model_dump() + invalid_candidate["algorithm"] = "svdquant" + + with pytest.raises(ValueError, match="only non-structural calibration algorithms"): + mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 8.0}, + [invalid_candidate], + [], + lambda *_: pytest.fail("Validation must run before model conversion."), + num_calib_steps=1, + num_score_steps=1, + ) + + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + valid_candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump() + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + with pytest.raises(ValueError, match="non-empty vocabulary dimension"): + mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 8.0}, + [valid_candidate], + data, + lambda *_: torch.ones(8), + num_calib_steps=1, + num_score_steps=1, + ) + + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + +def test_public_kv_autoquant_rejects_unmatched_or_unexportable_candidates_before_conversion(): + model = get_tiny_llama(num_hidden_layers=1) + original_types = {name: type(module) for name, module in model.named_modules()} + unmatched = _kv_config((4, 3), 8.0).model_dump() + unmatched["quant_cfg"].append({"quantizer_name": "q_proj.*_quantizer", "cfg": {"num_bits": 2}}) + invalid_candidates = [ + (unmatched, "does not match a supported qualified K/V quantizer"), + (_kv_config((4, 3), 8.0, algorithm=None).model_dump(), "no persistent export scale"), + ( + _kv_config(8, 8.0, algorithm=None, constant_amax=1.0).model_dump(), + "per-tensor FP8", + ), + ] + + for candidate, match in invalid_candidates: + with pytest.raises(ValueError, match=match): + mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 8.0}, + [candidate], + [], + lambda *_: pytest.fail("Validation must run before model conversion."), + num_calib_steps=1, + num_score_steps=1, + ) + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + +def test_public_kv_autoquant_rejects_distributed_execution_before_mutation(monkeypatch): + model = get_tiny_llama(num_hidden_layers=1) + original_types = {name: type(module) for name, module in model.named_modules()} + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + + with pytest.raises(RuntimeError, match="single-process only"): + mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 8.0}, + [], + [], + lambda *_: pytest.fail("Distributed validation must fail before search."), + ) + + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + +def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers( + monkeypatch, nvfp4_fake_quant_stub +): + torch.manual_seed(123) + model = get_tiny_llama(num_hidden_layers=2) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + fixed_kv_config = { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + "algorithm": None, + } + model = mtq.quantize(model, fixed_kv_config) + fixed_weight_quantizer = model.model.layers[0].self_attn.q_proj.weight_quantizer + fixed_weight_quantizer.enable() + fixed_weight_quantizer.amax = torch.tensor(1.0) + fixed_weight_quantizer.disable_quant() + fixed_weight_quantizer.disable_calib() + observed_fixed_states = [] + fixed_hook = fixed_weight_quantizer.register_forward_hook( + lambda module, _inputs, _output: observed_fixed_states.append( + (module.is_enabled, module._if_quant, module._if_calib) + ) + ) + fixed_qdq_quantizer = model.model.layers[1].self_attn.q_proj.weight_quantizer + fixed_qdq_quantizer.enable() + fixed_qdq_quantizer.amax = torch.tensor(1.0) + observed_qdq_states = [] + qdq_hook = fixed_qdq_quantizer.register_forward_hook( + lambda module, _inputs, _output: observed_qdq_states.append( + (module.is_enabled, module._if_quant, module._if_calib) + ) + ) + calibration_states = [] + real_calibrate = model_quant.calibrate + + def calibrate_with_state_check(*args, **kwargs): + calibration_states.append( + (fixed_weight_quantizer._if_quant, fixed_weight_quantizer._if_calib) + ) + result = real_calibrate(*args, **kwargs) + calibration_states.append( + (fixed_weight_quantizer._if_quant, fixed_weight_quantizer._if_calib) + ) + return result + + monkeypatch.setattr(model_quant, "calibrate", calibrate_with_state_check) + + try: + model, state = mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 4.5}, + [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + }, + }, + ], + "algorithm": "max", + "effective_bits": 4.5, + }, + "nvfp4", + ) + ], + data, + lambda search_model, batch: search_model(**batch).logits, + num_calib_steps=1, + num_score_steps=1, + disabled_layers="model.layers.1.self_attn", + ) + finally: + fixed_hook.remove() + qdq_hook.remove() + + assert set(state["layers"]) == {"model.layers.0.self_attn"} + assert observed_fixed_states + assert all(state == (True, False, False) for state in observed_fixed_states) + assert calibration_states == [(False, False), (False, False)] + assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == (2, 1) + assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled + assert not fixed_weight_quantizer._if_quant + assert not fixed_weight_quantizer._if_calib + assert fixed_weight_quantizer.amax.item() == pytest.approx(1.0) + assert observed_qdq_states + assert all(quantizer_state == (True, True, False) for quantizer_state in observed_qdq_states) + assert fixed_qdq_quantizer._if_quant + assert not fixed_qdq_quantizer._if_calib + assert fixed_qdq_quantizer.amax.item() == pytest.approx(1.0) + fixed_attention = model.model.layers[1].self_attn + assert fixed_attention.k_bmm_quantizer.is_enabled + assert fixed_attention.v_bmm_quantizer.is_enabled + assert fixed_attention.k_bmm_quantizer.num_bits == (4, 3) + assert fixed_attention.v_bmm_quantizer.num_bits == (4, 3)