diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 030c3d41fd8..b032dca0bd6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,7 @@ Changelog **New Features** +- Add normalized and group-boundary AutoQuant scoring. Set ``auto_quantize.constraints.score_model: per_element`` in an AutoQuantize recipe to normalize selector coefficients by represented weight elements, and set ``auto_quantize_method: group_recon`` with ``score_boundary: group`` to score projection recipes by normalized reconstruction error at their attention or MLP group output. ``quant_grouping_scheme`` independently controls whether self-attention q/k/v/o or linear-attention qkv/z/out share one recipe decision beyond mandatory runtime-fused groups. Shared-expert gate/up/down projections are grouped as one fused-MoE decision. The deprecated HF PTQ CLI flags map to the same recipe fields. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 3a7380a9f9e..99aa7445f8c 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -352,18 +352,23 @@ Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize `AutoQuantize` can be performed for Huggingface LLM models like [Qwen](https://huggingface.co/Qwen/Qwen3-8B) / [Nemotron](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) as shown below: `AutoQuantize` is driven by an **AutoQuantize recipe** passed with `--recipe`. The recipe defines the -candidate formats, the `effective_bits` target, cost model, scoring method, search-disabled layers, and -cost-excluded layers — see [`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in +candidate formats, the `effective_bits` target, cost model, scoring objective and boundary, +search-disabled layers, and cost-excluded layers — see +[`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in [`modelopt_recipes/general/auto_quantize/`](../../modelopt_recipes/general/auto_quantize); model-specific recipes (carrying architecture-specific disabled layers — e.g. VL vision towers) live under `modelopt_recipes/huggingface//auto_quantize/`. > *Migration: prefer an AutoQuantize `--recipe`. The `--auto_quantize_bits`, `--auto_quantize_method`, -> `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and `--auto_quantize_active_moe_expert_ratio` +> `--auto_quantize_score_size`, `--auto_quantize_score_model`, +> `--auto_quantize_score_boundary`, `--auto_quantize_cost_model`, and +> `--auto_quantize_active_moe_expert_ratio` > CLI flags are **deprecated but still work** — they are converted into an `AutoQuantizeConfig` on the fly > (with a `DeprecationWarning`) and will be removed in a future release. They map to recipe fields: > `--auto_quantize_bits` → `constraints.effective_bits`, `--auto_quantize_method` → `auto_quantize_method`, -> `--auto_quantize_score_size` → `score_size`, `--auto_quantize_cost_model` → `constraints.cost_model`, +> `--auto_quantize_score_size` → `score_size`, `--auto_quantize_score_model` → +> `constraints.score_model`, `--auto_quantize_score_boundary` → `score_boundary`, +> `--auto_quantize_cost_model` → `constraints.cost_model`, > `--auto_quantize_active_moe_expert_ratio` → `constraints.cost.active_moe_expert_ratio`, and the > `--qformat fp8,nvfp4` candidate list → `candidate_formats`. When converted, the shared base > `disabled_layers` and `cost_excluded_layers` patterns are appended automatically. `--auto_quantize_checkpoint` @@ -381,9 +386,16 @@ scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/n The recipe quantizes the less accuracy-sensitive layers with the more aggressive format (e.g. NVFP4) and keeps the more sensitive ones at higher precision (or unquantized), so the model meets the recipe's `effective_bits` target. To author your own, copy a shipped recipe and adjust `candidate_formats`, -`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div`), `score_size`, -`disabled_layers` (excluded from the search), and `cost_excluded_layers` (kept out of the bit-budget -accounting — e.g. VL vision towers). Recipes can splice a shared base `disabled_layers` set via +`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `group_recon` / `kl_div`), +`constraints.score_model` (`raw` / `per_element`), `score_boundary` (`local` / `group`), +`quant_grouping_scheme`, `score_size`, `disabled_layers` (excluded from the search), and +`cost_excluded_layers` (kept out +of the bit-budget accounting — e.g. VL vision towers). `group_recon` measures normalized +reconstruction error at shared attention/MLP outputs; group scoring changes the measurement +boundary but does not force those projections to share one recipe decision. Set +`quant_grouping_scheme` to an attention-layer variant when the search should make one recipe +decision for self-attention q/k/v/o and/or linear-attention qkv/z/out. Recipes can splice a +shared base `disabled_layers` set via `$import` (see `modelopt_recipes/configs/auto_quantize/units/base_disabled_layers`). bf16 (no quantization) is always an implicit per-layer choice, so `candidate_formats` need only list @@ -407,7 +419,8 @@ The example scripts above also have an additional flag `--tasks`, where the actu > *If GPU out-of-memory error is reported running the scripts, please try editing the scripts and reducing the max batch size to save GPU memory.* -> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` method. The `kl_div` method does not require backpropagation.* +> *NOTE: AutoQuantize requires backpropagation when using the `gradient` method. Models without +> backpropagation support can use `group_recon` or `kl_div`, which are forward-only.* ## Real Quant diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 8dcc78afa27..07d75fdce8c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -347,6 +347,8 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) - "disabled_layers": aq_config.disabled_layers, "kv_cache_quant_cfg": kv_cache_quant_cfg, "method": aq_config.auto_quantize_method, + "quant_grouping_scheme": aq_config.quant_grouping_scheme, + "score_boundary": aq_config.score_boundary, "score_size": aq_config.score_size, } @@ -383,9 +385,12 @@ def _auto_quantize_config_from_cli(args: argparse.Namespace): effective_bits=args.auto_quantize_bits, cost_model=args.auto_quantize_cost_model, cost=cost, + score_model=args.auto_quantize_score_model, ), candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES[q]) for q in args.qformat.split(",")], auto_quantize_method=args.auto_quantize_method, + quant_grouping_scheme=args.auto_quantize_grouping_scheme, + score_boundary=args.auto_quantize_score_boundary, score_size=args.auto_quantize_score_size, disabled_layers=disabled_layers, cost_excluded_layers=cost_excluded_layers, @@ -439,7 +444,7 @@ def loss_func(output, data): def loss_func(output, data): return output.loss - if inputs["method"] == "gradient": + if inputs["method"] in {"gradient", "group_recon"}: def forward_step(model, batch): inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch @@ -457,7 +462,8 @@ def forward_step(model, batch): else: raise ValueError( - f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'" + f"Invalid auto_quantize method: {inputs['method']}. " + "Must be 'gradient', 'group_recon', or 'kl_div'" ) language_model, _ = mtq.auto_quantize( @@ -465,13 +471,15 @@ def forward_step(model, batch): constraints=inputs["constraints"], data_loader=calib_dataloader, forward_step=forward_step, - loss_func=loss_func, + loss_func=loss_func if inputs["method"] == "gradient" else None, 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"], method=inputs["method"], + quant_grouping_scheme=inputs["quant_grouping_scheme"], + score_boundary=inputs["score_boundary"], checkpoint=args.auto_quantize_checkpoint, ) @@ -1476,7 +1484,7 @@ def parse_args() -> argparse.Namespace: "--auto_quantize_method", type=str, default="gradient", - choices=["gradient", "kl_div"], + choices=["gradient", "group_recon", "kl_div"], help="[Deprecated: use an AutoQuantize --recipe] Sensitivity scoring method.", ) parser.add_argument( @@ -1485,6 +1493,32 @@ def parse_args() -> argparse.Namespace: default=128, help="[Deprecated: use an AutoQuantize --recipe] Number of samples for sensitivity scoring.", ) + parser.add_argument( + "--auto_quantize_score_model", + type=str, + default="raw", + choices=["raw", "per_element"], + help="[Deprecated: use an AutoQuantize --recipe] Selector score model.", + ) + parser.add_argument( + "--auto_quantize_score_boundary", + type=str, + default=None, + choices=["local", "group"], + help="[Deprecated: use an AutoQuantize --recipe] Sensitivity score boundary.", + ) + parser.add_argument( + "--auto_quantize_grouping_scheme", + type=str, + default="runtime_fused", + choices=[ + "runtime_fused", + "runtime_fused+linear_attn_layer", + "runtime_fused+self_attn_layer", + "runtime_fused+linear_attn_layer+self_attn_layer", + ], + help="[Deprecated: use an AutoQuantize --recipe] Quantization decision grouping.", + ) parser.add_argument( "--auto_quantize_cost_model", type=str, diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index 84057e468c9..f4729231291 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -111,6 +111,11 @@ if [ -n "$AUTO_QUANTIZE_BITS" ]; then PTQ_ARGS+=" --auto_quantize_bits=$AUTO_QUANTIZE_BITS " PTQ_ARGS+=" --auto_quantize_method=${AUTO_QUANTIZE_METHOD:-gradient} " PTQ_ARGS+=" --auto_quantize_score_size=${AUTO_QUANTIZE_SCORE_SIZE:-128} " + PTQ_ARGS+=" --auto_quantize_score_model=${AUTO_QUANTIZE_SCORE_MODEL:-raw} " + if [ -n "$AUTO_QUANTIZE_SCORE_BOUNDARY" ]; then + PTQ_ARGS+=" --auto_quantize_score_boundary=$AUTO_QUANTIZE_SCORE_BOUNDARY " + fi + PTQ_ARGS+=" --auto_quantize_grouping_scheme=${AUTO_QUANTIZE_GROUPING_SCHEME:-runtime_fused} " PTQ_ARGS+=" --auto_quantize_cost_model=${AUTO_QUANTIZE_COST_MODEL:-weight} " if [ -n "$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO" ]; then PTQ_ARGS+=" --auto_quantize_active_moe_expert_ratio=$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO " diff --git a/examples/hf_ptq/scripts/parser.sh b/examples/hf_ptq/scripts/parser.sh index 03ed3a57631..0a176660155 100644 --- a/examples/hf_ptq/scripts/parser.sh +++ b/examples/hf_ptq/scripts/parser.sh @@ -41,7 +41,7 @@ parse_options() { CALIB_WITH_IMAGES=false # Parse command-line options - ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@") + ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_score_model:,auto_quantize_score_boundary:,auto_quantize_grouping_scheme:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@") eval set -- "$ARGS" while true; do @@ -76,6 +76,9 @@ parse_options() { --auto_quantize_bits ) AUTO_QUANTIZE_BITS="$2"; shift 2;; --auto_quantize_method ) AUTO_QUANTIZE_METHOD="$2"; shift 2;; --auto_quantize_score_size ) AUTO_QUANTIZE_SCORE_SIZE="$2"; shift 2;; + --auto_quantize_score_model ) AUTO_QUANTIZE_SCORE_MODEL="$2"; shift 2;; + --auto_quantize_score_boundary ) AUTO_QUANTIZE_SCORE_BOUNDARY="$2"; shift 2;; + --auto_quantize_grouping_scheme ) AUTO_QUANTIZE_GROUPING_SCHEME="$2"; shift 2;; --auto_quantize_cost_model ) AUTO_QUANTIZE_COST_MODEL="$2"; shift 2;; --auto_quantize_active_moe_expert_ratio ) AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO="$2"; shift 2;; --moe_calib_experts_ratio ) MOE_CALIB_EXPERTS_RATIO="$2"; shift 2;; @@ -180,6 +183,9 @@ parse_options() { echo "auto_quantize_bits: $AUTO_QUANTIZE_BITS" echo "auto_quantize_method: $AUTO_QUANTIZE_METHOD" echo "auto_quantize_score_size: $AUTO_QUANTIZE_SCORE_SIZE" + echo "auto_quantize_score_model: $AUTO_QUANTIZE_SCORE_MODEL" + echo "auto_quantize_score_boundary: $AUTO_QUANTIZE_SCORE_BOUNDARY" + echo "auto_quantize_grouping_scheme: $AUTO_QUANTIZE_GROUPING_SCHEME" echo "auto_quantize_cost_model: $AUTO_QUANTIZE_COST_MODEL" echo "auto_quantize_active_moe_expert_ratio: $AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO" echo "moe_calib_experts_ratio: $MOE_CALIB_EXPERTS_RATIO" diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index 2cf5a2f0cfb..2b6d0603273 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -182,6 +182,12 @@ class AutoQuantizeConstraints(ModeloptBaseConfig): title="Cost-model parameters", description="Extra cost-model parameters; omit for the 'weight' cost model.", ) + score_model: Literal["raw", "per_element"] = ModeloptField( + default="raw", + title="Selector score model", + description="'raw' uses sensitivity scores directly; 'per_element' normalizes each " + "score by the represented weight-element cost before budgeted selection.", + ) @field_validator("effective_bits") @classmethod @@ -206,10 +212,30 @@ class AutoQuantizeConfig(ModeloptBaseConfig): "(e.g. [fp8]) yields a {fp8, bf16} per-layer search.", validate_default=True, ) - auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField( + auto_quantize_method: Literal["gradient", "group_recon", "kl_div"] = ModeloptField( default="gradient", title="Sensitivity scoring method", - description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).", + description="'gradient' (Taylor + Fisher, needs labels), 'group_recon' (normalized " + "group-output reconstruction, no labels), or 'kl_div' (no labels).", + ) + score_boundary: Literal["local", "group"] | None = ModeloptField( + default=None, + title="Sensitivity score boundary", + description="'local' scores each quantized module output; 'group' scores attention and " + "MoE projection perturbations at their shared attention/MLP output. Defaults to 'group' " + "for group_recon and 'local' otherwise.", + ) + quant_grouping_scheme: Literal[ + "runtime_fused", + "runtime_fused+linear_attn_layer", + "runtime_fused+self_attn_layer", + "runtime_fused+linear_attn_layer+self_attn_layer", + ] = ModeloptField( + default="runtime_fused", + title="Quantization decision grouping", + description="Search-decision grouping beyond mandatory runtime-fused groups. Attention-" + "layer options make all listed projections in one layer share a quantization recipe; " + "this is independent of score_boundary.", ) score_size: int = ModeloptField( default=128, @@ -248,6 +274,22 @@ def _at_least_one_candidate(cls, v: list[QuantizeConfig]) -> list[QuantizeConfig ) return v + @model_validator(mode="after") + def _validate_scoring_configuration(self): + boundary = self.score_boundary or ( + "group" if self.auto_quantize_method == "group_recon" else "local" + ) + if self.auto_quantize_method == "group_recon" and boundary != "group": + raise ValueError("auto_quantize_method='group_recon' requires score_boundary='group'.") + if self.auto_quantize_method == "kl_div" and ( + self.constraints.score_model != "raw" or boundary != "local" + ): + raise ValueError( + "auto_quantize_method='kl_div' requires constraints.score_model='raw' and " + "score_boundary='local'." + ) + return self + class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase): """Our config class for AutoQuantize recipes.""" diff --git a/modelopt/torch/quantization/_auto_quantize_cost.py b/modelopt/torch/quantization/_auto_quantize_cost.py index 297e8cd8dc3..7505e778f14 100644 --- a/modelopt/torch/quantization/_auto_quantize_cost.py +++ b/modelopt/torch/quantization/_auto_quantize_cost.py @@ -26,7 +26,14 @@ # constraint is supplied. The value is intentionally kept for backward compatibility. DEFAULT_AUTO_QUANTIZE_EFFECTIVE_BITS: Final = 4.8 -AUTO_QUANTIZE_CONSTRAINT_KEYS: Final = frozenset({"effective_bits", "cost_model", "cost"}) +AUTO_QUANTIZE_CONSTRAINT_KEYS: Final = frozenset( + {"effective_bits", "cost_model", "cost", "score_model"} +) +AUTO_QUANTIZE_SCORE_MODEL_RAW: Final = "raw" +AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT: Final = "per_element" +AUTO_QUANTIZE_SCORE_MODELS: Final = frozenset( + {AUTO_QUANTIZE_SCORE_MODEL_RAW, AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT} +) ACTIVE_MOE_EXPERT_RATIO_KEY: Final = "active_moe_expert_ratio" EXCLUDED_MODULE_NAME_PATTERNS_KEY: Final = "excluded_module_name_patterns" COST_MODEL_WEIGHT: Final = "weight" @@ -235,8 +242,15 @@ def normalize_auto_quantize_constraints( if unexpected_constraint_keys: raise ValueError( f"Unsupported auto_quantize constraints: {unexpected_constraint_keys}. " - "Supported constraints are 'effective_bits', 'cost_model', and 'cost'." + "Supported constraints are 'effective_bits', 'cost_model', 'cost', and 'score_model'." + ) + + score_model = constraints.get("score_model", AUTO_QUANTIZE_SCORE_MODEL_RAW) + if score_model not in AUTO_QUANTIZE_SCORE_MODELS: + raise ValueError( + f"constraints['score_model'] must be one of {sorted(AUTO_QUANTIZE_SCORE_MODELS)}." ) + constraints["score_model"] = score_model cost_model_name = constraints.get("cost_model", COST_MODEL_WEIGHT) if not isinstance(cost_model_name, str): diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index c7803ecf838..4c2d09dfc29 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -43,6 +43,9 @@ from ._auto_quantize_cost import ( ACTIVE_MOE_EXPERT_RATIO_KEY, AUTO_QUANTIZE_CONSTRAINT_KEYS, + AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT, + AUTO_QUANTIZE_SCORE_MODEL_RAW, + AUTO_QUANTIZE_SCORE_MODELS, COST_MODEL_ACTIVE_MOE, COST_MODEL_WEIGHT, _get_module_weight_numel, @@ -443,8 +446,36 @@ def attrs(self) -> list[str]: return ["name", "cost_weight", *super().attrs] -_LINEAR_ATTN_QKVZ_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_qkv|in_proj_z)$") -_LINEAR_ATTN_BA_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_a|in_proj_b)$") +_LINEAR_ATTN_QKVZ_RE = re.compile(r"^((?:.*\.)?linear_attn)\.(?:in_proj_qkv|in_proj_z)$") +_LINEAR_ATTN_BA_RE = re.compile(r"^((?:.*\.)?linear_attn)\.(?:in_proj_a|in_proj_b)$") +_LINEAR_ATTN_LAYER_GROUP_RE = re.compile( + r"^((?:.*\.)?linear_attn)\.(?:in_proj_qkv|in_proj_z|out_proj)$" +) +_SELF_ATTN_GROUP_RE = re.compile(r"^((?:.*\.)?self_attn)\.(?:q_proj|k_proj|v_proj|o_proj)$") +_LINEAR_ATTN_GROUP_RE = re.compile( + r"^((?:.*\.)?linear_attn)\.(?:in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|out_proj)$" +) +_FUSED_EXPERTS_GROUP_RE = re.compile(r"^((?:.*\.)?mlp)\.experts$") + +AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL = "local" +AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP = "group" +AUTO_QUANTIZE_SCORE_BOUNDARIES = frozenset( + {AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL, AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP} +) +AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED = "runtime_fused" +AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED_LINEAR_ATTN_LAYER = "runtime_fused+linear_attn_layer" +AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED_SELF_ATTN_LAYER = "runtime_fused+self_attn_layer" +AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED_LINEAR_SELF_ATTN_LAYER = ( + "runtime_fused+linear_attn_layer+self_attn_layer" +) +AUTO_QUANTIZE_GROUPING_SCHEMES = frozenset( + { + AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED, + AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED_LINEAR_ATTN_LAYER, + AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED_SELF_ATTN_LAYER, + AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED_LINEAR_SELF_ATTN_LAYER, + } +) def _linear_attn_qkvz_group_key(_model, name: str) -> str | None: @@ -457,6 +488,37 @@ def _linear_attn_ba_group_key(_model, name: str) -> str | None: return f"{m.group(1)}/ba" if m else None +def _linear_attn_layer_group_key(_model, name: str) -> str | None: + match = _LINEAR_ATTN_LAYER_GROUP_RE.match(name) + return f"{match.group(1)}/layer" if match else None + + +def _self_attn_layer_group_key(_model, name: str) -> str | None: + match = _SELF_ATTN_GROUP_RE.match(name) + return f"{match.group(1)}/layer" if match else None + + +def _self_attn_group_score_module(_model, name: str) -> str | None: + match = _SELF_ATTN_GROUP_RE.match(name) + return match.group(1) if match else None + + +def _linear_attn_group_score_module(_model, name: str) -> str | None: + match = _LINEAR_ATTN_GROUP_RE.match(name) + return match.group(1) if match else None + + +def _fused_experts_group_score_module(model, name: str) -> str | None: + match = _FUSED_EXPERTS_GROUP_RE.match(name) + if match is None: + return None + try: + module = model.get_submodule(name) + except AttributeError: + return None + return match.group(1) if _is_hf_quant_fused_experts_module(module) else None + + class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): """Base searcher for AutoQuantize algorithm.""" @@ -474,6 +536,8 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): r"^(.*?)\.(q_proj|k_proj|v_proj)$", # q_proj, k_proj, v_proj for llama like models # gate_proj, up_proj, down_proj for Qwen3 like MoE models r"^(.*?\.mlp\.experts)\.\d+\.(gate_proj|up_proj|down_proj)$", + # Keep shared-expert projections in one deployable fused-MoE decision. + r"^((?:.*\.)?mlp\.shared_expert)\.(gate_proj|up_proj|down_proj)$", r"^(.*?\.mixer\.experts)\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts # NemotronH MoE experts in MCore naming (linear_fc1=gate+up fused, linear_fc2=down) r"^(.*?\.mlp\.experts\.local_experts)\.\d+\.(linear_fc1|linear_fc2)$", @@ -506,6 +570,8 @@ def default_search_config(self): "cost_model": COST_MODEL_WEIGHT, "cost": {}, "active_moe_expert_ratio": None, + "quant_grouping_scheme": AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED, + "score_boundary": AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL, } @property @@ -516,6 +582,9 @@ def default_state_dict(self) -> SearchStateDict: "cost_model": "weight", "cost": {}, "active_moe_expert_ratio": None, + "score_model": AUTO_QUANTIZE_SCORE_MODEL_RAW, + "quant_grouping_scheme": AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED, + "score_boundary": AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL, "cost_denominator": None, "disabled_layers": None, "candidate_stats": defaultdict(dict), @@ -533,8 +602,41 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: assert config["forward_step"] is not None, ( "`forward_step` must be provided for `auto_quantize`." ) + if config["score_boundary"] not in AUTO_QUANTIZE_SCORE_BOUNDARIES: + raise ValueError( + f"score_boundary must be one of {sorted(AUTO_QUANTIZE_SCORE_BOUNDARIES)}." + ) + if config["quant_grouping_scheme"] not in AUTO_QUANTIZE_GROUPING_SCHEMES: + raise ValueError( + f"quant_grouping_scheme must be one of {sorted(AUTO_QUANTIZE_GROUPING_SCHEMES)}." + ) return config + def _get_quant_grouping_rules(self): + rules: list[Any] = [] + scheme = self.config["quant_grouping_scheme"] + if "linear_attn_layer" in scheme: + # Keep A/B in their runtime-required pair; they are not part of the + # deployable qkv/z/out family decision and are commonly disabled. + rules.append(_linear_attn_layer_group_key) + if "self_attn_layer" in scheme: + rules.append(_self_attn_layer_group_key) + rules.extend(self.quant_grouping_rules) + return rules + + def _get_score_module_rules(self): + rules = [] + if self.config["score_boundary"] == AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP: + rules.extend( + [ + _self_attn_group_score_module, + _linear_attn_group_score_module, + _fused_experts_group_score_module, + ] + ) + rules.extend(self.score_module_rules) + return rules + def load_search_checkpoint(self) -> bool: return super().load_search_checkpoint(strict=False) @@ -653,7 +755,7 @@ def insert_hparams_after_merge_rules(self, model, quant_recipes, disabled_layers # Apply quant_grouping_rules to determine the group key group_key = name # Default: each module in its own group - for rule in self.quant_grouping_rules: + for rule in self._get_quant_grouping_rules(): result = self._apply_quant_group_rule(name, rule) if result is not None: group_key = result @@ -662,7 +764,7 @@ def insert_hparams_after_merge_rules(self, model, quant_recipes, disabled_layers # Apply score_module_rules to determine the score module name, then get the actual module score_module_name = name # Default: score from same module - for rule in self.score_module_rules: + for rule in self._get_score_module_rules(): result = self._apply_score_group_rule(name, rule) if result is not None: score_module_name = result @@ -725,22 +827,25 @@ def initialize_candidate_stats(self): if not isinstance(hparam, QuantRecipeHparam): continue - formats, scores, costs = [], [], [] + formats, scores, costs, element_costs = [], [], [], [] prev_score = float("inf") for recipe in hparam.choices: formats.append(recipe) score = hparam.get_score(recipe) # type: ignore [arg-type] cost = hparam.get_cost(recipe) # type: ignore [arg-type] + element_cost = hparam.get_cost(recipe, cost_weight=1.0) # type: ignore [arg-type] score = min(score, prev_score) # TODO: Should we get rid of this? scores.append(score) costs.append(cost) + element_costs.append(element_cost) prev_score = score self.candidate_stats[name]["formats"] = formats self.candidate_stats[name]["scores"] = scores self.candidate_stats[name]["costs"] = costs + self.candidate_stats[name]["element_costs"] = element_costs self.candidate_stats[name]["module_names"] = hparam.quant_module_names self.candidate_stats[name]["quantizer_attrs"] = hparam.quant_module_replay_attrs self.candidate_stats[name]["cost_weight"] = hparam.cost_weight @@ -763,6 +868,14 @@ def before_search(self): super().before_search() self.constraints = normalize_auto_quantize_constraints(self.model, self.constraints) + if self.method_name not in {"gradient", "group_recon"} and ( + self.constraints["score_model"] != AUTO_QUANTIZE_SCORE_MODEL_RAW + or self.config["score_boundary"] != AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL + ): + raise ValueError( + "score_model='per_element' and score_boundary='group' are supported only " + "with method='gradient' or method='group_recon'." + ) self.config["cost_model"] = self.constraints["cost_model"] self.config["cost"] = self.constraints.get("cost", {}) self.config["active_moe_expert_ratio"] = self.config["cost"].get( @@ -777,6 +890,14 @@ def before_search(self): ) restored_cost_model = getattr(self, "cost_model", "weight") restored_active_moe_expert_ratio = getattr(self, "active_moe_expert_ratio", None) + restored_score_boundary = getattr( + self, "score_boundary", AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL + ) + restored_quant_grouping_scheme = getattr( + self, + "quant_grouping_scheme", + AUTO_QUANTIZE_GROUPING_SCHEME_RUNTIME_FUSED, + ) if self.candidate_stats and ( restored_cost_model != self.config["cost_model"] or restored_active_moe_expert_ratio != self.config["active_moe_expert_ratio"] @@ -787,10 +908,28 @@ def before_search(self): f"current=({self.config['cost_model']}, {self.config['active_moe_expert_ratio']}). " "Use a different checkpoint path." ) + if self.candidate_stats and restored_score_boundary != self.config["score_boundary"]: + raise ValueError( + "Checkpoint AutoQuantize score boundary does not match current search config: " + f"checkpoint={restored_score_boundary}, current={self.config['score_boundary']}. " + "Use a different checkpoint path." + ) + if ( + self.candidate_stats + and restored_quant_grouping_scheme != self.config["quant_grouping_scheme"] + ): + raise ValueError( + "Checkpoint AutoQuantize quant grouping scheme does not match current search " + f"config: checkpoint={restored_quant_grouping_scheme}, " + f"current={self.config['quant_grouping_scheme']}. Use a different checkpoint path." + ) self.method = self.method_name self.cost_model = self.config["cost_model"] self.cost = self.config["cost"] self.active_moe_expert_ratio = self.config["active_moe_expert_ratio"] + self.score_model = self.constraints["score_model"] + self.quant_grouping_scheme = self.config["quant_grouping_scheme"] + self.score_boundary = self.config["score_boundary"] self.disabled_layers = self.config["disabled_layers"] self.cost_denominator = getattr(self, "cost_denominator", None) @@ -920,7 +1059,8 @@ def run_search(self): assert "effective_bits" in self.constraints and ( set(self.constraints) <= AUTO_QUANTIZE_CONSTRAINT_KEYS ), ( - "`constraints` must contain 'effective_bits' and may contain 'cost_model' and 'cost'. " + "`constraints` must contain 'effective_bits' and may contain 'cost_model', 'cost', " + "and 'score_model'. " f"Got {self.constraints.keys()}." ) @@ -975,7 +1115,10 @@ def run_search(self): effective_bits_from_search = (best_constraints / total_weight_size) * 16 self.best["recipe"] = best_recipe - self.best["constraints"] = {"effective_bits": effective_bits_from_search} + self.best["constraints"] = { + "effective_bits": effective_bits_from_search, + "score_model": self.constraints.get("score_model", AUTO_QUANTIZE_SCORE_MODEL_RAW), + } self.best["score"] = best_scores QuantRecipe.fold_pqs_to_weights(self.model) @@ -1027,6 +1170,10 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): score_module_rules = [ # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts) r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", + # Preserve the historical local-score boundary for shared experts. The + # projections are one semantic path, so score their combined effect at + # the parent MLP output even when attention uses leaf-local scoring. + r"^((?:.*\.)?mlp)\.shared_expert\.(gate_proj|up_proj|down_proj)$", r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts @@ -1105,12 +1252,12 @@ def forward_backward_step(model, data): @torch.enable_grad() def _estimate_auto_quantize_scores(self, is_param_grad_enabled): # TODO: remove the no-quant recipe - def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): + def auto_quantize_score_estimate_forward(module, *args, **kwargs): for hparam in module._hparams_for_scoring: if hparam.is_configurable: hparam.active = QuantRecipe(quant_cfg=None) - output = module._forward_original(input, *args, **kwargs) + output = module._forward_original(*args, **kwargs) # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass. # With gradient checkpointing, gradients are computed in the local forward pass during backward pass @@ -1128,7 +1275,7 @@ def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): if recipe == QuantRecipe(quant_cfg=None): continue hparam.active = recipe - output_diff = module._forward_original(input, *args, **kwargs) + output_diff = module._forward_original(*args, **kwargs) if isinstance(output_diff, tuple): output_diff = output_diff[0] - output[0] @@ -1258,6 +1405,32 @@ def run_search_with_stats(self, max_weight_size, verbose=False): max_weight_size, lower_bound ) + candidate_scores = [ + self._candidate_scores_for_search(candidate_stat) + for candidate_stat in self.candidate_stats.values() + ] + requires_objective_rescaling = ( + getattr(self, "constraints", {}).get("score_model", AUTO_QUANTIZE_SCORE_MODEL_RAW) + == AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT + or self.method_name == "group_recon" + ) + max_abs_score = ( + max( + (abs(float(score)) for scores in candidate_scores for score in scores), + default=0.0, + ) + if requires_objective_rescaling + else 0.0 + ) + if max_abs_score > 0.0: + # Per-element coefficients can fall below CBC's objective + # tolerance. Group reconstruction scores are normalized and + # can be similarly small. Global rescaling preserves the + # optimum without changing the default gradient objective. + candidate_scores = [ + [score / max_abs_score for score in scores] for scores in candidate_scores + ] + lps = LPS( name="AutoQuantize", constraints=constraints, @@ -1266,9 +1439,7 @@ def run_search_with_stats(self, max_weight_size, verbose=False): candidate_stat["costs"] for candidate_stat in self.candidate_stats.values() ] }, - candidate_scores=[ - candidate_stat["scores"] for candidate_stat in self.candidate_stats.values() - ], + candidate_scores=candidate_scores, objective_type="minimize", verbose=verbose, ) @@ -1294,6 +1465,160 @@ def run_search_with_stats(self, max_weight_size, verbose=False): return best_recipes, is_satisfied + def _candidate_scores_for_search(self, candidate_stat: dict[str, Any]) -> list[float]: + """Return objective coefficients for the configured score model.""" + score_model = getattr(self, "constraints", {}).get( + "score_model", AUTO_QUANTIZE_SCORE_MODEL_RAW + ) + if score_model == AUTO_QUANTIZE_SCORE_MODEL_RAW: + return candidate_stat["scores"] + if score_model == AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT: + element_costs = candidate_stat.get("element_costs") + if element_costs is None: + cost_weight = candidate_stat.get("cost_weight", 1.0) + element_costs = [ + cost / cost_weight if cost_weight > 0 else cost + for cost in candidate_stat["costs"] + ] + return [ + score / cost if cost > 0 else score + for score, cost in zip(candidate_stat["scores"], element_costs) + ] + raise ValueError( + f"Unsupported AutoQuantize score_model: {score_model}. " + f"Expected one of {sorted(AUTO_QUANTIZE_SCORE_MODELS)}." + ) + + +def _get_primary_output_tensor(output: Any) -> torch.Tensor: + """Return the tensor carrying a score module's hidden states.""" + if isinstance(output, torch.Tensor): + return output + if isinstance(output, tuple | list) and output and isinstance(output[0], torch.Tensor): + return output[0] + raise TypeError( + "Group reconstruction scoring expects a Tensor or a tuple/list whose first item " + f"is a Tensor, got {type(output)!r}." + ) + + +def _get_group_reconstruction_score(reference_output: Any, quantized_output: Any) -> torch.Tensor: + """Return normalized MSE between reference and quantized score-group outputs.""" + reference = _get_primary_output_tensor(reference_output) + quantized = _get_primary_output_tensor(quantized_output) + if reference.shape != quantized.shape: + raise ValueError( + "Reference and quantized score-group outputs must have the same shape, got " + f"{tuple(reference.shape)} and {tuple(quantized.shape)}." + ) + reference = reference.float() + quantized = quantized.float() + return (quantized - reference).square().mean() / reference.square().mean().clamp_min(1e-12) + + +def _get_model_config_objects(model: nn.Module) -> list[Any]: + """Return unique model configs that may carry generation cache defaults.""" + configs = [] + seen = set() + for module in model.modules(): + for attr_name in ("config", "generation_config"): + config = getattr(module, attr_name, None) + if config is None or id(config) in seen: + continue + seen.add(id(config)) + configs.append(config) + for nested_name in ("text_config", "language_config"): + nested = getattr(config, nested_name, None) + if nested is not None and id(nested) not in seen: + seen.add(id(nested)) + configs.append(nested) + return configs + + +def _set_model_use_cache(model: nn.Module, use_cache: bool) -> list[tuple[Any, Any]]: + """Set use_cache on model configs and return values to restore.""" + originals = [] + for config in _get_model_config_objects(model): + if hasattr(config, "use_cache"): + originals.append((config, config.use_cache)) + config.use_cache = use_cache + return originals + + +def _restore_model_use_cache(originals: list[tuple[Any, Any]]) -> None: + for config, use_cache in originals: + config.use_cache = use_cache + + +class AutoQuantizeGroupReconSearcher(AutoQuantizeGradientSearcher): + """AutoQuantize searcher using normalized score-group reconstruction error.""" + + method_name = "group_recon" + + def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: + """Ignore backward-only inputs and require the group score boundary.""" + config = dict(config or {}) + for ignored_key in ("score_func", "loss_func", "forward_backward_step"): + if config.get(ignored_key) is not None: + warnings.warn(f"`{ignored_key}` is ignored for group_recon auto_quantize.") + config.pop(ignored_key, None) + config = _AutoQuantizeBaseSearcher.sanitize_search_config(self, config) + if config["score_boundary"] != AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP: + raise ValueError("method='group_recon' requires score_boundary='group'.") + return config + + @torch.inference_mode() + def estimate_sensitivity_scores(self) -> None: + """Measure each recipe's normalized reconstruction error at score modules.""" + + def score_forward(module, *args, **kwargs): + hparams = [h for h in module._hparams_for_scoring if h.is_configurable] + no_quant = QuantRecipe(quant_cfg=None) + for hparam in hparams: + hparam.active = no_quant + reference_output = module._forward_original(*args, **kwargs) + + for hparam in hparams: + try: + for recipe in hparam.choices: + if recipe == no_quant: + continue + hparam.active = recipe + quantized_output = module._forward_original(*args, **kwargs) + score = _get_group_reconstruction_score(reference_output, quantized_output) + importance = hparam._importance_dict[recipe][module] + hparam._importance_dict[recipe][module] = ( + score if importance is None else importance + score + ) + finally: + hparam.active = no_quant + return reference_output + + score_modules = set() + for module in self.model.modules(): + if hasattr(module, "_hparams_for_scoring") and any( + hparam.is_configurable for hparam in module._hparams_for_scoring + ): + module._forward_original = module.forward + module.forward = types.MethodType(score_forward, module) + score_modules.add(module) + + self.model.eval() + cache_originals = _set_model_use_cache(self.model, False) + try: + self._run_func( + self.config["forward_step"], + num_iters=self.config["num_score_steps"], + desc="Estimating group reconstruction scores", + ) + finally: + _restore_model_use_cache(cache_originals) + for module in score_modules: + module.forward = module._forward_original + del module._forward_original + + gc.collect() + @torch.compile(dynamic=True) def _get_log_softmax_dist(logits: torch.Tensor, tp_group) -> torch.Tensor: @@ -1625,11 +1950,14 @@ def _resolve_best_recipe(search_state, constraints, verbose=False): if method == "gradient": searcher = AutoQuantizeGradientSearcher() + elif method == "group_recon": + searcher = AutoQuantizeGroupReconSearcher() elif method == "kl_div": searcher = AutoQuantizeKLDivSearcher() else: raise ValueError( - f"Unknown autoquant search method: {method!r}. Expected 'gradient' or 'kl_div'." + f"Unknown autoquant search method: {method!r}. " + "Expected 'gradient', 'group_recon', or 'kl_div'." ) searcher.candidate_stats = candidate_stats @@ -1648,6 +1976,17 @@ def _resolve_best_recipe(search_state, constraints, verbose=False): "cost": searcher.cost, "active_moe_expert_ratio": searcher.active_moe_expert_ratio, } + score_model = constraints.get( + "score_model", search_state.get("score_model", AUTO_QUANTIZE_SCORE_MODEL_RAW) + ) + if score_model not in AUTO_QUANTIZE_SCORE_MODELS: + raise ValueError( + f"constraints['score_model'] must be one of {sorted(AUTO_QUANTIZE_SCORE_MODELS)}." + ) + searcher.constraints = { + "effective_bits": effective_bits, + "score_model": score_model, + } best_recipe_info, _ = searcher.run_search_with_stats(max_weight_size, verbose=verbose) best_recipe = {name: info["format"] for name, info in best_recipe_info.items()} diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 7dbdd36d04e..e2cf63b11c9 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -36,7 +36,12 @@ ) from modelopt.torch.utils import atomic_print -from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe +from .algorithms import ( + AutoQuantizeGradientSearcher, + AutoQuantizeGroupReconSearcher, + AutoQuantizeKLDivSearcher, + QuantRecipe, +) from .algorithms import get_auto_quantize_config as _get_auto_quantize_config from .config import QuantizeAlgoCfgType from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg @@ -282,13 +287,16 @@ def auto_quantize( num_score_steps: int = 128, verbose: bool = False, method: str = "gradient", + quant_grouping_scheme: str = "runtime_fused", + score_boundary: str | None = None, checkpoint: str | None = None, ): r"""Perform optimal per-layer quantization by searching for the best quantization formats per-layer. ``auto_quantize`` uses sensitivity scores to rank the per-layer quantization formats and search for the best quantization formats per-layer. The sensitivity score can be computed using gradient-based - methods (default) or KL divergence loss, controlled by the ``method`` parameter. + methods (default), group-output reconstruction, or KL divergence loss, controlled by the + ``method`` parameter. Internally this API runs two main phases: @@ -318,12 +326,17 @@ def auto_quantize( constraints = { "effective_bits": 4.8, "cost_model": "active_moe", + "score_model": "per_element", "cost": { "active_moe_expert_ratio": 0.25, "excluded_module_name_patterns": ["*visual*", "*vision_tower*", "*mtp*"], }, } + ``score_model="per_element"`` normalizes each sensitivity score by the + number of weight elements represented by that search decision. The default + ``"raw"`` preserves the original objective. + quantization_formats: A list of quantization format config dictionaries or string names to search for. Each config dictionary should be valid as a ``config`` argument in :meth:`quantize `. @@ -440,9 +453,21 @@ def forward_backward_step(model, batch) -> None: verbose: If True, prints the search progress/intermediate results. method: Method to use for estimating sensitivity loss. Higher loss indicates greater sensitivity to quantization. Options are ``"gradient"`` (default; uses gradient-based loss estimation, - linear programming search, and requires ``loss_func`` or ``forward_backward_step``) and + linear programming search, and requires ``loss_func`` or ``forward_backward_step``), + ``"group_recon"`` (uses normalized group-output reconstruction and requires only + ``forward_step``), and ``"kl_div"`` (uses KL divergence between unquantized and quantized outputs, relies on threshold-based binary search, and only requires ``forward_step`` returning logits). + quant_grouping_scheme: Search-decision grouping. ``"runtime_fused"`` enforces only + mandatory runtime fusion groups. The optional ``linear_attn_layer`` and + ``self_attn_layer`` suffixes group each full attention projection family into one + recipe decision. This is independent of ``score_boundary``. + score_boundary: Boundary used to measure perturbations. ``"local"`` + preserves the existing per-module behavior. ``"group"`` scores attention + projections at their shared self-attention or linear-attention group output and scores + fused/shared MoE projections at the MLP group output. This does not group recipe + decisions or force modules to use the same quantization format. Defaults to + ``"group"`` for ``method="group_recon"`` and ``"local"`` otherwise. checkpoint: (Optional) Path to checkpoint file for saving/restoring auto_quantize search state. If the checkpoint file exists, the search state will be restored from it, skipping the expensive score estimation step. @@ -522,10 +547,16 @@ def forward_backward_step(model, batch) -> None: # Select the appropriate searcher based on method if method == "gradient": searcher = AutoQuantizeGradientSearcher() + elif method == "group_recon": + searcher = AutoQuantizeGroupReconSearcher() elif method == "kl_div": searcher = AutoQuantizeKLDivSearcher() else: - raise ValueError(f"Invalid method: {method}. Valid options are 'gradient' or 'kl_div'.") + raise ValueError( + f"Invalid method: {method}. Valid options are 'gradient', 'group_recon', or 'kl_div'." + ) + + score_boundary = score_boundary or ("group" if method == "group_recon" else "local") model = apply_mode( model, @@ -540,6 +571,8 @@ def forward_backward_step(model, batch) -> None: "forward_backward_step": forward_backward_step, "num_calib_steps": num_calib_steps, "num_score_steps": num_score_steps, + "quant_grouping_scheme": quant_grouping_scheme, + "score_boundary": score_boundary, "disabled_layers": disabled_layers, "verbose": verbose, "checkpoint": checkpoint, diff --git a/modelopt_recipes/configs/numerics/nvfp4_static.yaml b/modelopt_recipes/configs/numerics/nvfp4_static.yaml index 9f6ac62e11e..d4f47210a17 100644 --- a/modelopt_recipes/configs/numerics/nvfp4_static.yaml +++ b/modelopt_recipes/configs/numerics/nvfp4_static.yaml @@ -21,3 +21,6 @@ block_sizes: -1: 16 type: static scale_bits: e4m3 +# Autoquant LP cost: 4 value bits + an FP8 (8-bit) scale per 16-element block = 4.5 bits/element +# (the num_bits heuristic under-counts this as 4.0). Read only by autoquant. +effective_bits: 4.5 diff --git a/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_5p5bits-active_moe-group_recon.yaml b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_5p5bits-active_moe-group_recon.yaml new file mode 100644 index 00000000000..b54d397d9af --- /dev/null +++ b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_5p5bits-active_moe-group_recon.yaml @@ -0,0 +1,53 @@ +# 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. + +# Qwen3.6 MoE AutoQuantize: normalized group-output reconstruction scoring over +# mixed FP8 + NVFP4 weight-only candidates with active-MoE cost accounting. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + fp8: configs/ptq/presets/model/fp8 + w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4 + +metadata: + recipe_type: auto_quantize + description: >- + Qwen3.6 MoE: FP8 + NVFP4-weight-only search at 5.5 effective bits using + per-element normalized group reconstruction scores and active-MoE cost accounting. + +auto_quantize: + constraints: + effective_bits: 5.5 + cost_model: active_moe + cost: + active_moe_expert_ratio: 0.03125 + score_model: per_element + + candidate_formats: + - $import: fp8 + - $import: w4a16_nvfp4 + + auto_quantize_method: group_recon + score_boundary: group + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + + cost_excluded_layers: + - "*visual*" + - "*mtp*" + - "*vision_tower*" diff --git a/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient-local.yaml b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient-local.yaml new file mode 100644 index 00000000000..de4cf9c8507 --- /dev/null +++ b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient-local.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Qwen3.6 MoE P0 reproduction lane. The HF PTQ command selects the long-code +# calibration data; this recipe captures the old P0 score and allocation axes. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + fp8: configs/ptq/presets/model/fp8 + w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4 + +metadata: + recipe_type: auto_quantize + description: >- + Qwen3.6 MoE: local-gradient scores with grouped FP8/W4A16 NVFP4 decisions at + 6.0 active-MoE effective bits, matching the confirmed P0 score provenance. + +auto_quantize: + constraints: + effective_bits: 6.0 + cost_model: active_moe + cost: + active_moe_expert_ratio: 0.03125 + score_model: per_element + + candidate_formats: + - $import: fp8 + - $import: w4a16_nvfp4 + + auto_quantize_method: gradient + score_boundary: local + quant_grouping_scheme: runtime_fused+linear_attn_layer+self_attn_layer + score_size: 64 + + disabled_layers: + - $import: base_disabled_layers + + cost_excluded_layers: + - "*visual*" + - "*mtp*" + - "*vision_tower*" diff --git a/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient.yaml b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient.yaml new file mode 100644 index 00000000000..725cc277956 --- /dev/null +++ b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient.yaml @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Qwen3.6 MoE parent-boundary diagnostic: long-code calibration is selected by +# the HF PTQ command, while this recipe captures the grouped search policy. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + fp8: configs/ptq/presets/model/fp8 + w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4 + +metadata: + recipe_type: auto_quantize + description: >- + Qwen3.6 MoE: grouped parent-boundary gradient search over FP8 and W4A16 NVFP4 + at 6.0 active-MoE effective bits. This tests parent-output scoring; the confirmed + P0 source state used local gradient scores instead. + +auto_quantize: + constraints: + effective_bits: 6.0 + cost_model: active_moe + cost: + active_moe_expert_ratio: 0.03125 + score_model: per_element + + candidate_formats: + - $import: fp8 + - $import: w4a16_nvfp4 + + auto_quantize_method: gradient + score_boundary: group + quant_grouping_scheme: runtime_fused+linear_attn_layer+self_attn_layer + score_size: 64 + + disabled_layers: + - $import: base_disabled_layers + + cost_excluded_layers: + - "*visual*" + - "*mtp*" + - "*vision_tower*" diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 7160dd38cb4..a5900a9e679 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -53,9 +53,15 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): aq = load_recipe("general/auto_quantize/nvfp4_fp8_at_5p4bits").auto_quantize inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) - assert inputs["constraints"] == {"effective_bits": 5.4, "cost_model": "weight"} + assert inputs["constraints"] == { + "effective_bits": 5.4, + "cost_model": "weight", + "score_model": "raw", + } assert inputs["kv_cache_quant_cfg"] is None assert inputs["method"] == "gradient" + assert inputs["quant_grouping_scheme"] == "runtime_fused" + assert inputs["score_boundary"] is None assert inputs["score_size"] == 128 # disabled_layers come straight from the recipe (no model introspection). assert inputs["disabled_layers"] == aq.disabled_layers @@ -148,6 +154,14 @@ def test_autoquant_config_from_deprecated_cli_flags(monkeypatch): "active_moe", "--auto_quantize_active_moe_expert_ratio", "0.03125", + "--auto_quantize_method", + "group_recon", + "--auto_quantize_score_model", + "per_element", + "--auto_quantize_score_boundary", + "group", + "--auto_quantize_grouping_scheme", + "runtime_fused+linear_attn_layer+self_attn_layer", "--kv_cache_qformat", "none", ) @@ -156,7 +170,10 @@ def test_autoquant_config_from_deprecated_cli_flags(monkeypatch): assert aq.constraints.effective_bits == 5.4 assert aq.constraints.cost_model == "active_moe" assert aq.constraints.cost.active_moe_expert_ratio == 0.03125 - assert aq.auto_quantize_method == "gradient" + assert aq.constraints.score_model == "per_element" + assert aq.auto_quantize_method == "group_recon" + assert aq.quant_grouping_scheme == "runtime_fused+linear_attn_layer+self_attn_layer" + assert aq.score_boundary == "group" assert aq.score_size == 128 # candidates come from --qformat and resolve to their shipped presets. assert [hf_ptq._match_candidate_to_preset(f)[0] for f in aq.candidate_formats] == [ @@ -166,3 +183,114 @@ def test_autoquant_config_from_deprecated_cli_flags(monkeypatch): # base disabled + base cost-excluded appended from the shared units (no introspection). assert "*output_layer*" in aq.disabled_layers assert aq.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"] + + +def test_autoquant_group_scoring_recipe_matches_deprecated_cli(monkeypatch): + """The deprecated CLI shim and the shipped group-scoring recipe map to identical mtq inputs.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--qformat", + "fp8,w4a16_nvfp4", + "--auto_quantize_bits", + "5.5", + "--auto_quantize_method", + "group_recon", + "--auto_quantize_score_model", + "per_element", + "--auto_quantize_score_boundary", + "group", + "--auto_quantize_grouping_scheme", + "runtime_fused", + "--auto_quantize_cost_model", + "active_moe", + "--auto_quantize_active_moe_expert_ratio", + "0.03125", + "--kv_cache_qformat", + "none", + ) + recipe_config = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_5p5bits-active_moe-group_recon" + ).auto_quantize + cli_config = hf_ptq._auto_quantize_config_from_cli(args) + + assert hf_ptq._mtq_inputs_from_auto_quantize_config( + recipe_config, args + ) == hf_ptq._mtq_inputs_from_auto_quantize_config(cli_config, args) + + +def test_autoquant_grouped_gradient_recipe_matches_deprecated_cli(monkeypatch): + """The grouped parent-score recipe and deprecated CLI map to identical search inputs.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--qformat", + "fp8,w4a16_nvfp4", + "--auto_quantize_bits", + "6.0", + "--auto_quantize_method", + "gradient", + "--auto_quantize_score_model", + "per_element", + "--auto_quantize_score_boundary", + "group", + "--auto_quantize_grouping_scheme", + "runtime_fused+linear_attn_layer+self_attn_layer", + "--auto_quantize_score_size", + "64", + "--auto_quantize_cost_model", + "active_moe", + "--auto_quantize_active_moe_expert_ratio", + "0.03125", + "--kv_cache_qformat", + "none", + ) + recipe_config = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/" + "w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient" + ).auto_quantize + cli_config = hf_ptq._auto_quantize_config_from_cli(args) + + assert hf_ptq._mtq_inputs_from_auto_quantize_config( + recipe_config, args + ) == hf_ptq._mtq_inputs_from_auto_quantize_config(cli_config, args) + + +def test_autoquant_grouped_local_gradient_recipe_matches_deprecated_cli(monkeypatch): + """The local-score P0 recipe and deprecated CLI map to identical grouped search inputs.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--qformat", + "fp8,w4a16_nvfp4", + "--auto_quantize_bits", + "6.0", + "--auto_quantize_method", + "gradient", + "--auto_quantize_score_model", + "per_element", + "--auto_quantize_score_boundary", + "local", + "--auto_quantize_grouping_scheme", + "runtime_fused+linear_attn_layer+self_attn_layer", + "--auto_quantize_score_size", + "64", + "--auto_quantize_cost_model", + "active_moe", + "--auto_quantize_active_moe_expert_ratio", + "0.03125", + "--kv_cache_qformat", + "none", + ) + recipe_config = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/" + "w4a16_nvfp4_fp8_at_6p0bits-active_moe-grouped-gradient-local" + ).auto_quantize + cli_config = hf_ptq._auto_quantize_config_from_cli(args) + + assert hf_ptq._mtq_inputs_from_auto_quantize_config( + recipe_config, args + ) == hf_ptq._mtq_inputs_from_auto_quantize_config(cli_config, args) diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 3aaacaa3e0e..40541c14320 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -27,6 +27,8 @@ import modelopt.torch.quantization.config as qcfg from modelopt.recipe.config import ( + AutoQuantizeConfig, + AutoQuantizeConstraints, ModelOptAutoQuantizeRecipe, ModelOptDFlashRecipe, ModelOptEagleRecipe, @@ -35,7 +37,11 @@ ) from modelopt.recipe.loader import _apply_dotlist, load_config, load_recipe from modelopt.torch.opt.config_loader import _load_raw_config, _schema_type -from modelopt.torch.quantization.config import QuantizerAttributeConfig, normalize_quant_cfg_list +from modelopt.torch.quantization.config import ( + QuantizeConfig, + QuantizerAttributeConfig, + normalize_quant_cfg_list, +) # --------------------------------------------------------------------------- # Static YAML fixtures @@ -1720,11 +1726,13 @@ def test_load_recipe_autoquantize_minimal(tmp_path): assert isinstance(recipe, ModelOptAutoQuantizeRecipe) aq = recipe.auto_quantize assert aq.auto_quantize_method == "gradient" + assert aq.score_boundary is None assert aq.score_size == 128 assert aq.kv_cache is None assert aq.constraints.effective_bits == 4.8 assert aq.constraints.cost_model == "weight" assert aq.constraints.cost is None + assert aq.constraints.score_model == "raw" assert len(aq.candidate_formats) == 2 @@ -1753,6 +1761,7 @@ def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path): "effective_bits": 6.0, "cost_model": "active_moe", "cost": {"active_moe_expert_ratio": 0.03125}, + "score_model": "raw", } @@ -1798,6 +1807,57 @@ def test_load_recipe_autoquantize_effective_bits_out_of_range_raises(tmp_path): load_recipe(bad) +def test_load_recipe_autoquantize_group_scoring(tmp_path): + recipe_file = tmp_path / "group.yml" + recipe_file.write_text( + _AQ_MINIMAL_BODY.replace( + " effective_bits: 4.8\n", + " effective_bits: 4.8\n score_model: per_element\n", + ) + + " auto_quantize_method: group_recon\n" + + " score_boundary: group\n" + ) + aq = load_recipe(recipe_file).auto_quantize + assert aq.constraints.score_model == "per_element" + assert aq.auto_quantize_method == "group_recon" + assert aq.score_boundary == "group" + + +@pytest.mark.parametrize( + ("method", "score_model", "boundary", "message"), + [ + ("group_recon", "raw", "local", "requires score_boundary='group'"), + ("kl_div", "per_element", "local", "requires constraints.score_model='raw'"), + ("kl_div", "raw", "group", "requires constraints.score_model='raw'"), + ], +) +def test_load_recipe_autoquantize_rejects_invalid_scoring_combinations( + method, score_model, boundary, message +): + with pytest.raises(ValueError, match=message): + AutoQuantizeConfig( + constraints=AutoQuantizeConstraints( + effective_bits=5.4, + score_model=score_model, + ), + candidate_formats=[QuantizeConfig(quant_cfg=[])], + auto_quantize_method=method, + score_boundary=boundary, + ) + + +def test_load_recipe_autoquantize_builtin_group_scoring(): + recipe = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_5p5bits-active_moe-group_recon" + ) + aq = recipe.auto_quantize + assert aq.constraints.effective_bits == 5.5 + assert aq.constraints.cost_model == "active_moe" + assert aq.constraints.score_model == "per_element" + assert aq.auto_quantize_method == "group_recon" + assert aq.score_boundary == "group" + + def test_load_recipe_autoquantize_builtin_active_moe(): """The shipped active-MoE AutoQuantize recipe resolves to the expected values.""" recipe = load_recipe("general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe") diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index b85feb32649..f6e3641af64 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -24,6 +24,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.algorithms as quant_algorithms from modelopt.torch.quantization._auto_quantize_cost import ( EXCLUDED_MODULE_NAME_PATTERNS_KEY, _get_module_weight_numel, @@ -71,6 +72,40 @@ def get_input(self): return torch.randn(1, 4, 32) +class _LinearAttentionLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.in_proj_qkv = torch.nn.Linear(32, 32) + self.in_proj_z = torch.nn.Linear(32, 32) + self.in_proj_a = torch.nn.Linear(32, 32) + self.in_proj_b = torch.nn.Linear(32, 32) + self.out_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + x = self.in_proj_qkv(x) + self.in_proj_z(x) + x = x + self.in_proj_a(x) + self.in_proj_b(x) + return self.out_proj(x) + + +class _GroupBoundaryModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + use_cache=True, + text_config=SimpleNamespace(use_cache=True), + ) + self.use_cache_seen = [] + self.self_attn = _AttentionLayer() + self.linear_attn = _LinearAttentionLayer() + + def forward(self, x): + self.use_cache_seen.append((self.config.use_cache, self.config.text_config.use_cache)) + return self.linear_attn(x=self.self_attn(x=x)) + + def get_input(self): + return torch.randn(1, 4, 32) + + class _AutoQuantMoeModel(torch.nn.Module): def __init__(self, num_experts_attr="num_experts"): super().__init__() @@ -297,6 +332,221 @@ def test_active_moe_search_prefers_budget_lower_bound(): assert best_recipes["layers.0.mlp.quant_recipe"]["format"] == "near_budget" +def test_auto_quantize_per_element_changes_selector_objective(): + candidate_stats = { + "small.quant_recipe": { + "formats": ["compressed", "bf16"], + "costs": [4.0, 8.0], + "element_costs": [4.0, 8.0], + "scores": [8.0, 0.0], + }, + "large.quant_recipe": { + "formats": ["compressed", "bf16"], + "costs": [4.0, 8.0], + "element_costs": [400.0, 800.0], + "scores": [80.0, 0.0], + }, + } + + def solve(score_model): + searcher = AutoQuantizeGradientSearcher() + searcher.config = {"cost_model": "weight"} + searcher.cost_model = "weight" + searcher.constraints = {"effective_bits": 6.0, "score_model": score_model} + searcher.candidate_stats = copy.deepcopy(candidate_stats) + recipes, is_satisfied = searcher.run_search_with_stats(12.0) + assert is_satisfied + return {name: info["format"] for name, info in recipes.items()} + + assert solve("raw") == { + "small.quant_recipe": "compressed", + "large.quant_recipe": "bf16", + } + assert solve("per_element") == { + "small.quant_recipe": "bf16", + "large.quant_recipe": "compressed", + } + + +def test_auto_quantize_per_element_tiny_scores_are_order_invariant(): + candidate_stats = { + "small.quant_recipe": { + "formats": ["compressed", "bf16"], + "costs": [4.0, 8.0], + "element_costs": [4.0, 8.0], + "scores": [8.0e-12, 0.0], + }, + "large.quant_recipe": { + "formats": ["compressed", "bf16"], + "costs": [4.0, 8.0], + "element_costs": [400.0, 800.0], + "scores": [80.0e-12, 0.0], + }, + } + + def solve(candidate_items): + searcher = AutoQuantizeGradientSearcher() + searcher.config = {"cost_model": "weight"} + searcher.cost_model = "weight" + searcher.constraints = {"effective_bits": 6.0, "score_model": "per_element"} + searcher.candidate_stats = copy.deepcopy(dict(candidate_items)) + recipes, is_satisfied = searcher.run_search_with_stats(12.0) + assert is_satisfied + return {name: info["format"] for name, info in recipes.items()} + + expected = { + "small.quant_recipe": "bf16", + "large.quant_recipe": "compressed", + } + items = list(candidate_stats.items()) + assert solve(items) == expected + assert solve(reversed(items)) == expected + + +def test_auto_quantize_groups_shared_expert_projections(): + searcher = AutoQuantizeGradientSearcher() + group_keys = [] + for projection in ("gate_proj", "up_proj", "down_proj"): + name = f"model.layers.0.mlp.shared_expert.{projection}" + group_key = next( + ( + result + for rule in searcher.quant_grouping_rules + if (result := searcher._apply_quant_group_rule(name, rule)) is not None + ), + name, + ) + group_keys.append(group_key) + + assert group_keys == ["model.layers.0.mlp.shared_expert"] * 3 + + +def test_auto_quantize_local_boundary_scores_shared_expert_at_parent_mlp(): + searcher = AutoQuantizeGradientSearcher() + searcher.config = {"score_boundary": "local"} + + score_modules = [] + for projection in ("gate_proj", "up_proj", "down_proj"): + name = f"model.layers.0.mlp.shared_expert.{projection}" + score_module = next( + ( + result + for rule in searcher._get_score_module_rules() + if (result := searcher._apply_score_group_rule(name, rule)) is not None + ), + name, + ) + score_modules.append(score_module) + + assert score_modules == ["model.layers.0.mlp"] * 3 + + +def test_auto_quantize_group_reconstruction_score_is_normalized_mse(): + reference = torch.tensor([[1.0, 2.0]]) + quantized = torch.tensor([[2.0, 4.0]]) + + score = quant_algorithms._get_group_reconstruction_score(reference, quantized) + + expected = (quantized - reference).square().mean() / reference.square().mean() + assert score.item() == pytest.approx(expected.item()) + with pytest.raises(ValueError, match="same shape"): + quant_algorithms._get_group_reconstruction_score(torch.ones(2), torch.ones(3)) + with pytest.raises(TypeError, match="expects a Tensor"): + quant_algorithms._get_group_reconstruction_score({"hidden": reference}, quantized) + + +def test_auto_quantize_group_score_boundary_does_not_group_recipe_decisions(): + model = _GroupBoundaryModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0, "score_model": "per_element"}, + quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + num_calib_steps=1, + num_score_steps=1, + method="group_recon", + ) + + self_qkv = model.self_attn.q_proj.get_hparam("quant_recipe") + self_o = model.self_attn.o_proj.get_hparam("quant_recipe") + linear_qkvz = model.linear_attn.in_proj_qkv.get_hparam("quant_recipe") + linear_ba = model.linear_attn.in_proj_a.get_hparam("quant_recipe") + linear_out = model.linear_attn.out_proj.get_hparam("quant_recipe") + + assert self_qkv is not self_o + assert linear_qkvz is not linear_ba + assert linear_qkvz is not linear_out + assert linear_ba is not linear_out + assert self_qkv.score_modules == [model.self_attn] + assert self_o.score_modules == [model.self_attn] + assert linear_qkvz.score_modules == [model.linear_attn] + assert linear_ba.score_modules == [model.linear_attn] + assert linear_out.score_modules == [model.linear_attn] + assert search_state["score_boundary"] == "group" + assert search_state["score_model"] == "per_element" + assert search_state["method"] == "group_recon" + assert all("element_costs" in stats for stats in search_state["candidate_stats"].values()) + assert any( + score > 0 for stats in search_state["candidate_stats"].values() for score in stats["scores"] + ) + assert model.config.use_cache is True + assert model.config.text_config.use_cache is True + assert (False, False) in model.use_cache_seen + + +def test_auto_quantize_gradient_group_score_boundary_supports_keyword_calls(): + model = _GroupBoundaryModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0, "score_model": "per_element"}, + quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + score_boundary="group", + ) + + assert search_state["score_boundary"] == "group" + assert search_state["method"] == "gradient" + assert any( + score > 0 for stats in search_state["candidate_stats"].values() for score in stats["scores"] + ) + + +def test_auto_quantize_attention_layer_grouping_groups_recipe_decisions(): + model = _GroupBoundaryModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0, "score_model": "per_element"}, + quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + num_calib_steps=1, + num_score_steps=1, + method="group_recon", + quant_grouping_scheme="runtime_fused+linear_attn_layer+self_attn_layer", + ) + + self_qkv = model.self_attn.q_proj.get_hparam("quant_recipe") + self_o = model.self_attn.o_proj.get_hparam("quant_recipe") + linear_qkvz = model.linear_attn.in_proj_qkv.get_hparam("quant_recipe") + linear_ba = model.linear_attn.in_proj_a.get_hparam("quant_recipe") + linear_out = model.linear_attn.out_proj.get_hparam("quant_recipe") + + assert self_qkv is self_o + assert linear_qkvz is linear_out + assert linear_ba is not linear_qkvz + assert self_qkv.score_modules == [model.self_attn] + assert linear_qkvz.score_modules == [model.linear_attn] + assert search_state["quant_grouping_scheme"] == ( + "runtime_fused+linear_attn_layer+self_attn_layer" + ) + + # use this config to test custom quantization config INT8_CUSTOM_QUANT_TEST_CFG = { "quant_cfg": [ @@ -550,6 +800,10 @@ def test_estimate_quant_compression(): nvfp4_kv_rotate_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_KV_ROTATE_CFG) assert estimate_quant_compression(nvfp4_kv_rotate_cfg) == 0.28125 + # Static-scale NVFP4 candidates carry the same block-scale storage cost. + nvfp4_mse_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG) + assert estimate_quant_compression(nvfp4_mse_cfg) == 0.28125 + nvfp4_svdquant_default_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_SVDQUANT_DEFAULT_CFG) assert estimate_quant_compression(nvfp4_svdquant_default_cfg) == 0.28125 @@ -674,7 +928,7 @@ def test_estimate_quant_compression_per_entry_effective_bits(): ) -@pytest.mark.parametrize("method", ["gradient", "kl_div"]) +@pytest.mark.parametrize("method", ["gradient", "group_recon", "kl_div"]) def test_auto_quantize_checkpoint_resume(method, tmp_path, capsys): """Test that checkpoint can be used to resume an interrupted search.""" model = SimpleLinear() @@ -764,7 +1018,71 @@ def test_auto_quantize_checkpoint_resume(method, tmp_path, capsys): ) -@pytest.mark.parametrize("method", ["gradient", "kl_div"]) +def test_auto_quantize_checkpoint_rejects_score_boundary_mismatch(tmp_path): + checkpoint_path = str(tmp_path / "autoquant_group_boundary.pth") + model = _GroupBoundaryModel() + common = { + "constraints": {"effective_bits": 8.0, "score_model": "per_element"}, + "quantization_formats": [ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_DEFAULT_CFG, + ], + "forward_step": lambda model, batch: model(batch), + "loss_func": lambda output, data: output.sum(), + "num_calib_steps": 1, + "num_score_steps": 1, + "checkpoint": checkpoint_path, + } + mtq.auto_quantize( + model, + data_loader=[model.get_input()], + score_boundary="group", + **common, + ) + + resumed_model = _GroupBoundaryModel() + with pytest.raises(ValueError, match="score boundary does not match"): + mtq.auto_quantize( + resumed_model, + data_loader=[resumed_model.get_input()], + score_boundary="local", + **common, + ) + + +def test_auto_quantize_checkpoint_rejects_quant_grouping_mismatch(tmp_path): + checkpoint_path = str(tmp_path / "autoquant_grouping.pth") + model = _GroupBoundaryModel() + common = { + "constraints": {"effective_bits": 8.0, "score_model": "per_element"}, + "quantization_formats": [ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_DEFAULT_CFG, + ], + "forward_step": lambda model, batch: model(batch), + "num_calib_steps": 1, + "num_score_steps": 1, + "method": "group_recon", + "checkpoint": checkpoint_path, + } + mtq.auto_quantize( + model, + data_loader=[model.get_input()], + quant_grouping_scheme="runtime_fused+linear_attn_layer+self_attn_layer", + **common, + ) + + resumed_model = _GroupBoundaryModel() + with pytest.raises(ValueError, match="quant grouping scheme does not match"): + mtq.auto_quantize( + resumed_model, + data_loader=[resumed_model.get_input()], + quant_grouping_scheme="runtime_fused", + **common, + ) + + +@pytest.mark.parametrize("method", ["gradient", "group_recon", "kl_div"]) def test_get_auto_quantize_config(method): model = TransformerBlock()