diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 4063cea26cd..fa501489b96 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -425,11 +425,24 @@ For models without backprop support (e.g. Llama-4), use the `kl_div` scoring met 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 recipe searches FP8-cast K/V (8.0 bits/scalar) and NVFP4-cast K/V +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: @@ -441,9 +454,8 @@ python hf_ptq.py \ --export_path /path/to/qwen3-1.7b-mixed-kv ``` -Each candidate uses an explicit constant scale, avoiding an additional calibration pass while -keeping persistent K/V scales in the unified HF checkpoint. Unified export records the selected -formats in `kv_cache_quantized_layers` and writes +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. @@ -455,8 +467,9 @@ the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`; > Do not deploy them with the pinned runtime. Full FP8 K/V and full NVFP4 K/V use existing vLLM > kernels once the layer-wise metadata consumer is available. -The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an -interrupted search (skips re-scoring): +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 \ diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6daade6ed24..4ca92371298 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -386,7 +386,10 @@ def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]: 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. @@ -413,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 @@ -449,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, @@ -456,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. @@ -475,8 +498,14 @@ def auto_quantize( 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 = ( @@ -538,7 +567,7 @@ def forward_step(model, batch): ), verbose=True, disabled_layers=inputs["disabled_layers"], - checkpoint=args.auto_quantize_checkpoint, + checkpoint=checkpoint, ) return language_model @@ -556,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. @@ -843,6 +872,86 @@ 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 + primary_uses_kv_checkpoint = primary_is_kv and fixed_quantize_config is not None + + 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=( + "kv_auto_quantize_checkpoint" + if primary_uses_kv_checkpoint + else "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, @@ -1198,10 +1307,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): @@ -1277,7 +1384,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) + ) ), ) @@ -1289,16 +1400,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: @@ -1333,25 +1444,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( @@ -1417,7 +1510,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, ) @@ -1591,6 +1684,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 19e9ef7278b..5912cdd4010 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -333,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( @@ -343,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/quant_utils.py b/modelopt/torch/export/quant_utils.py index 2d8f94366a6..8cee1c9ed3a 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -1774,13 +1774,11 @@ def get_quant_config( quant_config["quantization"]["kv_cache_quant_algo"] = next(iter(kv_cache_formats)) elif kv_cache_quantized_layers: weight_quant_algo = quant_config["quantization"].get("quant_algo") - if weight_quant_algo not in (None, "MIXED_PRECISION"): - raise NotImplementedError( - "Mixed-precision KV-cache export with a uniform quantized-weight format is " - "not supported yet. Use BF16 weights or a mixed-weight AutoQuant recipe." - ) - quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION" - quant_config["quantization"].setdefault("quantized_layers", {}) + 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 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/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/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 95eaec41ed5..7f0c3d8077a 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,11 @@ 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 @@ -106,8 +110,33 @@ def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch): assert "kv_cache_quant_cfg" not in inputs -def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch): - """The HF entry point runs the real public KV AutoQuant path on an offline Qwen fixture.""" +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, @@ -115,7 +144,7 @@ def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch): lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs), ) model = get_tiny_qwen3(num_hidden_layers=1) - aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize + aq = getattr(load_recipe(recipe_path), kv_stage) args = SimpleNamespace( calib_with_images=False, inference_pipeline_parallel=1, @@ -135,6 +164,140 @@ def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch): 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) + original_auto_quantize = hf_ptq.auto_quantize + checkpoint_attrs = [] + + def recording_auto_quantize(*args, **kwargs): + checkpoint_attrs.append(kwargs["checkpoint_attr"]) + return original_auto_quantize(*args, **kwargs) + + monkeypatch.setattr(hf_ptq, "auto_quantize", recording_auto_quantize) + 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 + assert checkpoint_attrs == ["kv_auto_quantize_checkpoint"] + + +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) diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 1a4fabb627b..20b860c8001 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -1916,8 +1916,10 @@ 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", @@ -1961,6 +1963,94 @@ def test_load_recipe_kv_autoquantize_contract(): 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( diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 1202b1fcae4..0fe5636f1b8 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -161,6 +161,50 @@ def __init__(self): } +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):