From aa0df41e6bb37a7f55da8b29cb7b8396c1f187ee Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 1 Sep 2026 00:05:30 +0000 Subject: [PATCH 1/3] feat(quantization): scoped calibration pipelines via `algo_cfg` Prototype of the flexible-calibration design: assign an ordered calibration pipeline per scope instead of one model-wide `algorithm`. config = { "quant_cfg": [...], "algo_cfg": [ {"module_name": "*self_attn*", "cfg": ["awq_lite", "mse"]}, {"module_name": "*mlp*", "cfg": ["max", {"method": "gptq"}]}, {"quantizer_name": "*input_quantizer", "cfg": ["max"]}, ], "algorithm": "max", # fallback for anything no entry matches } `compile_algo_cfg` lowers the config into ordered scoped stages, reading the model's structure to resolve globs and validate but mutating nothing. The new `calibration_plan` mode executes those stages through the existing calibration functions, gated by a `should_process` write-mask, and records one mode. `algorithm` lowers through the same path as its all-`"*"` case, so there is no second engine; with no `algo_cfg` the old path and its saved state are untouched. Two upstream bugs found while making pipelines actually sequence, both reproducible on today's un-scoped `algorithm=[...]` list: - `_mse_calibrate_weights` never restored the search calibrator it installs, so any stage after `mse` crashed (`algorithm=['max','mse','max']` -> TypeError). Now restored in a `finally`. - `awq_lite` called `enable_stats_collection(model)` directly, so the write-mask had to reach helper calls inside an algorithm, not just its module loop. Validation rejects a config before anything runs: unknown algorithm, empty scope, role mismatch, fusible siblings split across pipelines, a stage whose every write is overwritten before being read, and repeating an algorithm whose own output violates its precondition. The last two are what make `awq_lite -> mse -> awq_lite` wrong. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/quantization/algo_cfg.py | 713 ++++++++++++++++++ modelopt/torch/quantization/config.py | 123 +++ modelopt/torch/quantization/mode.py | 91 ++- modelopt/torch/quantization/model_calib.py | 143 +++- modelopt/torch/quantization/model_quant.py | 30 +- .../unit/torch/quantization/test_algo_cfg.py | 413 ++++++++++ 6 files changed, 1476 insertions(+), 37 deletions(-) create mode 100644 modelopt/torch/quantization/algo_cfg.py create mode 100644 tests/unit/torch/quantization/test_algo_cfg.py diff --git a/modelopt/torch/quantization/algo_cfg.py b/modelopt/torch/quantization/algo_cfg.py new file mode 100644 index 00000000000..e51fb669dab --- /dev/null +++ b/modelopt/torch/quantization/algo_cfg.py @@ -0,0 +1,713 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Compile a quantize config into an ordered list of scoped calibration stages. + +This is the *compile* half of the calibration plan (the ``calibration_plan`` mode in +:mod:`~modelopt.torch.quantization.mode` is the *execute* half). It is deliberately +side-effect free: it reads the already-quantized model's **structure** — quantizer and +linear names — to resolve globs and validate, but it mutates nothing, runs no forward and +touches no data. Consequences the design leans on: + +* bad configs fail fast, before any expensive calibration runs; +* it is testable without running a model; +* the resulting plan is a pure function of ``(config, model structure)``, so it is + identical on every rank — which is what keeps predicate scoping from desynchronizing + collectives in distributed calibration. + +Both surfaces lower here. ``algorithm="max"`` becomes the single all-``"*"`` stage, so +the legacy whole-model path is a special case of the scoped one rather than a second +engine. +""" + +import fnmatch +import warnings +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field + +import torch.nn as nn + +from .config import AlgoCfgEntry, QuantizeAlgorithmConfig, QuantizeConfig + +__all__ = [ + "AlgoCapabilities", + "AlgoCfgValidationError", + "AlgoStage", + "CalibrationPlan", + "compile_algo_cfg", + "describe_plan", + "plan_hash", + "stage_predicate", +] + + +class AlgoCfgValidationError(ValueError): + """Raised when an ``algo_cfg`` cannot be lowered into a valid plan.""" + + +# -------------------------------------------------------------------------------------- +# Capabilities +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class AlgoCapabilities: + """What one calibration algorithm consumes, produces and needs to run. + + Only the fields the compiler actually uses today are declared. ``produces`` / + ``requires`` are a small open vocabulary of state tokens: + + ``weight``, ``weight_amax``, ``input_amax``, ``pre_quant_scale``, ``acts``. + """ + + granularity: str # "tensor" | "module" + role: str # "weight" | "input" | "both" — which quantizers it may write + requires: frozenset[str] = frozenset() + produces: frozenset[str] = frozenset() + needs_forward: bool = True + self_forwards: bool = False # needs its own forward pass, cannot share one + # State tokens that must NOT already be present for this algorithm to be correct. + # ``awq_lite`` folds ``1/s`` into the weight and assumes it is starting from an + # unsmoothed weight; running it twice folds twice while keeping only the last + # activation-side scale (see ``apply_pre_quant_scale_and_smooth``). + requires_absent: frozenset[str] = frozenset() + + @property + def shareable_forward(self) -> bool: + """Whether this stage could ride a forward pass shared with other stages.""" + return self.needs_forward and not self.self_forwards + + +_W_AMAX = "weight_amax" +_I_AMAX = "input_amax" +_PQS = "pre_quant_scale" +_W = "weight" + +#: Declared capabilities per algorithm. Phase 1 of the design defers this contract, but +#: the two failure modes that make ``awq_lite -> mse -> awq_lite`` wrong are only +#: detectable with it, so a minimal table ships here. ``None`` (no calibration) is absent +#: on purpose — it compiles to an empty plan. +ALGO_CAPABILITIES: dict[str, AlgoCapabilities] = { + "max": AlgoCapabilities( + granularity="tensor", role="both", produces=frozenset({_W_AMAX, _I_AMAX}) + ), + "mse": AlgoCapabilities( + granularity="tensor", + role="weight", + requires=frozenset({_W, _W_AMAX}), + produces=frozenset({_W_AMAX}), + ), + "nvfp4_act_headroom": AlgoCapabilities( + granularity="tensor", + role="input", + requires=frozenset({"acts"}), + produces=frozenset({_I_AMAX}), + ), + "local_hessian": AlgoCapabilities( + granularity="module", + role="weight", + requires=frozenset({_W, _W_AMAX, "acts"}), + produces=frozenset({_W_AMAX}), + ), + "smoothquant": AlgoCapabilities( + granularity="module", + role="both", + requires=frozenset({"acts"}), + produces=frozenset({_PQS, _I_AMAX, _W, _W_AMAX}), + requires_absent=frozenset({_PQS}), + ), + "awq_lite": AlgoCapabilities( + granularity="module", + role="both", + requires=frozenset({"acts", _W}), + produces=frozenset({_PQS, _W_AMAX, _I_AMAX}), + self_forwards=True, + requires_absent=frozenset({_PQS}), + ), + "awq_clip": AlgoCapabilities( + granularity="module", + role="weight", + requires=frozenset({"acts", _W, _W_AMAX}), + produces=frozenset({_W_AMAX}), + self_forwards=True, + ), + "awq_full": AlgoCapabilities( + granularity="module", + role="both", + requires=frozenset({"acts", _W}), + produces=frozenset({_PQS, _W_AMAX, _I_AMAX}), + self_forwards=True, + requires_absent=frozenset({_PQS}), + ), + "gptq": AlgoCapabilities( + granularity="module", + role="weight", + requires=frozenset({_W, "acts"}), + produces=frozenset({_W, _W_AMAX}), + self_forwards=True, + ), + "svdquant": AlgoCapabilities( + granularity="module", + role="both", + requires=frozenset({"acts", _W}), + produces=frozenset({_PQS, _W, _W_AMAX, _I_AMAX}), + self_forwards=True, + requires_absent=frozenset({_PQS}), + ), + "lsq": AlgoCapabilities( + granularity="tensor", + role="weight", + requires=frozenset({_W, _W_AMAX}), + produces=frozenset({_W_AMAX}), + ), +} + + +# -------------------------------------------------------------------------------------- +# Stages +# -------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class AlgoStage: + """One algorithm applied to one scope — the unit of work in a calibration plan.""" + + algo: str | None + cfg: dict # kwargs for the algorithm, including its "method" key + scope: str # the glob + selector: str # "module_name" | "quantizer_name" + order: int # position within its entry's pipeline + entry: int # which algo_cfg entry it came from (-1 = the `algorithm` fallback) + # Scopes this stage must NOT touch, as ``(selector, glob)`` pairs. The model-wide + # ``algorithm`` fallback covers "everything an algo_cfg entry did not match", which is a + # complement and so cannot be written as a glob. Keeping it as globs-minus-globs (rather + # than a resolved name list) preserves the property the plan depends on: it is derived + # from the config alone, so every rank computes the same thing. + exclude: tuple[tuple[str, str], ...] = () + + @property + def capabilities(self) -> AlgoCapabilities | None: + """Declared capabilities of this stage's algorithm, or ``None`` if undeclared. + + Looked up rather than copied onto the stage: capabilities describe the *algorithm*, + so a stage that carried its own copy could drift from the registry. + """ + return ALGO_CAPABILITIES.get(self.algo) if self.algo else None + + def key(self) -> tuple: + """Execution-relevant identity, used for the plan hash. + + ``entry`` is deliberately excluded: it records *where in the config* a stage came + from, which is provenance, not behaviour. Leaving it out is what makes the legacy + ``algorithm="max"`` plan and the explicit ``[{"quantizer_name": "*", "cfg": ["max"]}]`` + plan hash identical -- the same execution, written two ways. + """ + return ( + self.algo, + tuple(sorted(self.cfg.items(), key=str)), + self.scope, + self.selector, + self.order, + tuple(sorted(self.exclude)), + ) + + def __str__(self) -> str: + extra = {k: v for k, v in self.cfg.items() if k != "method"} + extra_s = f" {extra}" if extra else "" + excl = f" minus {[g for _, g in self.exclude]}" if self.exclude else "" + return ( + f"{self.algo or 'none'} @ {self.selector}={self.scope!r}{excl} (#{self.order}){extra_s}" + ) + + +CalibrationPlan = list[AlgoStage] + + +# -------------------------------------------------------------------------------------- +# Target resolution +# -------------------------------------------------------------------------------------- + + +@dataclass +class _ModelIndex: + """Names of the things a scope can select, read once from the model structure.""" + + linears: list[str] = field(default_factory=list) + quantizers: list[str] = field(default_factory=list) + quantizers_of: dict[str, list[str]] = field(default_factory=dict) # linear -> quantizers + parent_of: dict[str, str] = field(default_factory=dict) # quantizer -> linear + + +def _index_model(model: nn.Module) -> _ModelIndex: + # Imported lazily: `mode` imports this module while `modelopt.torch.quantization` is + # still initializing, and `.nn` pulls in the quantized-tensor backends. + from .nn import SequentialQuantizer, TensorQuantizer + from .utils import is_quantized_linear + + index = _ModelIndex() + for name, module in model.named_modules(): + if is_quantized_linear(module): + index.linears.append(name) + index.quantizers_of[name] = [] + elif isinstance(module, (TensorQuantizer, SequentialQuantizer)): + index.quantizers.append(name) + for q in index.quantizers: + parent = q.rsplit(".", 1)[0] if "." in q else "" + if parent in index.quantizers_of: + index.quantizers_of[parent].append(q) + index.parent_of[q] = parent + return index + + +def resolve_targets(model: nn.Module, scope: str, selector: str) -> tuple[set[str], set[str]]: + """Resolve a scope into ``(module names, quantizer names)``. + + A ``module_name`` scope pulls in that module's quantizers; a ``quantizer_name`` scope + pulls in the owning modules, so a stage's write-mask covers whichever name its + algorithm happens to iterate over. + """ + index = _index_model(model) + if selector == "module_name": + modules = {n for n in index.linears if fnmatch.fnmatch(n, scope)} + quantizers = {q for m in modules for q in index.quantizers_of[m]} + else: + quantizers = {n for n in index.quantizers if fnmatch.fnmatch(n, scope)} + modules = {index.parent_of[q] for q in quantizers if q in index.parent_of} + return modules, quantizers + + +#: Which quantizer a state token lives on. ``weight`` (the tensor itself) rides with the +#: weight quantizer for scoping purposes. +TOKEN_ROLE: dict[str, str] = { + "weight": "weight", + "weight_amax": "weight", + "input_amax": "input", + "pre_quant_scale": "input", +} + + +def stage_targets(model: nn.Module, stage: AlgoStage) -> tuple[set[str], set[str]]: + """``(modules, quantizers)`` a stage may write, after subtracting its exclusions.""" + modules, quantizers = resolve_targets(model, stage.scope, stage.selector) + for selector, glob in stage.exclude: + ex_modules, ex_quantizers = resolve_targets(model, glob, selector) + modules -= ex_modules + quantizers -= ex_quantizers + return modules, quantizers + + +def role_quantizers(model: nn.Module, stage: AlgoStage) -> dict[str, set[str]]: + """The quantizers a stage may write, split by role. + + A stage's scope can pull in quantizers its algorithm will never touch — a + ``module_name`` scope resolves to both the weight and the input quantizer, but ``mse`` + only writes weights. Overlap between stages has to be judged on what they can actually + write, otherwise two stages that share a module but write different roles look like they + conflict when they do not. + """ + _, quantizers = stage_targets(model, stage) + caps = stage.capabilities + role = caps.role if caps else "both" + weight = {q for q in quantizers if "weight_quantizer" in q} + inp = quantizers - weight + return { + "weight": weight if role in ("weight", "both") else set(), + "input": inp if role in ("input", "both") else set(), + } + + +def effective_produces(model: nn.Module, stage: AlgoStage) -> set[str]: + """Tokens a stage actually writes here — declared ``produces`` minus roles it cannot reach.""" + caps = stage.capabilities + if caps is None: + return set() + by_role = role_quantizers(model, stage) + return {t for t in caps.produces if by_role[TOKEN_ROLE.get(t, "weight")]} + + +def effective_requires(model: nn.Module, stage: AlgoStage) -> set[str]: + """Non-ambient tokens a stage reads here.""" + caps = stage.capabilities + if caps is None: + return set() + by_role = role_quantizers(model, stage) + return {t for t in caps.requires - AMBIENT_TOKENS if by_role[TOKEN_ROLE.get(t, "weight")]} + + +def _token_overlap(model: nn.Module, a: AlgoStage, b: AlgoStage, token: str) -> bool: + """Whether two stages can write the same ``token`` on the same quantizers.""" + role = TOKEN_ROLE.get(token, "weight") + return bool(role_quantizers(model, a)[role] & role_quantizers(model, b)[role]) + + +def stage_predicate(model: nn.Module, stage: AlgoStage) -> Callable[[str], bool]: + """Build the ``should_process`` write-mask for a stage. + + The predicate is AND-ed into each algorithm's existing ``is_enabled`` filter, so a + stage **writes only its targets and never toggles enable-state** — reads (and hence + the activations seen by search-based algorithms like AWQ and GPTQ) are unchanged. + """ + modules, quantizers = stage_targets(model, stage) + allowed = modules | quantizers + return lambda name: name in allowed + + +# -------------------------------------------------------------------------------------- +# Lowering +# -------------------------------------------------------------------------------------- + + +def _algo_to_name_and_cfg(algo) -> tuple[str | None, dict]: + """Normalize one pipeline element to ``(algo_name, kwargs)``.""" + if isinstance(algo, QuantizeAlgorithmConfig): + algo = algo.model_dump() + if algo is None or isinstance(algo, str): + return algo, {"method": algo} + if isinstance(algo, dict): + if "method" not in algo: + raise AlgoCfgValidationError( + f"Algorithm dict must have a 'method' key; got {sorted(algo)}. Entry: {algo!r}" + ) + return algo["method"], dict(algo) + raise AlgoCfgValidationError(f"Invalid algorithm config type {type(algo)}: {algo!r}") + + +def _lower(entries: Iterable[AlgoCfgEntry], algorithm) -> CalibrationPlan: + """Config -> stages. No model needed; validation of names happens separately.""" + plan: CalibrationPlan = [] + for e_idx, entry in enumerate(entries): + selector, scope = entry.selector + for order, algo in enumerate(entry.cfg): + name, cfg = _algo_to_name_and_cfg(algo) + plan.append(AlgoStage(name, cfg, scope, selector, order, e_idx)) + + # The model-wide `algorithm` is the same thing at scope "*" -- one engine, not two. + # It is the *fallback*, so it must not re-run over targets an entry already claimed; + # otherwise the default would silently overwrite every scoped pipeline. + if algorithm is not None: + claimed = tuple(entry.selector for entry in entries) + algos = algorithm if isinstance(algorithm, list) else [algorithm] + for order, algo in enumerate(algos): + name, cfg = _algo_to_name_and_cfg(algo) + if name is None: + continue + plan.append(AlgoStage(name, cfg, "*", "quantizer_name", order, -1, exclude=claimed)) + return plan + + +# -------------------------------------------------------------------------------------- +# Validation +# -------------------------------------------------------------------------------------- + +#: Sibling linears that are fused into one kernel at export and therefore must share a +#: single weight scale — so they must also share one pipeline. +FUSED_SIBLING_GROUPS: tuple[tuple[str, ...], ...] = ( + ("q_proj", "k_proj", "v_proj"), + ("gate_proj", "up_proj"), + ("w1", "w3"), +) + + +def _report(msg: str, strict: bool = True, sink: list[str] | None = None) -> None: + """Record a validation violation. + + Violations are collected rather than raised on the first hit so one compile reports + everything wrong with a config -- a config with three mistakes should not take three + round trips to fix. ``strict=False`` turns the whole set into warnings. + """ + if sink is not None: + sink.append(msg) + return + if strict: + raise AlgoCfgValidationError(msg) + warnings.warn(f"algo_cfg: {msg}", stacklevel=3) + + +def known_algorithms() -> list[str]: + """Algorithm names currently registered in the calibrate-mode registry.""" + from .mode import CalibrateModeRegistry + + names = getattr(CalibrateModeRegistry, "_name2descriptor", {}) + return sorted( + n.removesuffix("_calibrate") + for n in names + if n.endswith("_calibrate") and not n.startswith("_") + ) + + +def _validate_config_only(plan: CalibrationPlan, strict: bool) -> None: + from .mode import BaseCalibrateModeDescriptor, CalibrateModeRegistry + + for stage in plan: + mode_name = BaseCalibrateModeDescriptor._get_mode_name(stage.algo) + if mode_name not in CalibrateModeRegistry: + _report( + f"unknown algorithm {stage.algo!r}. Known algorithms: {known_algorithms()}", + strict=True, + ) + + +def _validate_scopes(model: nn.Module, plan: CalibrationPlan, sink: list[str]) -> None: + for stage in plan: + modules, quantizers = resolve_targets(model, stage.scope, stage.selector) + if not modules and not quantizers: + _report( + f"scope {stage.selector}={stage.scope!r} (stage {stage}) matches no target in " + "the model. Check the glob against the quantized module/quantizer names.", + sink=sink, + ) + continue + + caps = stage.capabilities + if caps is None: + continue + # Role check: a weight-only algorithm pointed at input quantizers writes nothing. + if stage.selector == "quantizer_name": + roles = {"weight" if "weight_quantizer" in q else "input" for q in quantizers} + if caps.role != "both" and roles and caps.role not in roles: + _report( + f"{stage.algo!r} only writes {caps.role} quantizers but " + f"{stage.selector}={stage.scope!r} matches only {sorted(roles)} quantizers " + "— the stage would be a no-op.", + sink=sink, + ) + + +def _validate_fused_siblings(model: nn.Module, plan: CalibrationPlan, sink: list[str]) -> None: + """Fusible siblings must share one pipeline: one fused kernel, one weight scale.""" + pipeline_of: dict[str, tuple] = {} + for stage in plan: + if stage.entry < 0: + continue # the "*" fallback covers everything equally + modules, _ = resolve_targets(model, stage.scope, stage.selector) + for m in modules: + pipeline_of.setdefault(m, ()) + for stage in plan: + if stage.entry < 0: + continue + modules, _ = resolve_targets(model, stage.scope, stage.selector) + for m in modules: + pipeline_of[m] = (*pipeline_of[m], (stage.algo, stage.entry)) + + index = _index_model(model) + for group in FUSED_SIBLING_GROUPS: + by_parent: dict[str, dict[str, tuple]] = {} + for linear in index.linears: + leaf = linear.rsplit(".", 1)[-1] + if leaf in group and linear in pipeline_of: + by_parent.setdefault(linear.rsplit(".", 1)[0], {})[leaf] = pipeline_of[linear] + for parent, members in by_parent.items(): + distinct = {tuple(a for a, _ in v) for v in members.values()} + if len(distinct) > 1: + _report( + f"fusible siblings under {parent!r} got different pipelines " + f"({ {k: [a for a, _ in v] for k, v in members.items()} }). They export to one " + "fused kernel and must share a single weight scale, so they must share one " + "pipeline.", + sink=sink, + ) + + +def _validate_dependencies(model: nn.Module, plan: CalibrationPlan, sink: list[str]) -> None: + """Capability-derived checks: non-composable repeats and dead stages. + + Both are judged **per token and per role**: two stages conflict only when they write the + same state token on the same quantizers. A stage whose scope happens to include a module + another stage also touches is not a conflict if the two write different roles. + """ + for i, stage in enumerate(plan): + caps = stage.capabilities + if caps is None: + continue + + # (1) Non-composable repeat: an earlier stage already produced a token this + # algorithm needs to be *absent* to be correct. + for j in range(i): + prev = plan[j] + if prev.capabilities is None: + continue + clash = { + t + for t in caps.requires_absent & effective_produces(model, prev) + if _token_overlap(model, stage, prev, t) + } + if clash: + _report( + f"stage {i} ({stage}) cannot follow stage {j} ({prev}) on overlapping " + f"targets: {stage.algo!r} assumes {sorted(clash)} is not already set, but " + f"{prev.algo!r} produces it. Re-running it folds the scale a second time " + "while keeping only the last activation-side scale. Insert an explicit " + "unfold (disable_pre_quant_scale_and_resmooth) between them, or drop the " + "repeat.", + sink=sink, + ) + + # (2) Dead stage: every token it writes is overwritten downstream before anyone + # reads it, so the stage cannot affect the final model. + produced = effective_produces(model, stage) + if not produced: + continue + overwriters: dict[str, AlgoStage] = {} + for token in produced: + for j in range(i + 1, len(plan)): + later = plan[j] + if later.capabilities is None or not _token_overlap(model, stage, later, token): + continue + if token in effective_requires(model, later): + break # somebody read it -- not dead + if token in effective_produces(model, later): + overwriters[token] = later + break + if set(overwriters) == produced: + first = next(iter(overwriters.values())) + _report( + f"stage {i} ({stage}) is dead: everything it produces ({sorted(produced)}) is " + f"overwritten by a later stage ({first}) on the same quantizers, without being " + "read in between. Remove it, or move it after the stage that overwrites it.", + sink=sink, + ) + + +# -------------------------------------------------------------------------------------- +# Entry point +# -------------------------------------------------------------------------------------- + + +def compile_algo_cfg( + config: QuantizeConfig | dict, + model: nn.Module | None = None, + strict: bool = True, +) -> CalibrationPlan: + """Lower a quantize config into an ordered, validated list of scoped stages. + + Pure: reads model structure only, mutates nothing, runs no forward. + + Args: + config: a :class:`QuantizeConfig` or a mapping with ``algo_cfg`` / ``algorithm``. + model: the already-quantized model. Required for the model-aware validation + (scope resolution, roles, fused siblings, dependencies); when ``None`` only + the config-only rules run. + strict: raise :class:`AlgoCfgValidationError` on violations. ``False`` downgrades + them to warnings, which is what lets a knowingly-broken pipeline be run for + demonstration purposes. + + Returns: + The ordered plan. Stages run in list order. + """ + if isinstance(config, QuantizeConfig): + entries, algorithm = config.algo_cfg or [], config.algorithm + else: + raw_entries = config.get("algo_cfg") or [] + entries = [e if isinstance(e, AlgoCfgEntry) else AlgoCfgEntry(**e) for e in raw_entries] + algorithm = config.get("algorithm", "max") + + # An explicit algo_cfg suppresses the implicit whole-model default: entries are the + # plan, and `algorithm` only fills in what they do not cover. Keeping the "*" stage + # unconditionally would silently re-calibrate every scoped target. + if entries and algorithm is not None: + covered = _coverage_is_total(model, entries) if model is not None else False + if covered: + algorithm = None + + plan = _lower(entries, algorithm) + _validate_config_only(plan, strict) + if model is not None: + violations: list[str] = [] + _validate_scopes(model, plan, violations) + _validate_fused_siblings(model, plan, violations) + _validate_dependencies(model, plan, violations) + if violations: + body = "\n".join(f" {i + 1}. {v}" for i, v in enumerate(violations)) + msg = f"invalid algo_cfg ({len(violations)} problem(s)):\n{body}" + if strict: + raise AlgoCfgValidationError(msg) + warnings.warn(f"algo_cfg: {msg}", stacklevel=2) + return plan + + +def _coverage_is_total(model: nn.Module, entries: list[AlgoCfgEntry]) -> bool: + """Whether the entries already cover every quantizer, making `algorithm` redundant.""" + index = _index_model(model) + covered: set[str] = set() + for entry in entries: + selector, scope = entry.selector + _, quantizers = resolve_targets(model, scope, selector) + covered |= quantizers + return covered >= set(index.quantizers) + + +#: State tokens that are always available and therefore never need to be *produced* by a +#: stage: the weight tensor is part of the model, and activations come from the forward +#: loop rather than from another algorithm. +AMBIENT_TOKENS = frozenset({"weight", "acts"}) + + +def derive_handoff(model: nn.Module, plan: CalibrationPlan, i: int) -> dict: + """Extra kwargs for stage ``i`` implied by what earlier stages already produced. + + This is the general form of the hard-coded ``skip_max_init`` flag: when every non-ambient + token a stage requires was already produced by an earlier stage on the same quantizers, + the stage should refine that state rather than re-initialize it. Derived from the declared + capabilities, not from a table of algorithm pairs. + """ + stage = plan[i] + if stage.capabilities is None: + return {} + needed = effective_requires(model, stage) + if not needed: + return {} + + satisfied = { + token + for token in needed + for j in range(i) + if plan[j].capabilities is not None + and token in effective_produces(model, plan[j]) + and _token_overlap(model, stage, plan[j], token) + } + return {"skip_max_init": True} if needed <= satisfied else {} + + +def _stage_targets(model: nn.Module, stage: AlgoStage) -> set[str]: + modules, quantizers = stage_targets(model, stage) + return modules | quantizers + + +def plan_hash(plan: CalibrationPlan) -> str: + """A stable hash of the plan. + + Distributed calibration is safe only if every rank runs the *same* stages over the + *same* scopes — otherwise a predicate skips a quantizer on one rank and its amax + all-reduce never matches, which hangs. Comparing this hash across ranks turns that + silent deadlock into a clear error. + """ + import hashlib + + payload = "|".join(str(s.key()) for s in plan) + return hashlib.sha256(payload.encode()).hexdigest()[:16] + + +def describe_plan(plan: CalibrationPlan, model: nn.Module | None = None) -> str: + """Human-readable plan dump, used by the demos and for debugging.""" + if not plan: + return " (empty plan — no calibration)" + lines = [] + for i, stage in enumerate(plan): + suffix = "" + if model is not None: + modules, quantizers = stage_targets(model, stage) + suffix = f" -> {len(modules)} module(s), {len(quantizers)} quantizer(s)" + lines.append(f" [{i}] {stage}{suffix}") + return "\n".join(lines) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 3e3f20a5999..fee87dbc180 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -996,6 +996,16 @@ class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): description="If True, the amax will be synced across the distributed processes.", ) + skip_max_init: bool = ModeloptField( + default=False, + title="Skip the max-calibration that initializes amax before the MSE search.", + description="MSE normally runs ``max_calibrate`` first to seed ``amax``. When a previous " + "stage of an ``algo_cfg`` pipeline already produced weights and an initial ``amax`` " + "(e.g. ``gptq`` or ``awq_lite``), re-running max calibration would discard nothing but " + "does cost a forward; more importantly the search should refine *that* stage's amax. " + "The calibration-plan executor sets this automatically for non-leading MSE stages.", + ) + class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): """Configuration for local Hessian-weighted MSE calibration. @@ -1577,6 +1587,99 @@ def _dict_to_entry(key: str, value) -> list[dict[str, Any]]: return result +class AlgoCfgEntry(ModeloptBaseConfig): + """A single entry in an ``algo_cfg`` list — one scope, one ordered algorithm pipeline. + + Deliberately shaped like :class:`QuantizerCfgEntry`: a selector plus a ``cfg``. Where + ``quant_cfg`` entries carry quantizer *attributes*, ``algo_cfg`` entries carry the ordered + list of calibration *algorithms* to run on the matched targets. + + Exactly one selector must be given: + + - ``module_name`` — glob over quantized-linear module names. Use for weight/module-level + algorithms (``gptq``, ``awq_lite``, ``smoothquant``), where the role is implied by the + algorithm itself. + - ``quantizer_name`` — glob over quantizer module names. Use when the role must be picked + explicitly, e.g. ``max`` on ``*input_quantizer`` only. + """ + + module_name: str | None = ModeloptField( + default=None, + title="Module name pattern.", + description="Glob matched against quantized-linear module names.", + ) + quantizer_name: str | None = ModeloptField( + default=None, + title="Quantizer name pattern.", + description="Glob matched against quantizer module names.", + ) + cfg: list[_QuantizeAlgoCfgType] = ModeloptField( + default=..., + title="Ordered calibration pipeline for the matched targets.", + description="A list of algorithms run in order, each consuming the previous one's " + 'mutated weights/scales. An element is an algorithm name (``"max"``), a dict keyed on ' + '``method`` (``{"method": "gptq", "block_size": 64}``), or a ' + ":class:`QuantizeAlgorithmConfig`.", + ) + + @model_validator(mode="before") + @classmethod + def _normalize_entry(cls, values): + """Accept a bare (non-list) ``cfg`` and enforce the exactly-one-selector rule.""" + if not isinstance(values, dict): + return values + values = dict(values) + if "cfg" in values and not isinstance(values["cfg"], list): + values["cfg"] = [values["cfg"]] + selectors = [k for k in ("module_name", "quantizer_name") if values.get(k) is not None] + if len(selectors) != 1: + raise ValueError( + "AlgoCfgEntry needs exactly one of 'module_name' / 'quantizer_name'; got " + f"{selectors or 'neither'}. Entry: {values!r}" + ) + if not values.get("cfg"): + raise ValueError( + f"AlgoCfgEntry 'cfg' must list at least one algorithm. Got: {values!r}" + ) + return values + + @property + def selector(self) -> tuple[str, str]: + """``(selector_kind, glob)`` for this entry.""" + if self.module_name is not None: + return "module_name", self.module_name + return "quantizer_name", self.quantizer_name # type: ignore[return-value] + + +class CalibrationPlanConfig(QuantizeAlgorithmConfig): + """Config for the ``calibration_plan`` mode — the compiled, scoped calibration plan. + + The saved config is the user's *intent* (``algo_cfg`` + ``algorithm``), not the compiled + stage list: the plan is a pure function of the config and the model structure, so it is + re-derivable, and keeping the intent makes the recorded state readable. + """ + + method: Literal["calibration_plan"] = ModeloptField("calibration_plan") + + algo_cfg: list[AlgoCfgEntry] | None = ModeloptField( + default=None, + title="Scoped calibration pipelines; see :class:`AlgoCfgEntry`.", + ) + + algorithm: QuantizeAlgoCfgType = ModeloptField( + default=None, + title="Model-wide fallback algorithm for targets no ``algo_cfg`` entry matches.", + ) + + strict: bool = ModeloptField( + default=True, + title="Fail on validation errors instead of warning.", + description="``False`` downgrades plan-validation errors to warnings, so a config the " + "compiler considers wrong can still be executed (used to demonstrate *why* a rule " + "exists). Leave at ``True`` outside experiments.", + ) + + class QuantizeConfig(ModeloptBaseConfig): """Default configuration for ``quantize`` mode.""" @@ -1593,6 +1696,23 @@ class QuantizeConfig(ModeloptBaseConfig): validate_default=True, ) + algo_cfg: list[AlgoCfgEntry] | None = ModeloptField( + default=None, + title="Scoped calibration pipelines.", + description="An ordered list of :class:`AlgoCfgEntry` dicts assigning a calibration " + "pipeline to a scope, e.g. ``[{'module_name': '*mlp*', 'cfg': ['awq_lite', 'mse']}]``. " + "Targets not matched by any entry fall back to the model-wide ``algorithm``. When " + "omitted, ``algorithm`` alone is used and behaviour is unchanged.", + ) + + strict: bool = ModeloptField( + default=True, + title="Fail on ``algo_cfg`` validation errors instead of warning.", + description="Only affects configs that use ``algo_cfg``. ``False`` downgrades plan " + "validation errors to warnings so a pipeline the compiler considers wrong can still be " + "run -- useful for checking whether a rule is justified, not for production recipes.", + ) + effective_bits: float | None = ModeloptField( default=None, title="Effective bits per element (autoquant cost override)", @@ -1833,6 +1953,9 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]: def need_calibration(config: QuantizeConfig | Mapping[str, Any]) -> bool: """Check if calibration is needed for the given config.""" + if config.get("algo_cfg"): + # Any scoped pipeline is an explicit request to calibrate. + return True if config["algorithm"] is not None and config["algorithm"] != "max": return True diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index db7704e89b5..e5228c6c996 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -31,11 +31,13 @@ ) from modelopt.torch.opt.searcher import ForwardLoop +from .algo_cfg import compile_algo_cfg, derive_handoff, describe_plan, plan_hash, stage_predicate from .compress import compress_convert, compress_restore, update_compress_metadata from .config import ( AWQClipCalibConfig, AWQFullCalibConfig, AWQLiteCalibConfig, + CalibrationPlanConfig, CompressConfig, GPTQCalibConfig, LocalHessianCalibConfig, @@ -71,6 +73,7 @@ smoothquant, svdquant, ) +from .utils import print_rank_0 __all__ = ["BaseCalibrateModeDescriptor"] @@ -174,7 +177,7 @@ def name(self) -> str: def next_modes(self) -> set[str] | None: """Real quantization should be the last mode in the chain.""" # TODO: update this to support QLoRA - return {"max_calibrate", "eagle"} + return {"max_calibrate", "calibration_plan", "eagle"} @property def config_class(self) -> type[ModeloptBaseConfig]: @@ -218,6 +221,7 @@ def wrapped_calib_func( forward_loop: ForwardLoop | None = None, func: Callable | None = None, supports_layerwise: bool = True, + should_process: Callable[[str], bool] | None = None, ) -> ConvertReturnType: """Wrap the calibration function to be compatible with the ModelOpt convert entrypoint. @@ -238,6 +242,11 @@ def wrapped_calib_func( # For backward compatibility kwargs["algorithm"] = method + # The scoping write-mask (see `algo_cfg.stage_predicate`). `None` means "whole model", + # which is exactly today's behaviour, so it is not forwarded in that case. + if should_process is not None: + kwargs["should_process"] = should_process + moe_calib_experts_ratio = kwargs.pop("moe_calib_experts_ratio", None) if moe_calib_experts_ratio is not None: assert ( @@ -558,6 +567,86 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: _calib_func = gptq +def calibration_plan_convert( + model: ModelLikeModule, + config: CalibrationPlanConfig, + forward_loop: ForwardLoop | None = None, +) -> ConvertReturnType: + """Run a compiled calibration plan and record it as a single mode. + + Compile (pure) then execute (effectful): + + 1. :func:`compile_algo_cfg ` lowers + ``algo_cfg`` + ``algorithm`` into one ordered stage list and validates it against the + model structure. + 2. Each stage runs in order through the same :func:`wrapped_calib_func` the whole-model + algorithms use, with a ``should_process`` write-mask built from the stage's scope and + any handoff kwargs implied by what earlier stages produced. + + Stages are serial in this phase — each runs its own forward. Batching independent + stages onto a shared forward is a later optimization the declared capabilities already + carry enough information for (``shareable_forward``). + """ + plan = compile_algo_cfg( + {"algo_cfg": config.algo_cfg, "algorithm": config.algorithm}, + model, + strict=config.strict, + ) + print_rank_0( + f"calibration_plan: {len(plan)} stage(s), hash {plan_hash(plan)}\n" + + describe_plan(plan, model) + ) + + for i, stage in enumerate(plan): + if stage.algo is None: + continue + descriptor = CalibrateModeRegistry[ + BaseCalibrateModeDescriptor._get_mode_name(stage.algo, check=True) + ] + stage_kwargs = {**stage.cfg, **derive_handoff(model, plan, i)} + stage_config = descriptor.config_class(**stage_kwargs) + wrapped_calib_func( + model, + stage_config, + forward_loop, + func=type(descriptor)._calib_func, + supports_layerwise=type(descriptor)._supports_layerwise, + should_process=stage_predicate(model, stage), + ) + + metadata = {} + update_quantize_metadata(model, config, metadata) + return model, metadata + + +@CalibrateModeRegistry.register_mode +class CalibrationPlanModeDescriptor(BaseCalibrateModeDescriptor): + """Mode for a compiled, scoped calibration plan. + + One mode covers an arbitrary number of stages: recording one mode per stage would + bloat the saved state (``auto_quantize`` would emit hundreds). Restore needs nothing + algorithm-specific — the generic quantizer-state snapshot already captures amax, + pre_quant_scale, num_bits and friends. + """ + + _calib_func = None + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "calibration_plan" + + @property + def config_class(self) -> type[QuantizeAlgorithmConfig]: + """Specifies the config class for the mode.""" + return CalibrationPlanConfig + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return calibration_plan_convert + + @CalibrateModeRegistry.register_mode class LSQModeDescriptor(BaseCalibrateModeDescriptor): """Mode for LSQ (Learned Scale Quantization) algorithm.""" diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d4266b289b9..cdb4479e564 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -185,12 +185,22 @@ def _uses_modelopt_fp8_weight_scales(weight_quantizer: TensorQuantizer) -> bool: return weight_quantizer.backend is None and weight_quantizer.is_nvfp4_static -def weight_only_quantize(model: nn.Module): +def _in_scope(should_process: Callable[[str], bool] | None, name: str) -> bool: + """Write-mask helper: ``None`` means "the whole model", i.e. today's behaviour. + + Calibration algorithms AND this into their existing ``is_enabled`` filter so a scoped + stage writes only its own targets. It never toggles enable-state, so reads -- and hence + the activations seen by search-based algorithms -- are identical either way. + """ + return should_process is None or should_process(name) + + +def weight_only_quantize(model: nn.Module, should_process: Callable[[str], bool] | None = None): """Just quantize the weights of the model.""" name_to_module = dict(model.named_modules()) seen_modules = set() - for module in name_to_module.values(): - if module in seen_modules: + for name, module in name_to_module.items(): + if module in seen_modules or not _in_scope(should_process, name): continue if isinstance(module, QuantModule): @@ -316,6 +326,7 @@ def max_calibrate( sync_expert_weight_amax=False, shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, skip_forward_without_activation_calib: bool = False, + should_process: Callable[[str], bool] | None = None, ): """Calibrate the model using max. @@ -349,8 +360,8 @@ def max_calibrate( # Always run weight calibration on the weight tensor directly so every weight # quantizer gets ``_amax``, regardless of MoE routing. Downstream algorithms # (MSE, AWQ, export) then no longer need to patch in a missing ``_amax``. - enable_stats_collection(model) - weight_only_quantize(model) + enable_stats_collection(model, should_process) + weight_only_quantize(model, should_process) if forward_loop is not None: if skip_forward_without_activation_calib and not _needs_activation_forward_for_max_calib( model @@ -361,7 +372,7 @@ def max_calibrate( ) else: forward_loop(model) - finish_stats_collection(model) + finish_stats_collection(model, should_process=should_process) # Sync quantizer amax across local experts within each rank (for SequentialMLP) for name, module in model.named_modules(): @@ -383,7 +394,11 @@ def max_calibrate( # Check MoE calibration completeness before sync for name, module in model.named_modules(): - if isinstance(module, QuantModule) and _has_expert_parallelism(module): + if ( + isinstance(module, QuantModule) + and _has_expert_parallelism(module) + and _in_scope(should_process, name) + ): for child in module.children(): if isinstance(child, AnyQuantizer): _check_moe_calibration_complete(child, module.parallel_state) @@ -402,7 +417,7 @@ def sync_quantizer_amax_across_dp_ep(quantizer, parallel_state, parent_name, chi # Step 2:Sync amax across data parallelism for name, module in model.named_modules(): - if isinstance(module, QuantModule): + if isinstance(module, QuantModule) and _in_scope(should_process, name): for child_name, child in module.named_children(): if isinstance(child, AnyQuantizer): sync_quantizer_amax_across_dp_ep(child, module.parallel_state, name, child_name) @@ -740,6 +755,8 @@ def mse_calibrate( stop_multiplier: float = 4.0, fp8_scale_sweep: bool = False, shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, + skip_max_init: bool = False, + should_process: Callable[[str], bool] | None = None, ): """Calibrate weight quantizers using MSE-based amax search. @@ -765,7 +782,17 @@ def mse_calibrate( details on the remaining arguments. """ # max_calibrate initializes activations and weights; MSE only refines weights below. - max_calibrate(model, forward_loop, distributed_sync, shared_states=shared_states) + # When a previous pipeline stage already produced amax (and possibly mutated the + # weights), re-running it would discard that stage's starting point, so the executor + # sets skip_max_init and the search refines what the previous stage left behind. + if not skip_max_init: + max_calibrate( + model, + forward_loop, + distributed_sync, + shared_states=shared_states, + should_process=should_process, + ) name_to_module = dict(model.named_modules()) _mse_calibrate_weights( model, @@ -774,6 +801,7 @@ def mse_calibrate( start_multiplier=start_multiplier, stop_multiplier=stop_multiplier, fp8_scale_sweep=fp8_scale_sweep, + should_process=should_process, ) @@ -787,6 +815,7 @@ def _mse_calibrate_weights( fp8_scale_sweep: bool, error_func_for: Callable[[TensorQuantizer], Callable | None] | None = None, hessian_for: Callable[[TensorQuantizer], torch.Tensor | None] | None = None, + should_process: Callable[[str], bool] | None = None, ): """Run MSE weight calibration over all eligible quantizers (shared by mse / local-Hessian). @@ -797,9 +826,11 @@ def _mse_calibrate_weights( """ seen_modules: set[int] = set() pbar = tqdm(desc="MSE weight calibration") - for parent_module in name_to_module.values(): + for parent_name, parent_module in name_to_module.items(): if id(parent_module) in seen_modules or not isinstance(parent_module, QuantModule): continue + if not _in_scope(should_process, parent_name): + continue seen_modules.add(id(parent_module)) with enable_weight_access_and_writeback(parent_module, model, name_to_module): for weight, weight_quantizer in parent_module.iter_weights_for_calibration(): @@ -816,12 +847,22 @@ def _mse_calibrate_weights( ) if cal is None: continue - weight_quantizer._calibrator = cal - _run_and_load_max_stats( - weight_quantizer, partial(_collect_weight_stats, weight=weight) - ) - if hasattr(cal, "reset"): - cal.reset() + # The MSE calibrator is a *search* calibrator installed for the duration of + # this amax search only. Restoring the original afterwards matters as soon as + # calibration is a pipeline: a later stage that collects stats (any `max`, and + # anything built on it) would otherwise re-enter this spent calibrator and + # fail on its cleared `_initial_amax`. The structural fix is to make the + # per-tensor strategy config-selected instead of swapping the object. + previous_calibrator = weight_quantizer._calibrator + try: + weight_quantizer._calibrator = cal + _run_and_load_max_stats( + weight_quantizer, partial(_collect_weight_stats, weight=weight) + ) + if hasattr(cal, "reset"): + cal.reset() + finally: + weight_quantizer._calibrator = previous_calibrator pbar.update(1) pbar.close() @@ -1126,10 +1167,14 @@ def capture(weight_quantizer, weight, input_tensor): print_rank_0("local_hessian: Calibration complete.") -def enable_stats_collection(model: nn.Module): +def enable_stats_collection(model: nn.Module, should_process: Callable[[str], bool] | None = None): """Enable stats collection for all quantizers in the model.""" for name, module in model.named_modules(): - if isinstance(module, TensorQuantizer) and not module._disabled: + if ( + isinstance(module, TensorQuantizer) + and not module._disabled + and _in_scope(should_process, name) + ): if module._use_constant_amax or module._constant_amax is not None: # Quantizers with a constant amax use a fixed amax and don't need calibration. # Disable quantization during calibration so it doesn't affect other quantizers. @@ -1142,11 +1187,18 @@ def enable_stats_collection(model: nn.Module): module.disable() -def finish_stats_collection(model: nn.Module, method: str | None = None, **kwargs): +def finish_stats_collection( + model: nn.Module, + method: str | None = None, + should_process: Callable[[str], bool] | None = None, + **kwargs, +): """Finish stats collection for all quantizers in the model.""" - for _, module in model.named_modules(): + for _name, module in model.named_modules(): if not isinstance(module, TensorQuantizer) or module._disabled: continue + if not _in_scope(should_process, _name): + continue if module._use_constant_amax or module._constant_amax is not None: # Re-enable quantization for constant-amax quantizers disabled in enable_stats_collection. @@ -1272,7 +1324,12 @@ def apply_pre_quant_scale_and_smooth( @torch.no_grad() -def smoothquant(model: nn.Module, forward_loop: ForwardLoop | None = None, alpha=1.0): +def smoothquant( + model: nn.Module, + forward_loop: ForwardLoop | None = None, + alpha=1.0, + should_process: Callable[[str], bool] | None = None, +): """Smooth-Quant variant with per-channel weight scaling. Args: @@ -1302,10 +1359,11 @@ def smoothquant(model: nn.Module, forward_loop: ForwardLoop | None = None, alpha is_quantized_linear(module) and module.input_quantizer.is_enabled and module.input_quantizer.axis is None + and _in_scope(should_process, name) ): module.input_quantizer.axis = -1 - max_calibrate(model, forward_loop) + max_calibrate(model, forward_loop, should_process=should_process) def postprocess(module): # It is important to keep scaling math in fp32 to be numerically safe @@ -1364,6 +1422,7 @@ def awq( model: nn.Module, forward_loop: ForwardLoop | None = None, algorithm: str = "awq_lite", + should_process: Callable[[str], bool] | None = None, **kwargs, ): """Apply AWQ to the model. @@ -1378,16 +1437,20 @@ def awq( """ with SequentialQuantizer.convert_to_single_quantizer(model): if algorithm in ["awq_full", "awq_lite"]: - awq_lite(model, forward_loop, **kwargs) + awq_lite(model, forward_loop, should_process=should_process, **kwargs) if algorithm in ["awq_full", "awq_clip"]: - awq_clip(model, forward_loop, **kwargs) + awq_clip(model, forward_loop, should_process=should_process, **kwargs) # Special handling for SequentialQuantizer # Pre-compute name_to_module dict to avoid O(n^2) complexity in enable_weight_access_and_writeback name_to_module = dict(model.named_modules()) for name, module in model.named_modules(): - if is_quantized_linear(module) and isinstance(module.weight_quantizer, SequentialQuantizer): + if ( + is_quantized_linear(module) + and isinstance(module.weight_quantizer, SequentialQuantizer) + and _in_scope(should_process, name) + ): with enable_weight_access_and_writeback(module, model, name_to_module): max_calibrate(module, lambda linear: linear.weight_quantizer(module.weight)) @@ -1398,6 +1461,7 @@ def awq_lite( forward_loop: ForwardLoop, alpha_step: float = 0.1, debug: bool = False, + should_process: Callable[[str], bool] | None = None, **kwargs, ): """Lite version of AWQ. @@ -1563,7 +1627,11 @@ def forward(self, input, *args, **kwargs): # Pre-compute name_to_module dict ONCE to avoid O(n^2) complexity in enable_weight_access_and_writeback name_to_module = dict(model.named_modules()) for name, module in name_to_module.items(): - if is_quantized_linear(module) and module.weight_quantizer.is_enabled: + if ( + is_quantized_linear(module) + and module.weight_quantizer.is_enabled + and _in_scope(should_process, name) + ): with enable_weight_access_and_writeback(module, model, name_to_module): module.awq_lite = AWQLiteHelper(module, name) module.awq_lite.setup() @@ -1574,7 +1642,7 @@ def forward(self, input, *args, **kwargs): # Lets enable stats collection # This will collect amax for input_quantizers and KV quantizers during the caching mode forward pass - enable_stats_collection(model) + enable_stats_collection(model, should_process) forward_loop(model) # Load the amax values collected during the caching mode forward pass @@ -1583,8 +1651,10 @@ def forward(self, input, *args, **kwargs): model, [{"quantizer_name": "*weight_quantizer", "enable": False}], ): - max_calibrate(model, lambda model: None, distributed_sync=True) - finish_stats_collection(model) + max_calibrate( + model, lambda model: None, distributed_sync=True, should_process=should_process + ) + finish_stats_collection(model, should_process=should_process) def sync_act_scale_across_dp(module, data_parallel_group): """Sync activation scale across Data Parallel (DP).""" @@ -1731,6 +1801,7 @@ def awq_clip( min_clip_ratio: float = 0.5, shrink_step: float = 0.05, debug: bool = False, + should_process: Callable[[str], bool] | None = None, **kwargs, ): """AWQ-Clip variant. @@ -1901,6 +1972,7 @@ def forward(name, self, input, *args, **kwargs): is_quantized_linear(module) and module.weight_quantizer.is_enabled and module.weight_quantizer.block_sizes is not None + and _in_scope(should_process, name) ): bind_forward_method(module, partial(forward, name), "_forward_no_awq") with enable_weight_access_and_writeback(module, model, name_to_module): @@ -1909,7 +1981,7 @@ def forward(name, self, input, *args, **kwargs): print_rank_0("awq_clip: Estimating parameters...") # Lets enable stats collection # This will collect amax for input_quantizers and KV quantizers during the caching mode forward pass - enable_stats_collection(model) + enable_stats_collection(model, should_process) forward_loop(model) # Load the amax values collected during the caching mode forward pass # This will also perform distributed amax sync for input_quantizers @@ -1917,8 +1989,10 @@ def forward(name, self, input, *args, **kwargs): model, [{"quantizer_name": "*weight_quantizer", "enable": False}], ): - max_calibrate(model, lambda model: None, distributed_sync=True) - finish_stats_collection(model) + max_calibrate( + model, lambda model: None, distributed_sync=True, should_process=should_process + ) + finish_stats_collection(model, should_process=should_process) def postprocess(module): update_best_params(module) @@ -2224,6 +2298,7 @@ def gptq( perc_damp: float = 0.01, block_size: int = 128, fused: bool = False, + should_process: Callable[[str], bool] | None = None, ): """GPTQ quantization. @@ -2254,12 +2329,12 @@ def gptq( total_start = time.time() # TODO: Add support for other scale setting strateiges like weight-mse or local-hessian - max_calibrate(model, forward_loop=forward_loop) + max_calibrate(model, forward_loop=forward_loop, should_process=should_process) quantized_layers = [ (n, m) for n, m in model.named_modules() - if is_quantized_linear(m) and m.weight_quantizer.is_enabled + if is_quantized_linear(m) and m.weight_quantizer.is_enabled and _in_scope(should_process, n) ] if not quantized_layers: print_rank_0("No quantized linear layers found, skipping GPTQ") diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 3f6040fd4ef..095dc67c2c2 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -65,6 +65,8 @@ def calibrate( model: nn.Module, algorithm: QuantizeAlgoCfgType = "max", forward_loop: ForwardLoop | None = None, + algo_cfg: list | None = None, + strict: bool = True, ) -> nn.Module: """Adjusts weights and scaling factors based on selected algorithms. @@ -86,6 +88,14 @@ def calibrate( forward_loop: A callable which takes the model as argument and forwards calibration data through the model. This is not required for weight-only quantization with the ``"max"`` algorithm. + algo_cfg: An optional ordered list of :class:`AlgoCfgEntry + ` dicts assigning a calibration + pipeline to a scope, e.g. + ``[{"module_name": "*mlp*", "cfg": ["awq_lite", "mse"]}]``. When given, the config is + compiled into an ordered list of scoped stages run by the ``"calibration_plan"`` mode, + and ``algorithm`` becomes the fallback for targets no entry matches. + strict: Fail on plan-validation errors instead of warning. Only relevant with + ``algo_cfg``. Returns: The calibrated pytorch model. """ @@ -109,10 +119,20 @@ def forward_loop(model): is_training = model.training model.eval() + # A scoped plan runs through the `calibration_plan` mode, which compiles `algo_cfg` and + # `algorithm` into one ordered stage list and records a single mode. Without `algo_cfg` + # the plan is the all-"*" single-stage case, which is exactly what the legacy per-algorithm + # modes already do -- so that path is left alone and its recorded state stays byte-identical. + mode = ( + [("calibration_plan", {"algo_cfg": algo_cfg, "algorithm": algorithm, "strict": strict})] + if algo_cfg + else get_modelike_from_algo_cfg(algorithm) + ) + with forward_with_reshard(model): apply_mode( model, - mode=get_modelike_from_algo_cfg(algorithm), + mode=mode, mode_kwargs={"forward_loop": forward_loop}, ) @@ -247,7 +267,13 @@ def forward_loop(model) -> None: # Already quantized, so lets apply the quant_cfg from the config quant_cfg = QuantizeConfig(**dict(config)).quant_cfg set_quantizer_by_cfg(model, quant_cfg) - return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) + return calibrate( + model, + config.get("algorithm"), + forward_loop=forward_loop, + algo_cfg=config.get("algo_cfg"), + strict=config.get("strict", True), + ) # TODO: create a config interface for auto_quantize and expose setting diff --git a/tests/unit/torch/quantization/test_algo_cfg.py b/tests/unit/torch/quantization/test_algo_cfg.py new file mode 100644 index 00000000000..a34458774da --- /dev/null +++ b/tests/unit/torch/quantization/test_algo_cfg.py @@ -0,0 +1,413 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for ``algo_cfg`` lowering, validation and scoped execution.""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.algo_cfg import ( + AlgoCfgValidationError, + compile_algo_cfg, + derive_handoff, + plan_hash, + resolve_targets, + stage_predicate, + stage_targets, +) +from modelopt.torch.quantization.config import AlgoCfgEntry + +QUANT_CFG = [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "block_sizes": {-1: 32}}}, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, +] + + +class _MLP(nn.Module): + def __init__(self, d=32, h=64): + super().__init__() + self.gate_proj = nn.Linear(d, h, bias=False) + self.up_proj = nn.Linear(d, h, bias=False) + self.down_proj = nn.Linear(h, d, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _Attn(nn.Module): + def __init__(self, d=32): + super().__init__() + self.q_proj = nn.Linear(d, d, bias=False) + self.k_proj = nn.Linear(d, d, bias=False) + self.v_proj = nn.Linear(d, d, bias=False) + self.o_proj = nn.Linear(d, d, bias=False) + + def forward(self, x): + return self.o_proj(self.q_proj(x) + self.k_proj(x) + self.v_proj(x)) + + +class _Block(nn.Module): + def __init__(self, d=32): + super().__init__() + self.self_attn = _Attn(d) + self.mlp = _MLP(d) + + def forward(self, x): + return x + self.mlp(x + self.self_attn(x)) + + +class _Model(nn.Module): + def __init__(self, d=32, n=2): + super().__init__() + self.layers = nn.ModuleList([_Block(d) for _ in range(n)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +def _model(seed=0): + torch.manual_seed(seed) + return _Model().eval() + + +def _forward_loop(model): + torch.manual_seed(1) + for _ in range(2): + model(torch.randn(2, 8, 32)) + + +@pytest.fixture +def quantized(): + """Quantizers inserted, nothing calibrated: compile only needs the model structure.""" + return mtq.quantize(_model(), {"quant_cfg": QUANT_CFG, "algorithm": None}, None) + + +def _weight_amax(model): + return { + name: module._amax.detach().clone() + for name, module in model.named_modules() + if name.endswith("weight_quantizer") and getattr(module, "_amax", None) is not None + } + + +# ---------------------------------------------------------------------------- config + + +def test_entry_requires_exactly_one_selector(): + with pytest.raises(ValueError, match="exactly one of"): + AlgoCfgEntry(module_name="*mlp*", quantizer_name="*weight_quantizer", cfg=["max"]) + with pytest.raises(ValueError, match="exactly one of"): + AlgoCfgEntry(cfg=["max"]) + + +def test_entry_requires_nonempty_cfg(): + with pytest.raises(ValueError, match="at least one algorithm"): + AlgoCfgEntry(module_name="*mlp*", cfg=[]) + + +def test_entry_wraps_a_bare_cfg_in_a_list(): + assert AlgoCfgEntry(module_name="*mlp*", cfg="max").cfg == ["max"] + + +# ---------------------------------------------------------------------------- lowering + + +def test_algorithm_and_equivalent_algo_cfg_compile_to_the_same_plan(quantized): + """The legacy whole-model path is the all-``"*"`` case, not a second engine.""" + legacy = compile_algo_cfg({"algorithm": "max"}, quantized) + explicit = compile_algo_cfg( + {"algo_cfg": [{"quantizer_name": "*", "cfg": ["max"]}], "algorithm": None}, quantized + ) + assert plan_hash(legacy) == plan_hash(explicit) + + +def test_pipeline_lowers_in_order_with_kwargs(quantized): + plan = compile_algo_cfg( + { + "algo_cfg": [ + {"module_name": "*mlp*", "cfg": ["max", {"method": "mse", "step_size": 0.05}]} + ], + "algorithm": None, + }, + quantized, + ) + assert [stage.algo for stage in plan] == ["max", "mse"] + assert [stage.order for stage in plan] == [0, 1] + assert plan[1].cfg["step_size"] == 0.05 + + +def test_fallback_algorithm_excludes_scopes_claimed_by_entries(quantized): + """Otherwise the model-wide default silently re-runs over every scoped pipeline.""" + plan = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["max"]}], "algorithm": "max"}, quantized + ) + fallback = plan[-1] + assert fallback.exclude == (("module_name", "*mlp*"),) + + _, fallback_quantizers = stage_targets(quantized, fallback) + _, mlp_quantizers = resolve_targets(quantized, "*mlp*", "module_name") + assert mlp_quantizers + assert not (fallback_quantizers & mlp_quantizers) + + +def test_plan_hash_ignores_provenance_only_differences(quantized): + a = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["max"]}], "algorithm": None}, quantized + ) + b = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["max"]}], "algorithm": None}, quantized + ) + assert plan_hash(a) == plan_hash(b) + + +# ---------------------------------------------------------------------------- validation + + +def test_unknown_algorithm_is_rejected(quantized): + with pytest.raises(AlgoCfgValidationError, match="unknown algorithm"): + compile_algo_cfg( + {"algo_cfg": [{"module_name": "*", "cfg": ["awq_supreme"]}], "algorithm": None}, + quantized, + ) + + +def test_scope_matching_nothing_is_rejected(quantized): + with pytest.raises(AlgoCfgValidationError, match="matches no target"): + compile_algo_cfg( + {"algo_cfg": [{"module_name": "*cross_attn*", "cfg": ["max"]}], "algorithm": None}, + quantized, + ) + + +def test_weight_only_algorithm_on_input_quantizers_is_rejected(quantized): + with pytest.raises(AlgoCfgValidationError, match="only writes weight quantizers"): + compile_algo_cfg( + { + "algo_cfg": [{"quantizer_name": "*input_quantizer", "cfg": ["mse"]}], + "algorithm": None, + }, + quantized, + ) + + +def test_fusible_siblings_must_share_one_pipeline(quantized): + with pytest.raises(AlgoCfgValidationError, match="fusible siblings"): + compile_algo_cfg( + { + "algo_cfg": [ + {"module_name": "*gate_proj", "cfg": ["awq_lite"]}, + {"module_name": "*up_proj", "cfg": ["max"]}, + ], + "algorithm": None, + }, + quantized, + ) + + +def test_stage_whose_output_is_overwritten_before_being_read_is_rejected(quantized): + with pytest.raises(AlgoCfgValidationError, match="is dead"): + compile_algo_cfg( + { + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max", "mse", "max"]}], + "algorithm": None, + }, + quantized, + ) + + +def test_repeating_a_smoothing_algorithm_is_rejected(quantized): + """``awq_lite`` folds ``1/s`` into the weight; a second pass folds again without unfolding.""" + with pytest.raises(AlgoCfgValidationError, match="pre_quant_scale"): + compile_algo_cfg( + { + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["awq_lite", "awq_lite"]}], + "algorithm": None, + }, + quantized, + ) + + +def test_awq_then_mse_then_awq_reports_both_problems(quantized): + with pytest.raises(AlgoCfgValidationError) as excinfo: + compile_algo_cfg( + { + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["awq_lite", "mse", "awq_lite"]}], + "algorithm": None, + }, + quantized, + ) + message = str(excinfo.value) + assert "2 problem(s)" in message + assert "is dead" in message + assert "pre_quant_scale" in message + + +def test_awq_then_mse_is_accepted(quantized): + plan = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["awq_lite", "mse"]}], "algorithm": None}, + quantized, + ) + assert [stage.algo for stage in plan] == ["awq_lite", "mse"] + + +def test_stages_sharing_a_module_but_writing_different_roles_do_not_conflict(quantized): + """A ``module_name`` scope resolves to both roles; overlap is judged on what is written.""" + plan = compile_algo_cfg( + { + "algo_cfg": [ + {"module_name": "*mlp*", "cfg": ["max", "mse"]}, + {"quantizer_name": "*input_quantizer", "cfg": ["max"]}, + ], + "algorithm": None, + }, + quantized, + ) + assert len(plan) == 3 + + +def test_strict_false_downgrades_violations_to_warnings(quantized): + config = { + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max", "mse", "max"]}], + "algorithm": None, + } + with pytest.warns(UserWarning, match="is dead"): + plan = compile_algo_cfg(config, quantized, strict=False) + assert len(plan) == 3 + + +# ---------------------------------------------------------------------------- handoff + + +def test_mse_after_a_stage_that_produced_amax_skips_its_own_max_init(quantized): + plan = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["max", "mse"]}], "algorithm": None}, + quantized, + ) + assert derive_handoff(quantized, plan, 0) == {} + assert derive_handoff(quantized, plan, 1) == {"skip_max_init": True} + + +def test_leading_mse_still_initializes_its_own_amax(quantized): + plan = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["mse"]}], "algorithm": None}, quantized + ) + assert derive_handoff(quantized, plan, 0) == {} + + +# ---------------------------------------------------------------------------- scoping + + +def test_stage_predicate_matches_only_its_own_targets(quantized): + plan = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["max"]}], "algorithm": None}, quantized + ) + should_process = stage_predicate(quantized, plan[0]) + assert should_process("layers.0.mlp.gate_proj") + assert should_process("layers.0.mlp.gate_proj.weight_quantizer") + assert not should_process("layers.0.self_attn.q_proj") + assert not should_process("layers.0.self_attn.q_proj.weight_quantizer") + + +def test_scoped_stage_writes_only_its_targets(): + model = mtq.quantize( + _model(), + { + "quant_cfg": QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max"]}], + }, + _forward_loop, + ) + calibrated = _weight_amax(model) + assert calibrated + assert all("mlp" in name for name in calibrated) + + +def test_scoping_never_toggles_enable_state(): + from modelopt.torch.quantization.nn import TensorQuantizer + + def flags(model): + return { + name: bool(module.is_enabled) + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) + } + + before = flags(mtq.quantize(_model(), {"quant_cfg": QUANT_CFG, "algorithm": None}, None)) + after = flags( + mtq.quantize( + _model(), + { + "quant_cfg": QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max"]}], + }, + _forward_loop, + ) + ) + assert before == after + + +def test_scoped_plan_records_a_single_calibration_mode(): + from modelopt.torch.opt.conversion import ModeloptStateManager + + model = mtq.quantize( + _model(), + { + "quant_cfg": QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max", "mse"]}], + }, + _forward_loop, + ) + modes = [str(mode) for mode, _, _ in ModeloptStateManager(model).modes_with_states()] + assert modes == ["quantize", "calibration_plan"] + + +def test_a_stage_can_follow_mse(quantized): + """``mse`` installs a search calibrator; it must not outlive its own stage.""" + model = mtq.quantize( + _model(), + { + "quant_cfg": QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max", "mse", "max"]}], + "strict": False, + }, + _forward_loop, + ) + assert _weight_amax(model) + + +def test_legacy_path_is_numerically_unchanged(): + legacy = mtq.quantize(_model(), {"quant_cfg": QUANT_CFG, "algorithm": "max"}, _forward_loop) + planned = mtq.quantize( + _model(), + { + "quant_cfg": QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"quantizer_name": "*", "cfg": ["max"]}], + }, + _forward_loop, + ) + legacy_amax, planned_amax = _weight_amax(legacy), _weight_amax(planned) + assert set(legacy_amax) == set(planned_amax) + assert all(torch.equal(legacy_amax[k], planned_amax[k]) for k in legacy_amax) From 233bbfd44e30329cdc4a13dd07929f2a35962294 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Tue, 1 Sep 2026 00:18:26 +0000 Subject: [PATCH 2/3] fix(quantization): only pass a derived handoff to algorithms that declare it `derive_handoff` reports what state earlier stages already produced; most algorithms have no knob to act on that. Handing `skip_max_init` to one of them made the stage config raise `extra_forbidden`, so every chain ending in `awq_clip` (which also consumes a prior stage's amax) failed to build. Found by sweeping all 121 ordered algorithm pairs through compile-then-run and cross-checking each outcome against the declared capability table: the five `* -> awq_clip` chains crashed on the scoped path while working on the legacy `algorithm=[...]` path, which located the bug in the executor rather than in any algorithm. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/quantization/mode.py | 11 +++++++++-- tests/unit/torch/quantization/test_algo_cfg.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index e5228c6c996..8fca3e7a407 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -603,8 +603,15 @@ def calibration_plan_convert( descriptor = CalibrateModeRegistry[ BaseCalibrateModeDescriptor._get_mode_name(stage.algo, check=True) ] - stage_kwargs = {**stage.cfg, **derive_handoff(model, plan, i)} - stage_config = descriptor.config_class(**stage_kwargs) + # A derived handoff only takes effect if the algorithm exposes a knob for it: + # `derive_handoff` reports what state earlier stages already produced, but most + # algorithms have no way to act on that and would reject the extra kwarg. + handoff = { + key: value + for key, value in derive_handoff(model, plan, i).items() + if key in descriptor.config_class.model_fields + } + stage_config = descriptor.config_class(**{**stage.cfg, **handoff}) wrapped_calib_func( model, stage_config, diff --git a/tests/unit/torch/quantization/test_algo_cfg.py b/tests/unit/torch/quantization/test_algo_cfg.py index a34458774da..88c4f70c39d 100644 --- a/tests/unit/torch/quantization/test_algo_cfg.py +++ b/tests/unit/torch/quantization/test_algo_cfg.py @@ -305,6 +305,20 @@ def test_mse_after_a_stage_that_produced_amax_skips_its_own_max_init(quantized): assert derive_handoff(quantized, plan, 1) == {"skip_max_init": True} +def test_handoff_is_dropped_for_algorithms_without_the_matching_knob(): + """`awq_clip` also consumes a prior stage's amax but has no `skip_max_init` field.""" + model = mtq.quantize( + _model(), + { + "quant_cfg": QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["max", "awq_clip"]}], + }, + _forward_loop, + ) + assert _weight_amax(model) + + def test_leading_mse_still_initializes_its_own_amax(quantized): plan = compile_algo_cfg( {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["mse"]}], "algorithm": None}, quantized From 2ad1796b230df39e3dcb032602810656f6c262f4 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:16:58 +0000 Subject: [PATCH 3/3] feat(quantization): let GPTQ consume a preceding range search `gptq()` unconditionally re-derived amax from max before its weight update, so a range search in front of it was discarded and the compiler correctly reported the search as a dead stage. But rounding error is only compensated consistently if GPTQ works against the grid the model actually uses, so the search belongs *before* GPTQ, not after -- which is the order DeepCompressor's QoQ recipes ship (`qoq-gchn.yaml`, `ooo.yaml`: `enable_calib_range` then `enable_kernel_gptq`). - `gptq(..., skip_max_init=False)` guards the initial `max_calibrate`. It also seeds the input quantizers, so it may only be skipped when an earlier stage calibrated them -- which is what the executor's handoff guarantees. - `GPTQCalibConfig.skip_max_init` exposes it. - `ALGO_CAPABILITIES["gptq"]` declares `weight_amax` as an input, so the executor derives the flag and the dead-stage rule stops firing. When nothing produced an amax, GPTQ still initializes its own, so an unsatisfied input is not an error. Tests cover both chains this enables: `mse -> gptq` keeps the searched amax bit-identically while differing from plain GPTQ, and `awq_lite -> mse` refines the amax on AWQ's smoothed weights. One declared token moved 5 of 121 ordered pairs and 105 of 1331 triples from "rejected" to "composing" in the algorithm sweep. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/quantization/algo_cfg.py | 6 +- modelopt/torch/quantization/config.py | 11 ++++ modelopt/torch/quantization/model_calib.py | 15 ++++- .../unit/torch/quantization/test_algo_cfg.py | 62 +++++++++++++++++++ 4 files changed, 90 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/algo_cfg.py b/modelopt/torch/quantization/algo_cfg.py index e51fb669dab..0c43a7725b0 100644 --- a/modelopt/torch/quantization/algo_cfg.py +++ b/modelopt/torch/quantization/algo_cfg.py @@ -154,7 +154,11 @@ def shareable_forward(self) -> bool: "gptq": AlgoCapabilities( granularity="module", role="weight", - requires=frozenset({_W, "acts"}), + # GPTQ rounds against an existing quantization grid; it declares `weight_amax` as an + # input so a preceding range search (`mse`, `local_hessian`) is recognized as feeding + # it rather than as dead work. When nothing produced one, GPTQ initializes it itself + # (`skip_max_init=False`), which is why an unsatisfied `weight_amax` is not an error. + requires=frozenset({_W, "acts", _W_AMAX}), produces=frozenset({_W, _W_AMAX}), self_forwards=True, ), diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index fee87dbc180..0f6ac1bd7bb 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1251,6 +1251,17 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig): per-column error propagation into one launch per GPTQ block.""", ) + skip_max_init: bool = ModeloptField( + default=False, + title="Skip the max-calibration that initializes amax before the GPTQ update.", + description="GPTQ normally runs ``max_calibrate`` first so every quantizer has an amax to " + "round against. When an earlier stage of an ``algo_cfg`` pipeline already produced that " + "amax -- e.g. an ``mse`` range search -- re-deriving it from max would discard the search " + "and make GPTQ compensate against a grid the model will not use. The calibration-plan " + "executor sets this automatically for non-leading GPTQ stages; it must not be set when " + "GPTQ runs first, since nothing else would initialize amax.", + ) + @model_validator(mode="after") def _gptq_qdq_default(self): """Inject ``get_qdq_activations_from_prev_layer=True`` unless the user set it. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index cdb4479e564..8ad7ce85488 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2298,6 +2298,7 @@ def gptq( perc_damp: float = 0.01, block_size: int = 128, fused: bool = False, + skip_max_init: bool = False, should_process: Callable[[str], bool] | None = None, ): """GPTQ quantization. @@ -2312,7 +2313,8 @@ def gptq( Per-module steps: - 1. ``max_calibrate`` to set amax values from the current activations. + 1. ``max_calibrate`` to set amax values from the current activations, unless + ``skip_max_init`` says an earlier pipeline stage already produced them. 2. Promote eligible quantizers to ``StaticBlockScaleQuantizer`` (two-level scaling). 3. Collect per-linear-layer Hessian matrices via forward hooks. 4. Blockwise weight updates using the inverse Hessian to compensate for @@ -2325,11 +2327,18 @@ def gptq( perc_damp: Percentage of avg Hessian diagonal for damping (default: 0.01). block_size: Block size for GPTQ weight update. fused: If True, use fused Triton kernel for NVFP4 static quantization. + skip_max_init: If True, keep the amax an earlier stage established instead of + re-deriving it from max. GPTQ compensates rounding error against a specific + quantization grid, so when a previous stage searched a better grid (e.g. ``mse``) + the compensation must be computed against *that* grid, not a fresh max one. """ total_start = time.time() - # TODO: Add support for other scale setting strateiges like weight-mse or local-hessian - max_calibrate(model, forward_loop=forward_loop, should_process=should_process) + # Scale setting: max by default, or whatever a previous pipeline stage established. + # Note this also seeds the input quantizers, so it may only be skipped when an earlier + # stage has already calibrated them -- which is what the executor's handoff guarantees. + if not skip_max_init: + max_calibrate(model, forward_loop=forward_loop, should_process=should_process) quantized_layers = [ (n, m) diff --git a/tests/unit/torch/quantization/test_algo_cfg.py b/tests/unit/torch/quantization/test_algo_cfg.py index 88c4f70c39d..6d9aab85fb4 100644 --- a/tests/unit/torch/quantization/test_algo_cfg.py +++ b/tests/unit/torch/quantization/test_algo_cfg.py @@ -108,6 +108,20 @@ def _weight_amax(model): } +def _run_chain(cfg, quant_cfg=None): + """Calibrate a fresh model with one scoped pipeline; returns weight amax by name.""" + model = mtq.quantize( + _model(), + { + "quant_cfg": quant_cfg or QUANT_CFG, + "algorithm": None, + "algo_cfg": [{"module_name": "*mlp*", "cfg": cfg}], + }, + _forward_loop, + ) + return _weight_amax(model) + + # ---------------------------------------------------------------------------- config @@ -319,6 +333,54 @@ def test_handoff_is_dropped_for_algorithms_without_the_matching_knob(): assert _weight_amax(model) +def test_range_search_then_gptq_is_recognized_as_a_handoff(quantized): + """GPTQ rounds against an existing grid, so a preceding search feeds it, not dies.""" + plan = compile_algo_cfg( + { + "algo_cfg": [{"module_name": "*mlp*", "cfg": ["mse", {"method": "gptq"}]}], + "algorithm": None, + }, + quantized, + ) + assert [stage.algo for stage in plan] == ["mse", "gptq"] + assert derive_handoff(quantized, plan, 0) == {} + assert derive_handoff(quantized, plan, 1) == {"skip_max_init": True} + + +def test_leading_gptq_still_initializes_its_own_amax(quantized): + plan = compile_algo_cfg( + {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["gptq"]}], "algorithm": None}, quantized + ) + assert derive_handoff(quantized, plan, 0) == {} + + +def test_gptq_preserves_a_preceding_range_search(): + """The point of the chain: GPTQ compensates against the grid MSE searched. + + Prior art for this ordering is DeepCompressor's QoQ recipes, which run their weight + range search before the GPTQ kernel rather than after. + """ + gptq = {"method": "gptq", "block_size": 32} + only_mse = _run_chain(["mse"]) + only_gptq = _run_chain([gptq]) + chained = _run_chain(["mse", gptq]) + + probe = "layers.0.mlp.gate_proj.weight_quantizer" + # GPTQ kept MSE's amax instead of re-deriving it from max ... + assert torch.equal(chained[probe], only_mse[probe]) + # ... and the resulting model is not the one plain GPTQ produces. + assert not torch.equal(chained[probe], only_gptq[probe]) + + +def test_awq_then_mse_refines_the_smoothed_weights(): + """MSE re-searches the amax on AWQ's smoothed weights: a forward-free awq_clip.""" + only_awq = _run_chain(["awq_lite"]) + chained = _run_chain(["awq_lite", "mse"]) + + assert set(only_awq) == set(chained) + assert any(not torch.equal(only_awq[k], chained[k]) for k in only_awq) + + def test_leading_mse_still_initializes_its_own_amax(quantized): plan = compile_algo_cfg( {"algo_cfg": [{"module_name": "*mlp*", "cfg": ["mse"]}], "algorithm": None}, quantized