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..fa501489b96 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -422,11 +422,54 @@ 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. + +To optimize GEMM and KV cache in one invocation, compose ordered stages in the same recipe. A fixed +`quantize` block followed by a KV-domain `auto_quantize` first calibrates the GEMM weight/activation +configuration, then searches K/V while the existing GEMM QDQ remains enabled with calibration +frozen. See `general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits`. + +A weight-domain `auto_quantize` can instead add a `kv_auto_quantize` follow-up with its own method, +constraints, candidates, score size, and disabled layers. This supports, for example, a +gradient-based GEMM search followed by a KL-divergence KV search; see +`general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits`. When the +follow-up is present, the recipe owns KV configuration and suppresses the CLI's uniform +`--kv_cache_qformat` fallback. Use `--auto_quantize_checkpoint` for the weight search and +`--kv_auto_quantize_checkpoint` for the KV search. + +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: -The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an -interrupted search (skips re-scoring): +```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. + +For a single-stage search, `--auto_quantize_checkpoint` saves/restores the search state to resume an +interrupted search (skips re-scoring). Composed weight-plus-KV recipes additionally use +`--kv_auto_quantize_checkpoint` for the independent KV search state: ```bash scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits \ @@ -460,6 +503,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..c6339a2dc55 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,8 +355,41 @@ 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 + nvfp4_quantizers = type(fmt)(**KV_QUANT_CFG_CHOICES["nvfp4"]).model_dump(exclude_none=True)[ + "quant_cfg" + ] + fp8_quantizers = type(fmt)(**KV_QUANT_CFG_CHOICES["fp8"]).model_dump(exclude_none=True)[ + "quant_cfg" + ] + if len(fp8_quantizers) != 1: + raise RuntimeError("The FP8 KV preset must contain exactly one quantizer entry.") + fp8_k_quantizer = copy.deepcopy(fp8_quantizers[0]) + fp8_k_quantizer["quantizer_name"] = "*.k_bmm_quantizer" + if candidate_quantizers == [*nvfp4_quantizers, fp8_k_quantizer]: + name = "fp8_k_nvfp4_v" + 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 + aq_config, + args: argparse.Namespace, + fixed_quantize_config=None, + allow_uniform_kv: bool = True, ) -> dict: """Map a resolved AutoQuantizeConfig to mtq.auto_quantize inputs. @@ -349,6 +399,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. @@ -356,7 +416,9 @@ def _mtq_inputs_from_auto_quantize_config( constraints.setdefault("cost", {})["excluded_module_name_patterns"] = ( aq_config.cost_excluded_layers ) - if aq_config.kv_cache is not None: + if not allow_uniform_kv: + kv_cache_quant_cfg = None + elif aq_config.kv_cache is not None: kv_cache_quant_cfg = aq_config.kv_cache.model_dump() elif args.kv_cache_qformat == KV_CACHE_NONE: kv_cache_quant_cfg = None @@ -380,6 +442,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, @@ -391,6 +454,22 @@ def _mtq_inputs_from_auto_quantize_config( } +def _assert_kv_autoquantize_input_is_clean(model: torch.nn.Module) -> None: + """Fail closed if an upstream stage left actual K/V quantizers enabled.""" + enabled = [ + name + for name, module in model.named_modules(remove_duplicate=False) + if name.endswith(("k_bmm_quantizer", "v_bmm_quantizer")) + and getattr(module, "is_enabled", False) + ] + if enabled: + raise ValueError( + "The preceding weight/activation stage left K/V quantizers enabled on the converted " + f"model: {enabled}. Disable them in that stage before running mixed-KV AutoQuant; " + "clearing them now would not undo its calibration or sensitivity measurements." + ) + + def auto_quantize( args: argparse.Namespace, language_model: torch.nn.Module, @@ -398,6 +477,8 @@ def auto_quantize( aq_config, full_model: torch.nn.Module | None = None, fixed_quantize_config=None, + allow_uniform_kv: bool = True, + checkpoint_attr: str = "auto_quantize_checkpoint", ): """Recipe-driven auto_quantize, organized around an AutoQuantizeConfig. @@ -414,15 +495,17 @@ 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 + aq_config, + args, + fixed_quantize_config=fixed_quantize_config, + allow_uniform_kv=allow_uniform_kv, ) + if inputs["search_domain"] == "kv_cache": + _assert_kv_autoquantize_input_is_clean(language_model) + checkpoint = getattr(args, checkpoint_attr, None) # base-model lm_head handling (mirrors the CLI helper) is_base_model = ( @@ -461,14 +544,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=checkpoint, + ) + return language_model + language_model, _ = mtq.auto_quantize( language_model, constraints=inputs["constraints"], @@ -483,7 +585,7 @@ def forward_step(model, batch): verbose=True, disabled_layers=inputs["disabled_layers"], method=inputs["method"], - checkpoint=args.auto_quantize_checkpoint, + checkpoint=checkpoint, ) # KV cache quantization is uniform; applied after the LP search. @@ -512,6 +614,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 @@ -768,6 +872,81 @@ def mono_quantize( warnings.warn("Skipping quantization: model is already quantized.") +def _prepare_quant_cfg( + args: argparse.Namespace, quant_cfg: dict[str, Any], full_model: torch.nn.Module +) -> dict[str, Any]: + """Apply shared checkpoint-local adjustments to a PTQ configuration.""" + mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None) + if mtp_layer_prefixes: + quant_cfg = copy.deepcopy(quant_cfg) + for prefix in mtp_layer_prefixes: + pattern = f"*{prefix}*" + quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) + print(f"Excluding MTP layer from quantization: {pattern}") + + if needs_checkpoint_path_update(quant_cfg): + quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path) + print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}") + + if args.cast_mxfp4_to_nvfp4: + quant_cfg = copy.deepcopy(quant_cfg) + force_weight_quantizers_static(quant_cfg["quant_cfg"]) + return quant_cfg + + +def _run_auto_quantize_recipe( + args: argparse.Namespace, + recipe: ModelOptAutoQuantizeRecipe, + full_model: torch.nn.Module, + language_model: torch.nn.Module, + model_type: str | None, + calibration_only: bool, + calib_dataloader: DataLoader, + is_nemotron_vl_model: bool, +) -> None: + """Run the recipe's fixed PTQ, weight search, and KV search in order.""" + primary = recipe.auto_quantize + followup_kv = recipe.kv_auto_quantize + primary_is_kv = primary.constraints.kv_effective_bits is not None + fixed_quantize_config = recipe.quantize + + if primary_is_kv and fixed_quantize_config is not None: + quant_cfg = _prepare_quant_cfg(args, fixed_quantize_config.model_dump(), full_model) + mono_quantize( + args, + quant_cfg, + full_model, + language_model, + model_type, + calibration_only, + calib_dataloader, + is_nemotron_vl_model, + ) + fixed_quantize_config = None + + auto_quantize( + args, + full_model, + calib_dataloader, + aq_config=primary, + full_model=full_model, + fixed_quantize_config=fixed_quantize_config, + allow_uniform_kv=followup_kv is None, + checkpoint_attr="auto_quantize_checkpoint", + ) + + if followup_kv is not None: + auto_quantize( + args, + full_model, + calib_dataloader, + aq_config=followup_kv, + full_model=full_model, + allow_uniform_kv=False, + checkpoint_attr="kv_auto_quantize_checkpoint", + ) + + def export_quantized( args: argparse.Namespace, full_model: torch.nn.Module, @@ -1123,10 +1302,8 @@ def quantize_main( # AutoQuantize is recipe-driven: everything downstream reads the resolved AutoQuantizeConfig. if isinstance(recipe, ModelOptAutoQuantizeRecipe): aq_config = recipe.auto_quantize - fixed_quantize_config = recipe.quantize else: aq_config = None - fixed_quantize_config = None def _is_layerwise(obj): if isinstance(obj, ModelOptPTQRecipe): @@ -1202,7 +1379,11 @@ def _is_layerwise(obj): device, model_type, autoquant_gradient_recipe=( - aq_config is not None and aq_config.auto_quantize_method == "gradient" + isinstance(recipe, ModelOptAutoQuantizeRecipe) + and any( + config is not None and config.auto_quantize_method == "gradient" + for config in (recipe.auto_quantize, recipe.kv_auto_quantize) + ) ), ) @@ -1214,16 +1395,16 @@ def _is_layerwise(obj): ) if aq_config is not None: - # AutoQuantize (recipe-driven). For VL models the search walks the OUTER CausalLM (which - # carries lm_head and the LM-head forward path); architecture-specific exclusions come - # from aq_config.disabled_layers. - auto_quantize( + assert isinstance(recipe, ModelOptAutoQuantizeRecipe) + _run_auto_quantize_recipe( args, + recipe, full_model, + language_model, + model_type, + calibration_only, calib_dataloader, - aq_config, - full_model=full_model, - fixed_quantize_config=fixed_quantize_config, + is_nemotron_vl_model, ) else: @@ -1258,25 +1439,7 @@ def _is_layerwise(obj): KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"], ) - # Exclude MTP layers from quantization if detected (e.g., GLM-4.7's layer 92). - # These layers are typically speculative decoding layers that should be exported as-is. - # Complementary to recipe `*mtp*` wildcards (name-match); this catches MTP layers - # identified by index. - mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None) - if mtp_layer_prefixes: - quant_cfg = copy.deepcopy(quant_cfg) - for prefix in mtp_layer_prefixes: - pattern = f"*{prefix}*" - quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) - print(f"Excluding MTP layer from quantization: {pattern}") - - if needs_checkpoint_path_update(quant_cfg): - quant_cfg, resolved_dir = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path) - print(f"Auto-resolved layerwise checkpoint_dir: {resolved_dir}") - - if args.cast_mxfp4_to_nvfp4: - quant_cfg = copy.deepcopy(quant_cfg) - force_weight_quantizers_static(quant_cfg["quant_cfg"]) + quant_cfg = _prepare_quant_cfg(args, quant_cfg, full_model) if quant_cfg: mono_quantize( @@ -1342,7 +1505,7 @@ def parse_args() -> argparse.Namespace: "general/ptq/nvfp4_default-kv_fp8_cast, general/auto_quantize/nvfp4_fp8_at_4p8bits). " "KV cache source depends on the recipe type: PTQ recipes bake KV cache into quant_cfg " "and --kv_cache_qformat is ignored; AutoQuantize recipes fall back to --kv_cache_qformat " - "unless the recipe sets an explicit kv_cache field." + "unless the recipe sets an explicit kv_cache or kv_auto_quantize field." ), default=None, ) @@ -1516,6 +1679,15 @@ def parse_args() -> argparse.Namespace: "(sensitivity scores, costs, etc.). Used with an AutoQuantize --recipe." ), ) + parser.add_argument( + "--kv_auto_quantize_checkpoint", + type=str, + default=None, + help=( + "Path for saving/restoring the KV-cache search checkpoint in a composed recipe. " + "Use a new path whenever the preceding weight/activation quantization stage changes." + ), + ) parser.add_argument( "--moe_calib_experts_ratio", type=float, diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index 2ca8f4f462b..5912cdd4010 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -149,10 +149,19 @@ 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( + effective_bits: float | None = ModeloptField( default=4.8, 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."), + ) + 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 +174,35 @@ class AutoQuantizeConstraints(ModeloptBaseConfig): description="Extra cost-model parameters; omit for the 'weight' cost model.", ) - @field_validator("effective_bits") + @model_validator(mode="before") + @classmethod + def _select_kv_constraint(cls, data): + if isinstance(data, dict) and "kv_effective_bits" in data and "effective_bits" not in data: + data = dict(data) + data["effective_bits"] = None + return data + + @field_validator("effective_bits", "kv_effective_bits") @classmethod - def _validate_effective_bits(cls, v: float) -> float: - if not (0 < v <= 16): + 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 +303,25 @@ 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." + ) + if self.cost_excluded_layers: + raise ValueError( + "KV-cache AutoQuant does not support cost_excluded_layers; use " + "disabled_layers to exclude non-KV-cache modules from the search." + ) return self @@ -283,9 +333,9 @@ class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase): quantize: QuantizeConfig | None = ModeloptField( default=None, title="Fixed PTQ baseline", - description="Optional normal PTQ QuantizeConfig for modules outside the explicit " - "AutoQuantize module_search_spaces. Fixed and searched modules are calibrated, scored, " - "costed, and exported in one integrated AutoQuantize operation.", + description="Optional normal PTQ QuantizeConfig. A weight AutoQuantize stage uses it for " + "modules outside explicit module_search_spaces; a KV AutoQuantize stage applies it first " + "as the fixed GEMM weight/activation configuration.", ) auto_quantize: AutoQuantizeConfig = Field( @@ -293,22 +343,42 @@ class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase): description="AutoQuantize search configuration. Required.", ) + kv_auto_quantize: AutoQuantizeConfig | None = ModeloptField( + default=None, + title="Follow-up KV-cache AutoQuantize config", + description="Optional KV-cache search run after the primary weight AutoQuantize search.", + ) + @model_validator(mode="after") def _validate_fixed_and_searched_spaces(self): + primary_is_kv = self.auto_quantize.constraints.kv_effective_bits is not None + if self.kv_auto_quantize is not None: + if primary_is_kv: + raise ValueError( + "kv_auto_quantize cannot follow an auto_quantize stage that already searches " + "the KV cache." + ) + if self.kv_auto_quantize.constraints.kv_effective_bits is None: + raise ValueError("kv_auto_quantize must use a kv_effective_bits constraint.") + if self.auto_quantize.kv_cache is not None: + raise ValueError( + "A weight AutoQuantize stage followed by kv_auto_quantize must omit the " + "uniform auto_quantize.kv_cache post-step." + ) has_fixed_baseline = self.quantize is not None has_global_search = bool(self.auto_quantize.candidate_formats) - if has_fixed_baseline and has_global_search: + if not primary_is_kv and has_fixed_baseline and has_global_search: raise ValueError( "An AutoQuantize recipe with a fixed quantize baseline must omit top-level " "auto_quantize.candidate_formats and explicitly list searched modules under " "auto_quantize.module_search_spaces." ) - if has_fixed_baseline and not self.auto_quantize.module_search_spaces: + if not primary_is_kv and has_fixed_baseline and not self.auto_quantize.module_search_spaces: raise ValueError( "An AutoQuantize recipe with a fixed quantize baseline requires at least one " "auto_quantize.module_search_spaces entry." ) - if not has_fixed_baseline and not has_global_search: + if not primary_is_kv and not has_fixed_baseline and not has_global_search: raise ValueError( "An AutoQuantize recipe without a fixed quantize baseline requires top-level " "auto_quantize.candidate_formats for unmatched modules." 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/model_utils.py b/modelopt/torch/export/model_utils.py index f27405ec83f..cf89ddc2e5b 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -108,7 +108,7 @@ def is_multimodal_model(model): config = model.config # Check for Nemotron-Parse encoder-decoder architecture - architectures = getattr(config, "architectures", []) + architectures = getattr(config, "architectures", None) or [] is_nemotron_parse = any("nemotronparse" in arch.lower() for arch in architectures) return ( @@ -137,12 +137,17 @@ def get_language_model_from_vl(model) -> list[nn.Module] | None: >>> # lineage[0] is vlm_model >>> # lineage[1] is vlm_model.language_model """ - # always prioritize model.model.langauge_model + candidates = [] if hasattr(model, "model") and hasattr(model.model, "language_model"): - return [model, model.model, model.model.language_model] - + candidates.append([model, model.model, model.model.language_model]) if hasattr(model, "language_model"): - return [model, model.language_model] + candidates.append([model, model.language_model]) + if len(candidates) > 1: + raise ValueError( + "Found multiple language-model roots; refusing to select one by traversal order." + ) + if candidates: + return candidates[0] # Pattern 3: For encoder-decoder VL models (e.g., Nemotron-Parse), the decoder is the language model. # Only match if the model is detected as multimodal to avoid matching non-VLM encoder-decoder 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..8cee1c9ed3a 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -18,7 +18,6 @@ import logging from collections import defaultdict from collections.abc import Generator -from types import SimpleNamespace from typing import Any from warnings import warn @@ -50,6 +49,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 +389,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 +433,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 +1033,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 +1067,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 +1086,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 +1128,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 +1167,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 +1670,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 +1688,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)) @@ -1672,21 +1743,18 @@ def get_quant_config( layer_config_dict[name + ".quantization"] = quantization_format layer_config_dict[name + ".awq_block_size"] = block_size - not_enabled = SimpleNamespace(is_enabled=False) - # Find kv cache quant format - 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" - ) + 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 module.k_bmm_quantizer.is_enabled and module.v_bmm_quantizer.is_enabled: + module_kv_quant = get_kv_cache_dtype(module) + 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 +1767,21 @@ 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 is None: + quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION" + quant_config["quantization"]["quantized_layers"] = {} + elif weight_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..ecf6e52ffbe --- /dev/null +++ b/modelopt/torch/quantization/kv_cache_auto_quant.py @@ -0,0 +1,735 @@ +# 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 _candidate_kv_bits(config: QuantizeConfig) -> tuple[float, float]: + quantizers = _candidate_quantizers(config) + return ( + _deployable_kv_bits(quantizers["k_bmm_quantizer"]), + _deployable_kv_bits(quantizers["v_bmm_quantizer"]), + ) + + +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(): + if quantizer._dynamic: + raise ValueError( + f"KV-cache AutoQuant candidate {attr} uses top-level dynamic quantization, " + "which does not retain a persistent export scale." + ) + 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 _validate_candidate_cost_geometry( + candidates: list[tuple[str, QuantizeConfig]], + layers: list[tuple[str, nn.Module, int]], +) -> None: + candidate_bits = [_candidate_kv_bits(config) for _, config in candidates] + if all(k_bits == v_bits for k_bits, v_bits in candidate_bits): + return + + unequal_width_layers = [] + for name, module, _ in layers: + k_width = _projection_width(module, "k") + v_width = _projection_width(module, "v") + if k_width != v_width: + unequal_width_layers.append(f"{name} (K={k_width}, V={v_width})") + if unequal_width_layers: + raise ValueError( + "KV-cache AutoQuant cannot cost asymmetric K/V candidates on layers with unequal " + "K/V widths: " + ", ".join(unequal_width_layers) + "." + ) + + +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 [] + boundaries = [] + names_by_identity: dict[int, list[str]] = {} + for name, module in model.named_modules(remove_duplicate=False): + if not all(hasattr(module, attr) for attr in _KV_QUANTIZER_ATTRS): + continue + boundaries.append((name, module)) + names_by_identity.setdefault(id(module), []).append(name) + aliases = [names for names in names_by_identity.values() if len(names) > 1] + if aliases: + raise ValueError( + f"KV-cache attention boundaries are registered through aliases: {aliases}." + ) + + layers = [] + for name, module in boundaries: + 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 _validate_persistent_candidate_scales( + candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]], +) -> None: + """Require every calibrated candidate scale to be persistent in its state dict.""" + 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(): + if "_amax" not in quantizer.state_dict(): + raise ValueError( + f"KV-cache AutoQuant candidate {candidate_name!r} for " + f"{layer_name!r}/{attr} has no persistent export scale after calibration." + ) + + +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) + _validate_candidate_cost_geometry(candidates, 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) + _validate_persistent_candidate_scales(candidate_quantizers) + 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]) + + _validate_persistent_candidate_scales(candidate_quantizers) + 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/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml new file mode 100644 index 00000000000..925c1bd6972 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Fixed FP8 GEMM PTQ followed by layer-wise KV-cache AutoQuantize. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + fp8: configs/numerics/fp8 + nvfp4: configs/numerics/nvfp4 + w8a8_fp8_fp8: configs/ptq/units/w8a8_fp8_fp8 + +metadata: + recipe_type: auto_quantize + description: Fixed FP8 GEMM PTQ followed by mixed FP8/NVFP4 KV-cache search. + +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - $import: w8a8_fp8_fp8 + - $import: default_disabled_quantizers + +auto_quantize: + constraints: + kv_effective_bits: 5.4 + + candidate_formats: + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: fp8 + constant_amax: 448.0 + algorithm: + effective_bits: 8.0 + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: nvfp4 + constant_amax: 448.0 + algorithm: + effective_bits: 4.5 + + auto_quantize_method: kl_div + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + - "*mtp*" 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..ba191e87bc8 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Layer-wise KV-cache search over FP8-cast and NVFP4-cast 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 + fp8: configs/numerics/fp8 + nvfp4: configs/numerics/nvfp4 + +metadata: + recipe_type: auto_quantize + description: Layer-wise FP8-cast/NVFP4-cast KV-cache search at 5.4 bits using forward KL. + +auto_quantize: + constraints: + kv_effective_bits: 5.4 + + candidate_formats: + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: fp8 + constant_amax: 448.0 + algorithm: + effective_bits: 8.0 + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: nvfp4 + constant_amax: 448.0 + algorithm: + effective_bits: 4.5 + + auto_quantize_method: kl_div + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + - "*mtp*" diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml new file mode 100644 index 00000000000..41214ddeb71 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Gradient-based GEMM AutoQuantize followed by layer-wise KV-cache AutoQuantize. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_cost_excluded_layers: configs/auto_quantize/units/base_cost_excluded_layers + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + fp8: configs/ptq/presets/model/fp8 + kv_fp8: configs/numerics/fp8 + kv_nvfp4: configs/numerics/nvfp4 + nvfp4: configs/ptq/presets/model/nvfp4 + +metadata: + recipe_type: auto_quantize + description: Gradient GEMM search followed by KL-divergence mixed-KV search at 5.4 bits. + +auto_quantize: + constraints: + effective_bits: 5.4 + + candidate_formats: + - $import: nvfp4 + - $import: fp8 + + auto_quantize_method: gradient + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + + cost_excluded_layers: + - $import: base_cost_excluded_layers + +kv_auto_quantize: + constraints: + kv_effective_bits: 5.4 + + candidate_formats: + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: kv_fp8 + constant_amax: 448.0 + algorithm: + effective_bits: 8.0 + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: kv_nvfp4 + constant_amax: 448.0 + algorithm: + effective_bits: 4.5 + + auto_quantize_method: kl_div + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + - "*mtp*" diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index cf75e50e107..dc70c180057 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -220,6 +220,7 @@ def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel: "head_dim": 8, "max_position_embeddings": 32, "vocab_size": 32, + "rope_scaling": {"rope_type": "default", "mrope_section": [1, 1, 2]}, } text_kwargs.update(config_kwargs) # Pass as dicts — transformers 5.3.0 Qwen3VLConfig.__init__ only handles diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 6a2c36e4a7c..652eecaa665 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import getpass import importlib import sys @@ -20,11 +21,18 @@ from types import SimpleNamespace import pytest +import torch import yaml +from _test_utils.torch.transformers_models import get_tiny_qwen3 from modelopt.recipe import load_recipe -from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints +from modelopt.recipe.config import ( + AutoQuantizeConfig, + AutoQuantizeConstraints, + ModelOptAutoQuantizeRecipe, +) from modelopt.recipe.presets import QUANT_CFG_CHOICES +from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -83,6 +91,260 @@ 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 aq.cost_excluded_layers == [] + assert "*mtp*" in inputs["disabled_layers"] + assert "kv_cache_quant_cfg" not in inputs + + +def test_followup_kv_autoquant_suppresses_uniform_kv_fallback(monkeypatch): + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "fp8_cast" + ) + aq = load_recipe("general/auto_quantize/nvfp4_fp8_at_5p4bits").auto_quantize + + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args, allow_uniform_kv=False) + + assert inputs["kv_cache_quant_cfg"] is None + + +@pytest.mark.parametrize( + ("recipe_path", "kv_stage"), + [ + ("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits", "auto_quantize"), + ( + "general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits", + "auto_quantize", + ), + ( + "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits", + "kv_auto_quantize", + ), + ], +) +def test_hf_ptq_shipped_kv_autoquant_recipes_invoke_public_api(monkeypatch, recipe_path, kv_stage): + """Every shipped recipe runs the real public KV AutoQuant path on an offline Qwen fixture.""" + hf_ptq = _import_hf_ptq(monkeypatch) + monkeypatch.setattr( + tensor_quant, + "dynamic_block_quantize_op", + lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs), + ) + model = get_tiny_qwen3(num_hidden_layers=1) + aq = getattr(load_recipe(recipe_path), kv_stage) + args = SimpleNamespace( + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=False, + kv_cache_qformat="none", + batch_size=1, + auto_quantize_checkpoint=None, + ) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + + hf_ptq.auto_quantize(args, model, data, aq, full_model=model) + + attention = model.model.layers[0].self_attn + assert attention.k_bmm_quantizer.num_bits == (2, 1) + assert attention.v_bmm_quantizer.num_bits == (2, 1) + assert attention.k_bmm_quantizer.amax == 448.0 + assert attention.v_bmm_quantizer.amax == 448.0 + + +def test_hf_ptq_runs_weight_then_kv_autoquantize_stages(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + weight_aq = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=8.0), + candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES["fp8"])], + auto_quantize_method="gradient", + ) + kv_aq = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(kv_effective_bits=8.0), + candidate_formats=[ + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + algorithm=None, + effective_bits=8.0, + ) + ], + auto_quantize_method="kl_div", + ) + recipe = ModelOptAutoQuantizeRecipe(auto_quantize=weight_aq, kv_auto_quantize=kv_aq) + calls = [] + monkeypatch.setattr( + hf_ptq, + "auto_quantize", + lambda *_args, **kwargs: calls.append(kwargs), + ) + + hf_ptq._run_auto_quantize_recipe( + SimpleNamespace(), recipe, torch.nn.Module(), torch.nn.Module(), None, False, [], False + ) + + assert [call["aq_config"] for call in calls] == [weight_aq, kv_aq] + assert calls[0]["allow_uniform_kv"] is False + assert calls[0]["checkpoint_attr"] == "auto_quantize_checkpoint" + assert calls[1]["checkpoint_attr"] == "kv_auto_quantize_checkpoint" + + +def test_hf_ptq_runs_fixed_ptq_before_kv_autoquantize(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + fixed = QuantizeConfig( + quant_cfg=[ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*q_proj.weight_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None, "constant_amax": 1.0}, + }, + ], + algorithm=None, + ) + kv_aq = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(kv_effective_bits=8.0), + candidate_formats=[ + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + algorithm=None, + effective_bits=8.0, + ) + ], + auto_quantize_method="kl_div", + score_size=1, + ) + recipe = ModelOptAutoQuantizeRecipe(quantize=fixed, auto_quantize=kv_aq) + model = get_tiny_qwen3(num_hidden_layers=1) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + args = SimpleNamespace( + qformat="fp8", + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=False, + batch_size=1, + auto_quantize_checkpoint=None, + kv_auto_quantize_checkpoint=None, + pyt_ckpt_path="dummy", + cast_mxfp4_to_nvfp4=False, + ) + + hf_ptq._run_auto_quantize_recipe(args, recipe, model, model, None, False, data, False) + + attention = model.model.layers[0].self_attn + assert attention.q_proj.weight_quantizer.is_enabled + assert attention.q_proj.weight_quantizer.num_bits == (4, 3) + assert attention.k_bmm_quantizer.is_enabled + assert attention.v_bmm_quantizer.is_enabled + + +def test_composed_kv_autoquantize_rejects_enabled_actual_kv_quantizers(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + model = get_tiny_qwen3(num_hidden_layers=1) + hf_ptq.mtq.quantize( + model, + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + "algorithm": None, + }, + ) + + args = SimpleNamespace( + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=False, + kv_cache_qformat="none", + batch_size=1, + auto_quantize_checkpoint=None, + ) + aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize + + with pytest.raises(ValueError, match="preceding weight/activation stage left K/V"): + hf_ptq.auto_quantize(args, model, [], aq, full_model=model) + assert model.model.layers[0].self_attn.k_bmm_quantizer.is_enabled + + +def test_kv_autoquant_names_asymmetric_export_format(monkeypatch): + """The supported FP8-K/NVFP4-V candidate has a stable semantic name.""" + hf_ptq = _import_hf_ptq(monkeypatch) + mixed_config = copy.deepcopy(hf_ptq.KV_QUANT_CFG_CHOICES["nvfp4"]) + fp8_k_quantizer = copy.deepcopy(hf_ptq.KV_QUANT_CFG_CHOICES["fp8"]["quant_cfg"][0]) + fp8_k_quantizer["quantizer_name"] = "*.k_bmm_quantizer" + mixed_config["quant_cfg"].append(fp8_k_quantizer) + mixed_config["effective_bits"] = 6.25 + + candidates = hf_ptq._mtq_kv_candidate_formats([QuantizeConfig(**mixed_config)]) + + assert candidates[0][1] == "fp8_k_nvfp4_v" + + +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..20b860c8001 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -28,6 +28,8 @@ import modelopt.torch.quantization.config as qcfg from modelopt.recipe.config import ( + AutoQuantizeConfig, + AutoQuantizeConstraints, ModelOptAutoQuantizeRecipe, ModelOptDFlashRecipe, ModelOptEagleRecipe, @@ -1768,6 +1770,15 @@ def test_load_recipe_autoquantize_minimal(tmp_path): assert aq.module_search_spaces == [] +def test_autoquantize_constraints_preserve_default_with_kv_override(): + assert AutoQuantizeConstraints.model_fields["effective_bits"].default == 4.8 + assert AutoQuantizeConstraints().effective_bits == 4.8 + + constraints = AutoQuantizeConstraints(kv_effective_bits=5.4) + assert constraints.effective_bits is None + assert constraints.kv_effective_bits == 5.4 + + def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path): """cost_model + cost.active_moe_expert_ratio parse and dump to the mtq constraints dict shape.""" recipe_file = tmp_path / "aq.yml" @@ -1905,8 +1916,11 @@ def test_load_recipe_autoquantize_fixed_baseline_requires_explicit_search(tmp_pa @pytest.mark.parametrize( "recipe_path", [ + "general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits", "general/auto_quantize/nvfp4_fp8_at_5p4bits", "general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits", + "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_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", @@ -1918,11 +1932,133 @@ def test_load_recipe_autoquantize_builtin_general(recipe_path): assert isinstance(recipe, ModelOptAutoQuantizeRecipe) assert len(recipe.auto_quantize.candidate_formats) >= 2 assert recipe.auto_quantize.auto_quantize_method in ("gradient", "kl_div") - # Both shared base units must be spliced in: the removed --auto_quantize_* CLI shim appended - # them unconditionally, so a general recipe is the migration target and must match it. Without - # cost_excluded_layers a VL/MTP model counts its vision tower in the effective-bits denominator. assert "*output_layer*" in recipe.auto_quantize.disabled_layers - assert recipe.auto_quantize.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"] + if recipe.auto_quantize.constraints.kv_effective_bits is not None: + assert "*mtp*" in recipe.auto_quantize.disabled_layers + assert recipe.auto_quantize.cost_excluded_layers == [] + else: + 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 "*mtp*" in aq.disabled_layers + assert aq.cost_excluded_layers == [] + assert all(candidate.algorithm is None for candidate in aq.candidate_formats) + 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 == 448.0 + assert fmt.algorithm is None + + +@pytest.mark.parametrize( + ("recipe_path", "kv_stage"), + [ + ( + "general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits", + "auto_quantize", + ), + ( + "general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits", + "kv_auto_quantize", + ), + ], +) +def test_builtin_kv_autoquantize_recipes_use_calibration_free_cast_candidates( + recipe_path, kv_stage +): + aq = getattr(load_recipe(recipe_path), kv_stage) + + assert all(candidate.algorithm is None for candidate in aq.candidate_formats) + assert all( + candidate.quant_cfg[0].cfg.constant_amax == 448.0 for candidate in aq.candidate_formats + ) + + +def test_load_recipe_fixed_ptq_then_kv_autoquantize(tmp_path): + recipe_file = tmp_path / "ptq-then-kv.yml" + recipe_file.write_text( + "metadata:\n recipe_type: auto_quantize\n" + "quantize:\n algorithm: max\n quant_cfg:\n" + " - quantizer_name: '*'\n enable: false\n" + " - quantizer_name: '*.weight_quantizer'\n" + " cfg: {num_bits: [4, 3], axis: null}\n" + "auto_quantize:\n constraints:\n kv_effective_bits: 8.0\n" + " candidate_formats:\n" + " - algorithm: null\n effective_bits: 8.0\n quant_cfg:\n" + " - quantizer_name: '*[kv]_bmm_quantizer'\n" + " cfg: {num_bits: [4, 3], constant_amax: 1.0}\n" + " auto_quantize_method: kl_div\n" + ) + + recipe = load_recipe(recipe_file) + + assert recipe.quantize is not None + assert recipe.auto_quantize.constraints.kv_effective_bits == 8.0 + assert recipe.kv_auto_quantize is None + + +def test_load_recipe_weight_autoquantize_then_kv_autoquantize(tmp_path): + recipe_file = tmp_path / "weight-then-kv.yml" + recipe_file.write_text( + _AQ_MINIMAL_BODY + "kv_auto_quantize:\n constraints:\n kv_effective_bits: 8.0\n" + " candidate_formats:\n" + " - algorithm: null\n effective_bits: 8.0\n quant_cfg:\n" + " - quantizer_name: '*[kv]_bmm_quantizer'\n" + " cfg: {num_bits: [4, 3], constant_amax: 1.0}\n" + " auto_quantize_method: kl_div\n" + ) + + recipe = load_recipe(recipe_file) + + assert recipe.auto_quantize.auto_quantize_method == "gradient" + assert recipe.auto_quantize.constraints.effective_bits == 4.8 + assert recipe.kv_auto_quantize is not None + assert recipe.kv_auto_quantize.auto_quantize_method == "kl_div" + assert recipe.kv_auto_quantize.constraints.kv_effective_bits == 8.0 + + +def test_composed_kv_autoquantize_accepts_scoped_gemm_rule(tmp_path): + recipe_file = tmp_path / "scoped-gemm.yml" + recipe_file.write_text( + "metadata:\n recipe_type: auto_quantize\n" + "quantize:\n algorithm: max\n quant_cfg:\n" + " - quantizer_name: 'model.layers.*.mlp.*'\n" + " cfg: {num_bits: [4, 3], constant_amax: 1.0}\n" + "auto_quantize:\n constraints:\n kv_effective_bits: 8.0\n" + " candidate_formats:\n" + " - algorithm: null\n effective_bits: 8.0\n quant_cfg:\n" + " - quantizer_name: '*[kv]_bmm_quantizer'\n" + " cfg: {num_bits: [4, 3], constant_amax: 1.0}\n" + " auto_quantize_method: kl_div\n" + ) + + recipe = load_recipe(recipe_file) + + assert recipe.quantize is not None + assert recipe.quantize.quant_cfg[0].quantizer_name == "model.layers.*.mlp.*" + + +def test_kv_autoquantize_rejects_cost_excluded_layers(): + with pytest.raises(ValueError, match=r"cost_excluded_layers.*disabled_layers"): + AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(kv_effective_bits=8.0), + candidate_formats=[qcfg.QuantizeConfig(quant_cfg=[], effective_bits=8.0)], + auto_quantize_method="kl_div", + cost_excluded_layers=["*mtp*"], + ) def _all_shipped_ptq_recipe_paths(): 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..0fe5636f1b8 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,205 @@ def test_nvfp4_static_quantizer_export(): assert quant_config["quantization"]["group_size"] == 16 +def test_projection_output_quantizers_are_not_exported_as_kv_cache(): + model = ToyModel() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*.weight_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*.input_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*.output_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda x: x(torch.randn(1, 4, 10))) + + quantization = get_quant_config(model)["quantization"] + + assert quantization["quant_algo"] == "FP8" + assert quantization["kv_cache_quant_algo"] is None + assert "kv_cache_quantized_layers" not in quantization + + +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_uniform_weight_quantization_exports_mixed_kv_cache_map(): + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + model = ToyModel() + mtq.quantize(model, partial_fp8_config, lambda x: x(torch.randn(1, 4, 10))) + model.attn0 = FakeAttention() + model.attn1 = 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, + }, + } + ], + ) + + quantization = get_quant_config(model)["quantization"] + + assert quantization["quant_algo"] == "FP8" + assert quantization["kv_cache_quant_algo"] == "MIXED_PRECISION" + assert quantization["kv_cache_quantized_layers"] == { + "attn0": {"quant_algo": "FP8"}, + "attn1": {"quant_algo": "NVFP4"}, + } + + +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/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index b9fa29238d7..ca7147d61da 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -15,6 +15,8 @@ """Tests for tied-weight helpers in unified_export_hf.""" +from types import SimpleNamespace + import pytest import torch from _test_utils.torch.quantization.tied_modules import ( @@ -23,7 +25,11 @@ ) import modelopt.torch.quantization as mtq -from modelopt.torch.export.model_utils import TiedWeightMap +from modelopt.torch.export.model_utils import ( + TiedWeightMap, + get_language_model_from_vl, + is_multimodal_model, +) from modelopt.torch.export.quant_utils import ( fuse_prequant_layernorm, postprocess_state_dict, @@ -32,6 +38,25 @@ from modelopt.torch.quantization.nn import TensorQuantizer +def test_multimodal_detection_accepts_null_architectures(): + """Unified export treats absent architecture metadata as an empty list.""" + model = SimpleNamespace(config=SimpleNamespace(architectures=None)) + + assert not is_multimodal_model(model) + + +@pytest.mark.parametrize("aliased", [False, True]) +def test_language_model_extraction_rejects_competing_or_aliased_roots(aliased): + """Language-model extraction must not select ambiguous roots by traversal order.""" + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.language_model = torch.nn.Module() + model.language_model = model.model.language_model if aliased else torch.nn.Module() + + with pytest.raises(ValueError, match="multiple language-model roots"): + get_language_model_from_vl(model) + + def test_hf_all_tied_weights_keys_contract(): """Pin the transformers API we build tied_map from, so a version bump fails loud here. 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..f93ce8913af --- /dev/null +++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py @@ -0,0 +1,780 @@ +# 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 pytest +import torch +import torch.nn as nn +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3, get_tiny_qwen3vl + +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, + _eligible_layers, + _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 _asymmetric_kv_config(): + return QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": _quantizer_cfg((2, 1), constant_amax=1.0), + }, + { + "quantizer_name": "*.k_bmm_quantizer", + "cfg": _quantizer_cfg((4, 3), constant_amax=1.0), + }, + ], + algorithm=None, + effective_bits=6.25, + ) + + +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) + + +def test_kv_candidate_rejects_top_level_dynamic_fp8(): + config = QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic"}, + } + ], + algorithm="max", + effective_bits=8.0, + ) + + with pytest.raises(ValueError, match="top-level dynamic"): + _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_rejects_missing_scale_after_calibration(monkeypatch): + model = _ToyKVModel() + original_quantizers = { + name: (module.k_bmm_quantizer, module.v_bmm_quantizer) + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)) + } + monkeypatch.setattr(model_quant, "calibrate", lambda *_args, **_kwargs: None) + + with pytest.raises(ValueError, match="no persistent export scale after calibration"): + auto_quantize_kv_cache( + model, + {"kv_effective_bits": 8.0}, + [(_kv_config((4, 3), 8.0).model_dump(), "fp8")], + [torch.randn(1, 2, 8)], + lambda search_model, batch: search_model(batch), + 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_kv_eligible_layers_supports_hybrid_attention_mixers_only(): + """Hybrid decoders include attention mixers but exclude nonattention mixers.""" + model = nn.Module() + model.layers = nn.ModuleList([nn.Module(), nn.Module()]) + model.layers[0].mixer = nn.Linear(8, 8, bias=False) + model.layers[1].mixer = _ToyKVAttention(8, gain=1.0) + + layers = _eligible_layers(model, disabled_layers=None) + + assert [(name, width) for name, _, width in layers] == [("layers.1.mixer", 16)] + + +def test_kv_eligible_layers_rejects_aliased_attention_boundary(): + """An attention object registered at multiple paths must not be selected by traversal order.""" + model = nn.Module() + attention = _ToyKVAttention(8, gain=1.0) + model.attention = attention + model.attention_alias = attention + + with pytest.raises(ValueError, match="registered through aliases"): + _eligible_layers(model, disabled_layers=None) + + +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 = (_asymmetric_kv_config().model_dump(exclude_none=True), "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_asymmetric_candidate_for_unequal_kv_widths( + nvfp4_fake_quant_stub, +): + model = _ToyKVModel() + model.attn0.k_proj = nn.Linear(8, 12, bias=False) + model.attn0.v_proj = nn.Linear(8, 8, bias=False) + + with pytest.raises(ValueError, match=r"asymmetric K/V candidates.*unequal K/V widths"): + auto_quantize_kv_cache( + model, + {"kv_effective_bits": 6.25}, + [(_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")], + [torch.randn(1, 2, 8)], + lambda search_model, batch: search_model(batch), + num_calib_steps=1, + num_score_steps=1, + ) + + +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) + ) + + +@pytest.mark.parametrize( + ("model_factory", "expected_layer", "disabled_layers"), + [ + (get_tiny_qwen3, "model.layers.0.self_attn", None), + (get_tiny_qwen3vl, "model.language_model.layers.0.self_attn", "*visual*"), + ], +) +def test_public_kv_autoquant_selects_qwen_causal_attention_only( + model_factory, expected_layer, disabled_layers +): + """Plain and conditional Qwen models expose only causal attention to the KV search.""" + model = model_factory(num_hidden_layers=1) + text_config = getattr(model.config, "text_config", model.config) + data = [{"input_ids": torch.randint(0, text_config.vocab_size, (1, 8))}] + candidate = ( + _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump(), + "fp8", + ) + + model, state = mtq.auto_quantize_kv_cache( + model, + {"kv_effective_bits": 8.0}, + [candidate], + data, + lambda search_model, batch: search_model(**batch).logits, + num_calib_steps=1, + num_score_steps=1, + disabled_layers=disabled_layers, + ) + + assert set(state["layers"]) == {expected_layer} + assert json.loads(json.dumps(state))["layers"][expected_layer]["selected"] == "fp8" + exported = get_quant_config(model)["quantization"] + if "kv_cache_quantized_layers" in exported: + assert set(exported["kv_cache_quantized_layers"]) == {expected_layer} + else: + assert exported["kv_cache_quant_algo"] == "FP8" + + +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)