From 08497cacb316d68ea38cbb2df4c28a049b043806 Mon Sep 17 00:00:00 2001
From: Copilot CLI <223556219+Copilot@users.noreply.github.com>
Date: Fri, 15 May 2026 22:01:47 +0000
Subject: [PATCH 01/18] WIP: MoE quantization (intermediate,
QuantLinear/Embedding still present)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
This commit checkpoints the in-progress MoE quantization work before a
larger refactor that deletes QuantLinear/QuantEmbedding in favour of
storing every quantized weight (2D linear, 2D embedding, 3D MoE experts)
as a QuantTensor nn.Parameter on the original host module.
Included so far:
- New olive/common/quant/patterns.py for re: prefix matching in
modules_to_not_convert / overrides.
- New olive/common/quant/tensor.py with QuantTensor wrapper subclass
(_make_wrapper_subclass + __torch_function__ + __torch_dispatch__),
supporting 2D and 3D layouts.
- LayerWrapper.get_experts() / get_router() accessors.
- 3D quantize helpers in olive/common/quant/utils.py.
- moe field on OliveHfQuantizationConfig.
- _process_model_before_weight_loading skips ModuleList(Expert) subtrees
when moe=False, fixing a latent silent-quantization bug for
Mixtral / PhiMoE / Qwen2/3-MoE.
- Fused-3D MoE support in prepare_model / finalize via QuantTensor
parameters; current save layout uses _qweight buffer suffixes — to be
replaced in the upcoming refactor with the canonical
.qweight/.scales/.qzeros layout.
- ModelBuilder raises NotImplementedError for Olive-quantized MoE
checkpoints (Mobius is the intended consumer).
- Test additions:
test/common/quant/test_patterns.py, test/common/quant/test_tensor.py,
TestOliveHfQuantizerMoE / TestRegexOverrides in test_hf_utils.py,
test/passes/pytorch/test_quant_utils.py for flatten helper,
test_olive_quantized_model_raises_for_moe in test_model_builder.py.
- 294 tests pass; lintrunner clean (--skip PYLINT).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/wrapper.py | 61 +++
olive/common/quant/hf_utils.py | 62 ++-
olive/common/quant/patterns.py | 80 ++++
olive/common/quant/tensor.py | 490 ++++++++++++++++++++++++
olive/common/quant/utils.py | 75 +++-
olive/passes/onnx/model_builder.py | 8 +
olive/passes/pytorch/quant_utils.py | 158 +++++++-
olive/passes/pytorch/rtn.py | 2 +-
test/common/quant/test_hf_utils.py | 149 +++++++
test/common/quant/test_patterns.py | 83 ++++
test/common/quant/test_tensor.py | 138 +++++++
test/passes/onnx/test_model_builder.py | 30 ++
test/passes/pytorch/test_quant_utils.py | 87 +++++
13 files changed, 1408 insertions(+), 15 deletions(-)
create mode 100644 olive/common/quant/patterns.py
create mode 100644 olive/common/quant/tensor.py
create mode 100644 test/common/quant/test_patterns.py
create mode 100644 test/common/quant/test_tensor.py
create mode 100644 test/passes/pytorch/test_quant_utils.py
diff --git a/olive/common/hf/wrapper.py b/olive/common/hf/wrapper.py
index 8bde44ef58..f8c7fadb23 100644
--- a/olive/common/hf/wrapper.py
+++ b/olive/common/hf/wrapper.py
@@ -138,6 +138,29 @@ class LayerWrapper:
"opt": ["fc2"],
"qwen": ["c_proj"],
}
+ # MoE-block conventions. These are resolved relative to ``self.mlp``
+ # (i.e., the layer's MLP attribute) because every modern HF MoE
+ # transformer block lives at ``layer.mlp``.
+ #
+ # ``EXPERTS`` is the experts sub-module:
+ # - For fused-3D MoEs (Mixtral, Qwen3-MoE, GPT-OSS) it owns 3D
+ # ``nn.Parameter`` tensors such as ``gate_up_proj`` of shape
+ # ``(num_experts, ...)``.
+ # - For ``ModuleList(Expert)`` MoEs (PhiMoE, DeepSeek-V3, classic
+ # Mixtral) it is the ``ModuleList`` whose children are the
+ # per-expert ``nn.Module`` blocks (each with their own
+ # ``nn.Linear``s).
+ #
+ # ``ROUTER`` is the routing module ("gate" in most, "router" in
+ # GPT-OSS). It is usually an ``nn.Linear`` (or a small custom module
+ # containing one) and should typically be kept in full precision.
+ EXPERTS = {
+ "default": "experts",
+ }
+ ROUTER = {
+ "default": "gate",
+ "gpt_oss": "router",
+ }
def __init__(self, layer: nn.Module, model_type: str):
# TODO(jambayk): use _layer and property to get the layer?
@@ -189,6 +212,44 @@ def get_mlp_outputs(self, return_name: bool = True):
self.mlp, self.MLP_OUTPUTS, self.model_type, return_name=return_name, return_name_prefix=f"{self.mlp_name}."
)
+ def get_experts(self, return_name: bool = True):
+ """Return the experts sub-module of this layer (or ``None`` if not MoE).
+
+ The experts sub-module is the parent of every per-expert weight
+ (fused 3D ``nn.Parameter``s, or a ``ModuleList`` of per-expert
+ ``nn.Module``s). The caller can use the returned module to (a)
+ collect ids of every ``nn.Module`` under the experts subtree
+ (for the ``moe=False`` skip set) and (b) iterate the
+ ``nn.Parameter``s to quantize when ``moe=True``.
+
+ Returns ``None`` (and an empty name when ``return_name=True``)
+ for layers without an experts sub-module.
+ """
+ if self.mlp is None:
+ return (None, "") if return_name else None
+ module = get_submodules(self.mlp, self.EXPERTS, self.model_type, return_name=False, fail_on_not_found=False)
+ if module is None:
+ return (None, "") if return_name else None
+ name = f"{self.mlp_name}.{self.EXPERTS.get(self.model_type, self.EXPERTS['default'])}"
+ return (module, name) if return_name else module
+
+ def get_router(self, return_name: bool = True):
+ """Return the router sub-module of this layer (or ``None`` if not MoE).
+
+ Routers are typically small modules (e.g., a single
+ ``nn.Linear``) and should usually be kept in full precision.
+ Olive does not quantize routers automatically — this accessor is
+ used by callers that want to skip the router via
+ ``modules_to_not_convert``.
+ """
+ if self.mlp is None:
+ return (None, "") if return_name else None
+ module = get_submodules(self.mlp, self.ROUTER, self.model_type, return_name=False, fail_on_not_found=False)
+ if module is None:
+ return (None, "") if return_name else None
+ name = f"{self.mlp_name}.{self.ROUTER.get(self.model_type, self.ROUTER['default'])}"
+ return (module, name) if return_name else module
+
class ModelWrapper:
"""Wrapper for transformer model."""
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index c1e4c5065b..67d82172f7 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -11,6 +11,8 @@
from transformers.quantizers.base import HfQuantizer
from transformers.utils.quantization_config import QuantizationConfigMixin
+from olive.common.hf.wrapper import ModelWrapper
+from olive.common.quant.patterns import match_override, match_skip
from olive.common.utils import StrEnumBase
if TYPE_CHECKING:
@@ -58,8 +60,21 @@ class OliveHfQuantizationConfig(QuantizationConfigMixin):
-1 = per-channel, >0 = groupwise.
lm_head: Whether to quantize the language model head.
embeds: Whether to quantize the input embeddings.
- modules_to_not_convert : List of module names to exclude from quantization.
+ moe: Whether to quantize MoE expert modules / parameters.
+ When ``False`` (default), every ``nn.Module`` under each
+ experts subtree returned by ``LayerWrapper.get_experts()``
+ is added to the skip set — this both leaves fused-3D
+ experts alone *and* fixes the previous silent quantization
+ of per-expert ``nn.Linear``s in ``ModuleList(Expert)``
+ blocks (Mixtral, PhiMoE, Qwen2/3-MoE).
+ modules_to_not_convert: List of module name patterns to exclude
+ from quantization. Plain strings use **substring** matching
+ (preserving HF semantics); entries prefixed with ``re:`` use
+ ``re.fullmatch``.
overrides: Per-module overrides for quantization parameters.
+ Keys use **literal equality** matching by default; entries
+ prefixed with ``re:`` use ``re.fullmatch``. Among matching
+ keys, the longest pattern wins (ties broken lexically).
"""
@@ -71,6 +86,7 @@ def __init__(
group_size: int,
lm_head: bool = False,
embeds: bool = False,
+ moe: bool = False,
modules_to_not_convert: list | None = None,
overrides: dict | None = None,
tie_word_embeddings: bool = False,
@@ -84,6 +100,7 @@ def __init__(
self.group_size = group_size
self.lm_head = lm_head
self.embeds = embeds
+ self.moe = moe
self.modules_to_not_convert = modules_to_not_convert
self.overrides = {
module_name: OliveHfQuantizationOverrideConfig(**override)
@@ -127,7 +144,9 @@ def get_qlinear_init_args(self, module_name: str) -> dict:
"symmetric": self.symmetric,
"group_size": self.group_size,
}
- if override := self.overrides.get(module_name):
+ best = match_override(module_name, list(self.overrides.keys())) if self.overrides else None
+ if best is not None:
+ override = self.overrides[best]
init_args.update({k: v for k, v in override.__dict__.items() if v is not None})
return init_args
@@ -153,24 +172,49 @@ def _process_model_before_weight_loading(
):
from olive.common.quant.nn import QuantEmbedding, QuantLinear
- ids_to_skip = []
+ ids_to_skip: list[int] = []
if not self.quantization_config.lm_head:
ids_to_skip.append(id(model.get_output_embeddings()))
if not self.quantization_config.embeds:
ids_to_skip.append(id(model.get_input_embeddings()))
- self.modules_to_not_convert = (
+ if not self.quantization_config.moe:
+ # Skip every nn.Module under each experts subtree — this both
+ # leaves fused-3D experts alone *and* fixes the previous silent
+ # quantization of per-expert nn.Linears inside
+ # ModuleList(Expert) blocks (Mixtral, PhiMoE, Qwen2/3-MoE).
+ try:
+ wrapper = ModelWrapper.from_model(model)
+ for lw in wrapper.get_layer_wrappers():
+ experts = lw.get_experts(return_name=False)
+ if experts is None:
+ continue
+ for sub in experts.modules():
+ ids_to_skip.append(id(sub))
+ except Exception: # pylint: disable=broad-except
+ # Not every model is wrappable (e.g., random tests). Falling
+ # back to the previous behaviour (no experts skip) is safe.
+ pass
+
+ modules_to_not_convert: list[str] = (
[name for name, module in model.named_modules() if id(module) in ids_to_skip] if ids_to_skip else []
)
+ # Pattern-aware skip list. Plain strings keep their HF substring
+ # semantics; ``re:`` opts into regex fullmatch.
+ skip_patterns: list[str] = []
if self.quantization_config.modules_to_not_convert:
- self.modules_to_not_convert.extend(self.quantization_config.modules_to_not_convert)
+ skip_patterns.extend(self.quantization_config.modules_to_not_convert)
if keep_in_fp32_modules:
- self.modules_to_not_convert.extend(keep_in_fp32_modules)
+ skip_patterns.extend(keep_in_fp32_modules)
+ self.modules_to_not_convert = modules_to_not_convert
+ self._skip_patterns = skip_patterns
def should_quantize(module: nn.Module, name: str) -> bool:
"""Check if a module should be quantized."""
- return isinstance(module, (nn.Linear, nn.Embedding)) and not any(
- key in name for key in self.modules_to_not_convert
- )
+ if not isinstance(module, (nn.Linear, nn.Embedding)):
+ return False
+ if any(literal == name for literal in self.modules_to_not_convert):
+ return False
+ return not match_skip(name, self._skip_patterns)
def create_quantized_module(module: nn.Linear | nn.Embedding, name: str) -> QuantLinear | QuantEmbedding:
"""Create a quantized version of a module."""
diff --git a/olive/common/quant/patterns.py b/olive/common/quant/patterns.py
new file mode 100644
index 0000000000..d2fed4af44
--- /dev/null
+++ b/olive/common/quant/patterns.py
@@ -0,0 +1,80 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Pattern matching helpers for Olive quantization module selection.
+
+`overrides` keys and `modules_to_not_convert` entries use the following
+semantics:
+
+- Plain string keys are *literal* for ``overrides`` (equality match) and
+ *substring* for ``modules_to_not_convert`` (matches the existing HF
+ ``modules_to_not_convert`` semantics).
+- Keys prefixed with ``re:`` opt into regular-expression matching via
+ ``re.fullmatch``.
+
+These helpers are the single source of truth for the matching logic and
+are used by both ``OliveHfQuantizationConfig`` and the Olive walker.
+"""
+
+from __future__ import annotations
+
+import re
+from functools import cache
+
+REGEX_PREFIX = "re:"
+
+
+def is_regex_pattern(pattern: str) -> bool:
+ """Return True if ``pattern`` opts into regex matching."""
+ return isinstance(pattern, str) and pattern.startswith(REGEX_PREFIX)
+
+
+@cache
+def _compiled(pattern: str) -> re.Pattern:
+ return re.compile(pattern[len(REGEX_PREFIX) :])
+
+
+def match_override(name: str, patterns) -> str | None:
+ """Find the best override pattern for ``name``.
+
+ Plain string patterns match by **literal equality**. ``re:`` patterns
+ match by ``re.fullmatch``. The most specific match (longest pattern
+ string) wins; ties are broken lexicographically so the choice is
+ deterministic.
+
+ Returns the original pattern (with prefix preserved) so the caller
+ can look it up in the underlying overrides dict, or ``None`` if no
+ pattern matched.
+ """
+ if not patterns:
+ return None
+ matches: list[str] = []
+ for pattern in patterns:
+ if is_regex_pattern(pattern):
+ if _compiled(pattern).fullmatch(name):
+ matches.append(pattern)
+ elif name == pattern:
+ matches.append(pattern)
+ if not matches:
+ return None
+ # longest pattern first, then lexical
+ return sorted(matches, key=lambda p: (-len(p), p))[0]
+
+
+def match_skip(name: str, patterns) -> bool:
+ """Return True if ``name`` is matched by any skip pattern.
+
+ Plain string patterns use **substring** matching to preserve the
+ existing HF ``modules_to_not_convert`` semantics. ``re:`` patterns
+ use ``re.fullmatch``.
+ """
+ if not patterns:
+ return False
+ for pattern in patterns:
+ if is_regex_pattern(pattern):
+ if _compiled(pattern).fullmatch(name):
+ return True
+ elif pattern and pattern in name:
+ return True
+ return False
diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py
new file mode 100644
index 0000000000..7cb181a7a5
--- /dev/null
+++ b/olive/common/quant/tensor.py
@@ -0,0 +1,490 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""``QuantTensor`` — a wrapper ``torch.Tensor`` subclass that stores weight
+quantization buffers (``qweight``, ``scales``, ``qzeros``) but presents the
+shape / dtype / device of the dequantized full-precision weight.
+
+Design notes (see ``/datadisks/jambaykinley/archive/qmoe/research.md`` and
+``plan.md``):
+
+* The class is a **wrapper** subclass (``_make_wrapper_subclass``) — it
+ carries no real storage of its own, so the dense FP weight is never
+ materialised in memory; only the packed buffers are allocated.
+* ``F.linear`` and ``F.embedding`` are dispatched via
+ ``__torch_function__``:
+ - Eager: unpack + dequantize on the fly and forward to the dense op.
+ - Under ``torch.onnx.is_in_onnx_export()``: raise — Olive's ONNX
+ conversion pass swaps any ``nn.Linear`` / ``nn.Embedding`` whose
+ weight is a ``QuantTensor`` for the existing exportable
+ ``QuantLinearNbit`` / ``QuantEmbeddingNbit`` ``nn.Module``s (see
+ ``olive/common/quant/nn.py`` and ``olive/common/hf/quant.py``)
+ *before* the tracer ever inspects the parameter. This keeps the
+ legacy ``com.microsoft::MatMulNBits`` /
+ ``com.microsoft::GatherBlockQuantized`` symbolic emission intact.
+* All other ops (including ``model.to(dtype/device)``, ``.detach()``,
+ ``.contiguous()``, ``.clone()``) are routed through ``_apply_fn_to_data``
+ via ``__torch_dispatch__`` so the inner buffers move with the wrapper.
+* For 3D fused MoE experts (``(num_experts, out, in)``) the same buffers
+ carry an additional leading dim. ``__getitem__`` / ``index_select`` on
+ the leading dim return a 2D ``QuantTensor`` (so per-expert
+ ``F.linear(current_state, weight[expert_idx])`` continues to dispatch
+ through the same code path).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import torch
+import torch.nn.functional as F
+
+from olive.common.quant.utils import (
+ WeightQuantizer,
+ pack_to_uint8,
+ unpack_from_uint8,
+ unpack_from_uint8_along_last,
+)
+
+__all__ = ["QuantTensor", "implements"]
+
+
+_TORCH_FN_TABLE: dict[Callable, Callable] = {}
+
+
+def implements(*torch_fns: Callable) -> Callable[[Callable], Callable]:
+ """Decorator to register a torch-function override for ``QuantTensor``."""
+
+ def decorator(fn: Callable) -> Callable:
+ for torch_fn in torch_fns:
+ _TORCH_FN_TABLE[torch_fn] = fn
+ return fn
+
+ return decorator
+
+
+def _midq(bits: int) -> int:
+ return 1 << (bits - 1)
+
+
+def _zero_points_or_default(weight: QuantTensor) -> torch.Tensor:
+ """Unpack zero_points or return a tensor full of the symmetric mid-q value."""
+ if weight.qzeros is not None:
+ return unpack_from_uint8_along_last(weight.qzeros, weight.bits, tuple(weight.scales.shape)).to(torch.int32)
+ return torch.full(weight.scales.shape, _midq(weight.bits), dtype=torch.int32, device=weight.scales.device)
+
+
+def _dequantize(weight: QuantTensor) -> torch.Tensor:
+ """Unpack + dequantize ``weight`` into a dense tensor of ``weight.dtype``."""
+ quantizer = WeightQuantizer(
+ bits=weight.bits, symmetric=weight.symmetric, group_size=weight.group_size, signed=False
+ )
+ if weight.dim() == 2:
+ qw = unpack_from_uint8(weight.qweight, weight.bits, tuple(weight.shape))
+ zp = _zero_points_or_default(weight)
+ return quantizer.dequantize(qw, weight.scales, zp).to(weight.dtype)
+ if weight.dim() == 3:
+ qw = unpack_from_uint8_along_last(weight.qweight, weight.bits, tuple(weight.shape))
+ zp = _zero_points_or_default(weight)
+ leading = weight.shape[0]
+ out = torch.empty(weight.shape, dtype=weight.dtype, device=weight.qweight.device)
+ for i in range(leading):
+ out[i] = quantizer.dequantize(qw[i], weight.scales[i], zp[i]).to(weight.dtype)
+ return out
+ raise NotImplementedError(f"QuantTensor only supports 2D / 3D layouts, got {weight.dim()}D")
+
+
+class QuantTensor(torch.Tensor):
+ """A weight-quantized tensor.
+
+ Holds:
+ qweight: ``torch.uint8`` packed quantized values along the last
+ dim. Shape ``(*, math.ceil(in_features * bits / 8))``.
+ scales: per-group scales, dtype matches the dequantized dtype.
+ qzeros: ``torch.uint8`` packed zero-points, or ``None`` for
+ symmetric quantization.
+
+ Attributes (non-tensor):
+ bits, group_size, symmetric.
+
+ The shape / dtype / device exposed via the wrapper subclass are
+ those of the **dequantized** weight, so the host ``nn.Linear`` /
+ ``nn.Embedding`` continues to see the right metadata.
+ """
+
+ qweight: torch.Tensor
+ scales: torch.Tensor
+ qzeros: torch.Tensor | None
+ bits: int
+ group_size: int
+ symmetric: bool
+
+ @staticmethod
+ def __new__(
+ cls,
+ qweight: torch.Tensor,
+ scales: torch.Tensor,
+ qzeros: torch.Tensor | None,
+ bits: int,
+ group_size: int,
+ symmetric: bool,
+ shape: torch.Size | tuple[int, ...],
+ dtype: torch.dtype,
+ ) -> QuantTensor:
+ return torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined]
+ cls,
+ tuple(shape),
+ dtype=dtype,
+ device=qweight.device,
+ requires_grad=False,
+ )
+
+ def __init__(
+ self,
+ qweight: torch.Tensor,
+ scales: torch.Tensor,
+ qzeros: torch.Tensor | None,
+ bits: int,
+ group_size: int,
+ symmetric: bool,
+ shape: torch.Size | tuple[int, ...],
+ dtype: torch.dtype,
+ ) -> None:
+ self.qweight = qweight
+ self.scales = scales
+ self.qzeros = qzeros
+ self.bits = int(bits)
+ self.group_size = int(group_size)
+ self.symmetric = bool(symmetric)
+
+ # ------------------------------------------------------------------
+ # Construction helpers
+ # ------------------------------------------------------------------
+
+ @classmethod
+ def from_float(
+ cls,
+ weight: torch.Tensor,
+ bits: int = 4,
+ symmetric: bool = True,
+ group_size: int = -1,
+ scales: torch.Tensor | None = None,
+ zero_points: torch.Tensor | None = None,
+ ) -> QuantTensor:
+ """Quantize a 2D or 3D FP weight tensor and produce a ``QuantTensor``.
+
+ For 3D weights the leading dim is iterated and each slice is
+ quantized independently — the resulting qweight retains the
+ leading dim and ``scales`` / ``qzeros`` gain one as well.
+ """
+ if weight.dim() not in (2, 3):
+ raise ValueError(f"QuantTensor only supports 2D and 3D weights, got shape {tuple(weight.shape)}")
+
+ quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False)
+ if weight.dim() == 2:
+ if scales is None or zero_points is None:
+ scales, zero_points = quantizer.find_qparams(weight)
+ else:
+ scales = scales.to(weight.device).to(weight.dtype)
+ zero_points = zero_points.to(weight.device).to(torch.int32)
+ qweight_int = quantizer.quantize(weight, scales, zero_points)
+ qweight_packed = pack_to_uint8(qweight_int, bits).contiguous()
+ scales_packed = scales.reshape(quantizer.get_qparam_shape(weight.shape)).contiguous()
+ if symmetric:
+ qzeros_packed = None
+ if not torch.all(zero_points == quantizer.midq):
+ raise ValueError("Zero points must equal midq for symmetric quantization")
+ else:
+ qzeros_packed = pack_to_uint8(
+ zero_points.reshape(quantizer.get_qparam_shape(weight.shape)), bits
+ ).contiguous()
+ else:
+ # 3D — iterate the leading dim
+ if scales is not None or zero_points is not None:
+ raise NotImplementedError("Pre-computed scales/zero_points are not yet supported for 3D weights")
+ qweight_list, scales_list, zp_list = [], [], []
+ for i in range(weight.shape[0]):
+ s, zp = quantizer.find_qparams(weight[i])
+ qw = quantizer.quantize(weight[i], s, zp)
+ qweight_list.append(pack_to_uint8(qw, bits))
+ scales_list.append(s.reshape(quantizer.get_qparam_shape(weight[i].shape)))
+ zp_list.append(zp.reshape(quantizer.get_qparam_shape(weight[i].shape)))
+ qweight_packed = torch.stack(qweight_list, dim=0).contiguous()
+ scales_packed = torch.stack(scales_list, dim=0).contiguous()
+ if symmetric:
+ qzeros_packed = None
+ for zp in zp_list:
+ if not torch.all(zp == quantizer.midq):
+ raise ValueError("Zero points must equal midq for symmetric quantization")
+ else:
+ qzeros_packed = torch.stack([pack_to_uint8(zp, bits) for zp in zp_list], dim=0).contiguous()
+
+ return cls(
+ qweight=qweight_packed,
+ scales=scales_packed,
+ qzeros=qzeros_packed,
+ bits=bits,
+ group_size=group_size,
+ symmetric=symmetric,
+ shape=tuple(weight.shape),
+ dtype=scales_packed.dtype,
+ )
+
+ @classmethod
+ def from_packed(
+ cls,
+ qweight: torch.Tensor,
+ scales: torch.Tensor,
+ qzeros: torch.Tensor | None,
+ bits: int,
+ group_size: int,
+ symmetric: bool,
+ shape: tuple[int, ...],
+ ) -> QuantTensor:
+ """Reconstruct a ``QuantTensor`` from already-packed buffers."""
+ return cls(
+ qweight=qweight,
+ scales=scales,
+ qzeros=qzeros,
+ bits=bits,
+ group_size=group_size,
+ symmetric=symmetric,
+ shape=shape,
+ dtype=scales.dtype,
+ )
+
+ # ------------------------------------------------------------------
+ # Dequantization
+ # ------------------------------------------------------------------
+
+ def to_dense(self) -> torch.Tensor:
+ """Unpack + dequantize into a dense FP tensor of ``self.dtype``."""
+ return _dequantize(self)
+
+ # ------------------------------------------------------------------
+ # Flatten / Unflatten for torch.compile and friends
+ # ------------------------------------------------------------------
+
+ def __tensor_flatten__(self):
+ names = ["qweight", "scales"]
+ if self.qzeros is not None:
+ names.append("qzeros")
+ meta = {
+ "bits": self.bits,
+ "group_size": self.group_size,
+ "symmetric": self.symmetric,
+ "shape": tuple(self.shape),
+ "dtype": self.dtype,
+ "has_qzeros": self.qzeros is not None,
+ }
+ return names, meta
+
+ @classmethod
+ def __tensor_unflatten__(cls, inner_tensors, meta, outer_size, outer_stride):
+ return cls(
+ qweight=inner_tensors["qweight"],
+ scales=inner_tensors["scales"],
+ qzeros=inner_tensors["qzeros"] if meta["has_qzeros"] else None,
+ bits=meta["bits"],
+ group_size=meta["group_size"],
+ symmetric=meta["symmetric"],
+ shape=meta["shape"],
+ dtype=meta["dtype"],
+ )
+
+ # ------------------------------------------------------------------
+ # _apply_fn_to_data — propagate per-tensor transforms (.to, detach…)
+ # through every inner buffer
+ # ------------------------------------------------------------------
+
+ def _apply_fn_to_data(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> QuantTensor:
+ new_qweight = fn(self.qweight)
+ new_scales = fn(self.scales)
+ new_qzeros = fn(self.qzeros) if self.qzeros is not None else None
+ return QuantTensor(
+ qweight=new_qweight,
+ scales=new_scales,
+ qzeros=new_qzeros,
+ bits=self.bits,
+ group_size=self.group_size,
+ symmetric=self.symmetric,
+ shape=tuple(self.shape),
+ dtype=new_scales.dtype,
+ )
+
+ # ------------------------------------------------------------------
+ # Dispatch
+ # ------------------------------------------------------------------
+
+ @classmethod
+ def __torch_function__(cls, func, types, args=(), kwargs=None):
+ kwargs = kwargs or {}
+ handler = _TORCH_FN_TABLE.get(func)
+ if handler is not None:
+ return handler(*args, **kwargs)
+ # Fall through to __torch_dispatch__ for everything else.
+ return super().__torch_function__(func, types, args, kwargs)
+
+ @classmethod
+ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+ kwargs = kwargs or {}
+ aten = torch.ops.aten
+
+ if func in (aten.detach.default, aten.clone.default, aten.alias.default, aten.contiguous.default):
+ (self_,) = args
+ return self_._apply_fn_to_data(lambda x: func(x))
+
+ if func is aten._to_copy.default:
+ (self_,) = args
+ dtype = kwargs.get("dtype")
+ device = kwargs.get("device")
+
+ def _move(x: torch.Tensor) -> torch.Tensor:
+ copy_kwargs: dict[str, Any] = {}
+ if device is not None:
+ copy_kwargs["device"] = device
+ # only scales are real-dtype; keep qweight/qzeros as uint8
+ if dtype is not None and x.is_floating_point():
+ copy_kwargs["dtype"] = dtype
+ return func(x, **copy_kwargs) if copy_kwargs else x
+
+ return self_._apply_fn_to_data(_move)
+
+ if func is aten.copy_.default:
+ self_, src = args
+ if not isinstance(src, QuantTensor):
+ raise TypeError(f"Cannot copy_ a non-QuantTensor source into a QuantTensor (got {type(src)})")
+ self_.qweight.copy_(src.qweight)
+ self_.scales.copy_(src.scales)
+ if self_.qzeros is not None and src.qzeros is not None:
+ self_.qzeros.copy_(src.qzeros)
+ return self_
+
+ # Fallback: dequantize any QuantTensor args and re-dispatch.
+ new_args = [_maybe_dense(a) for a in args]
+ new_kwargs = {k: _maybe_dense(v) for k, v in kwargs.items()}
+ return func(*new_args, **new_kwargs)
+
+ # Friendlier repr — full dequant would defeat the purpose.
+ def __repr__(self) -> str: # pragma: no cover - trivial
+ return (
+ f"QuantTensor(shape={tuple(self.shape)}, dtype={self.dtype}, device={self.device}, "
+ f"bits={self.bits}, group_size={self.group_size}, symmetric={self.symmetric})"
+ )
+
+
+def _maybe_dense(x: Any) -> Any:
+ if isinstance(x, QuantTensor):
+ return x.to_dense()
+ return x
+
+
+# ----------------------------------------------------------------------
+# Torch-function overrides
+# ----------------------------------------------------------------------
+
+
+@implements(F.linear)
+def _linear(input: torch.Tensor, weight: QuantTensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+ if torch.onnx.is_in_onnx_export():
+ raise RuntimeError(
+ "Olive QuantTensor cannot be traced by torch.onnx.export directly. "
+ "Use olive.common.hf.quant.make_export_compatible_quant(model, dynamo=...) "
+ "before exporting, which replaces nn.Linear modules backed by a "
+ "QuantTensor with an exportable QuantLinearNbit nn.Module."
+ )
+ if weight.dim() != 2:
+ raise RuntimeError(
+ "F.linear expects a 2D weight; got a "
+ f"{weight.dim()}D QuantTensor. For 3D fused MoE experts, slice the leading "
+ "dim first (e.g. `weight[expert_idx]`)."
+ )
+ dense = weight.to_dense().to(input.dtype)
+ return F.linear(input, dense, bias)
+
+
+@implements(F.embedding)
+def _embedding(
+ input: torch.Tensor,
+ weight: QuantTensor,
+ padding_idx: int | None = None,
+ max_norm: float | None = None,
+ norm_type: float = 2.0,
+ scale_grad_by_freq: bool = False,
+ sparse: bool = False,
+) -> torch.Tensor:
+ if torch.onnx.is_in_onnx_export():
+ raise RuntimeError(
+ "Olive QuantTensor cannot be traced by torch.onnx.export directly. "
+ "Use olive.common.hf.quant.make_export_compatible_quant(model, dynamo=...) "
+ "before exporting, which replaces nn.Embedding modules backed by a "
+ "QuantTensor with an exportable QuantEmbeddingNbit nn.Module."
+ )
+ if weight.dim() != 2:
+ raise RuntimeError(f"F.embedding expects a 2D weight; got a {weight.dim()}D QuantTensor.")
+ dense = weight.to_dense()
+ return F.embedding(input, dense, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
+
+
+@implements(torch.matmul, torch.Tensor.matmul)
+def _matmul(a, b):
+ if torch.onnx.is_in_onnx_export() and (isinstance(a, QuantTensor) or isinstance(b, QuantTensor)):
+ raise RuntimeError(
+ "ONNX export of matmul on a QuantTensor is not supported in Olive. "
+ "Olive's MoE quantization is storage-only; use Mobius to emit "
+ "com.microsoft.QMoE or a per-expert MatMulNBits loop."
+ )
+ return torch.matmul(_maybe_dense(a), _maybe_dense(b))
+
+
+@implements(torch.bmm)
+def _bmm(a, b):
+ if torch.onnx.is_in_onnx_export() and (isinstance(a, QuantTensor) or isinstance(b, QuantTensor)):
+ raise RuntimeError(
+ "ONNX export of bmm on a QuantTensor is not supported in Olive. "
+ "Olive's MoE quantization is storage-only; use Mobius to emit "
+ "com.microsoft.QMoE or a per-expert MatMulNBits loop."
+ )
+ return torch.bmm(_maybe_dense(a), _maybe_dense(b))
+
+
+@implements(torch.Tensor.__getitem__)
+def _getitem(self: QuantTensor, idx):
+ # Slicing the leading dim of a 3D QuantTensor returns a 2D
+ # QuantTensor when possible — keeps MoE per-expert forwards
+ # (`weight[expert_idx]`) on the quantized fast path.
+ if self.dim() == 3 and isinstance(idx, int):
+ new_shape = tuple(self.shape[1:])
+ return QuantTensor(
+ qweight=self.qweight[idx],
+ scales=self.scales[idx],
+ qzeros=self.qzeros[idx] if self.qzeros is not None else None,
+ bits=self.bits,
+ group_size=self.group_size,
+ symmetric=self.symmetric,
+ shape=new_shape,
+ dtype=self.dtype,
+ )
+ return self.to_dense()[idx]
+
+
+@implements(torch.Tensor.to)
+def _to(self: QuantTensor, *args, **kwargs):
+ # Use torch's own _parse_to to robustly resolve the (device, dtype,
+ # non_blocking, convert_to_format) tuple — covers every signature
+ # including nn.Module.to's ``t.to(None, dtype, non_blocking)``.
+ device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs) # type: ignore[attr-defined]
+
+ if device is None and dtype is None:
+ return self
+
+ def _move(x: torch.Tensor) -> torch.Tensor:
+ move_kwargs: dict[str, Any] = {}
+ if device is not None:
+ move_kwargs["device"] = device
+ if dtype is not None and x.is_floating_point():
+ move_kwargs["dtype"] = dtype
+ return x.to(**move_kwargs) if move_kwargs else x
+
+ return self._apply_fn_to_data(_move)
diff --git a/olive/common/quant/utils.py b/olive/common/quant/utils.py
index 8163e2b4fa..2aac15477f 100644
--- a/olive/common/quant/utils.py
+++ b/olive/common/quant/utils.py
@@ -9,7 +9,14 @@
# TODO(jambayk): consider supporting transposed weights, useful for onnx weights
class WeightQuantizer:
- """Class to quantize weight tensors."""
+ """Class to quantize weight tensors.
+
+ Operates on 2D tensors of shape ``(out_features, in_features)`` with
+ the last dimension quantized (groupwise / per-channel / per-tensor).
+ Use :func:`quantize_along_leading_dim` to handle 3D tensors (e.g.
+ fused MoE experts of shape ``(num_experts, out, in)``) by iterating
+ the leading dimension.
+ """
def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, signed: bool = False):
"""Initialize the quantizer with parameters.
@@ -257,3 +264,69 @@ def unpack_from_uint8(packed_tensor: torch.Tensor, bits: int, shape: tuple[int,
unpacked_tensor = unpacked_tensor.reshape(shape[0], -1)
unpacked_tensor = unpacked_tensor[:, : shape[1]]
return torch.bitwise_and(unpacked_tensor, maxq).to(torch.int32)
+
+
+@torch.no_grad()
+def quantize_along_leading_dim(
+ quantizer: WeightQuantizer, tensor: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Quantize a 2D or 3D tensor with a :class:`WeightQuantizer`.
+
+ For a 2D input of shape ``(out, in)`` this is equivalent to
+ ``quantizer.find_qparams`` + ``quantizer.quantize``.
+
+ For a 3D input of shape ``(leading, out, in)`` the leading dimension
+ is iterated and each slice is quantized independently. The returned
+ ``qweight`` shares the same shape as the input; ``scales`` /
+ ``zero_points`` carry an additional leading dimension matching the
+ input's leading dim.
+ """
+ if tensor.dim() == 2:
+ scales, zero_points = quantizer.find_qparams(tensor)
+ qweight = quantizer.quantize(tensor, scales, zero_points)
+ return qweight, scales, zero_points
+ if tensor.dim() != 3:
+ raise ValueError(f"Only 2D and 3D tensors are supported, got shape {tuple(tensor.shape)}")
+
+ leading = tensor.shape[0]
+ qweights, scales_list, zp_list = [], [], []
+ for i in range(leading):
+ s, zp = quantizer.find_qparams(tensor[i])
+ qweights.append(quantizer.quantize(tensor[i], s, zp))
+ scales_list.append(s)
+ zp_list.append(zp)
+ return (
+ torch.stack(qweights, dim=0),
+ torch.stack(scales_list, dim=0),
+ torch.stack(zp_list, dim=0),
+ )
+
+
+@torch.no_grad()
+def pack_to_uint8_along_last(tensor: torch.Tensor, bits: int) -> torch.Tensor:
+ """Pack 2/4/8 bit values into uint8 along the last dim, supporting 2D or 3D.
+
+ For 3D inputs (e.g. fused MoE experts) the leading dim is preserved.
+ """
+ if tensor.dim() == 2:
+ return pack_to_uint8(tensor, bits)
+ if tensor.dim() != 3:
+ raise ValueError(f"Only 2D and 3D tensors are supported, got shape {tuple(tensor.shape)}")
+
+ leading = tensor.shape[0]
+ return torch.stack([pack_to_uint8(tensor[i], bits) for i in range(leading)], dim=0)
+
+
+@torch.no_grad()
+def unpack_from_uint8_along_last(packed_tensor: torch.Tensor, bits: int, shape: tuple[int, ...]) -> torch.Tensor:
+ """Unpack uint8-packed values into 2/4/8 bit ints, supporting 2D or 3D."""
+ if len(shape) == 2:
+ return unpack_from_uint8(packed_tensor, bits, shape)
+ if len(shape) != 3:
+ raise ValueError(f"Only 2D and 3D tensors are supported, got shape {shape}")
+
+ leading = shape[0]
+ return torch.stack(
+ [unpack_from_uint8(packed_tensor[i], bits, shape[1:]) for i in range(leading)],
+ dim=0,
+ )
diff --git a/olive/passes/onnx/model_builder.py b/olive/passes/onnx/model_builder.py
index e1062e02b9..6f86201bd2 100644
--- a/olive/passes/onnx/model_builder.py
+++ b/olive/passes/onnx/model_builder.py
@@ -437,6 +437,14 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
config = quant_attrs["config"]
+ if config.get("moe"):
+ raise NotImplementedError(
+ "ModelBuilder does not support loading Olive-quantized MoE checkpoints "
+ "(``quantization_config.moe == True``). Use the Mobius model builder for "
+ "MoE models or rerun the RTN pass with ``moe=False`` to leave experts in "
+ "their original precision."
+ )
+
self.quant_type = quant_type
self.embedding = QuantizedTensorModule() if config["embeds"] else TensorModule()
self.final_norm = TensorModule()
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index bdfd08e0da..b4c112708a 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -18,6 +18,8 @@
tie_quant_word_embeddings,
)
from olive.common.quant.nn import QuantEmbedding, QuantLinear
+from olive.common.quant.patterns import match_skip
+from olive.common.quant.tensor import QuantTensor
from olive.common.quant.utils import WeightQuantizer
from olive.common.utils import tensor_data_to_device
from olive.constants import PrecisionBits
@@ -33,7 +35,7 @@
logger = logging.getLogger(__name__)
-def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigParam]:
+def get_quantizer_config(allow_embeds: bool = False, allow_moe: bool = False) -> dict[str, PassConfigParam]:
return {
"bits": PassConfigParam(
type_=PrecisionBits,
@@ -66,13 +68,37 @@ def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigPara
if allow_embeds
else {}
),
+ **(
+ {
+ "moe": PassConfigParam(
+ type_=bool,
+ default_value=False,
+ description=(
+ "Whether to quantize MoE expert modules / parameters. When False (default), every "
+ "nn.Module under each experts subtree is skipped. Defaults reuse the pass-level "
+ "bits/group_size/sym settings; use ``overrides`` to tune experts independently."
+ ),
+ )
+ }
+ if allow_moe
+ else {}
+ ),
+ "modules_to_not_convert": PassConfigParam(
+ type_=list,
+ default_value=None,
+ description=(
+ "Optional list of module name patterns to exclude from quantization. Plain strings use"
+ " substring matching (HF semantics); entries prefixed with 're:' use re.fullmatch."
+ ),
+ ),
"overrides": PassConfigParam(
type_=dict,
default_value=None,
description=(
- "Optional dictionary to specify overrides for specific modules. The keys are module names and the"
- " values are dictionaries with any of the following keys: 'bits', 'symmetric', 'group_size'. These"
- " overrides take precedence over the overrides provided in the mixed precision info."
+ "Optional dictionary to specify overrides for specific modules. The keys are module names"
+ " (literal match) or 're:' patterns (regex fullmatch). Values are dictionaries with"
+ " any of the following keys: 'bits', 'symmetric', 'group_size'. These overrides take"
+ " precedence over the overrides provided in the mixed precision info."
),
),
}
@@ -126,9 +152,35 @@ def prepare_model(
else:
excluded_attn_inputs.update(attn_inputs[:2])
+ # Collect every ``nn.Module`` under any experts subtree, so we can
+ # honour the ``moe`` category flag the same way ``lm_head`` /
+ # ``embeds`` are honoured today.
+ expert_module_ids: set[int] = set()
+ expert_owners: list[tuple[torch.nn.Module, str]] = [] # (experts_module, dotted_name)
+ for layer_wrapper in wrapper.get_layer_wrappers():
+ experts_module = layer_wrapper.get_experts(return_name=False)
+ if experts_module is None:
+ continue
+ # Find the dotted name of this experts module relative to the model root.
+ experts_name = None
+ for name, mod in wrapper.model.named_modules():
+ if mod is experts_module:
+ experts_name = name
+ break
+ expert_owners.append((experts_module, experts_name or ""))
+ for sub in experts_module.modules():
+ expert_module_ids.add(id(sub))
+
+ skip_patterns = list(getattr(qcfg, "modules_to_not_convert", None) or [])
+
def should_quantize(module: torch.nn.Module, name: str) -> bool:
if module in excluded_attn_inputs:
return False
+ if match_skip(name, skip_patterns):
+ return False
+ # category-flag skips (lm_head / embeds / moe) — first rule wins
+ if id(module) in expert_module_ids and not getattr(qcfg, "moe", False):
+ return False
if isinstance(module, torch.nn.Linear):
return name != lm_head_name or qcfg.lm_head
if qcfg.embeds and isinstance(module, torch.nn.Embedding):
@@ -144,8 +196,30 @@ def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module:
replace_matching_submodules(wrapper.model, should_quantize, add_quant_info, description="Preparing model")
+ # Fused-3D MoE: experts modules expose 3D nn.Parameters directly (e.g.
+ # ``gate_up_proj`` of shape ``(num_experts, *, *)``). Annotate the
+ # experts module with a per-parameter quant_info_3d dict so
+ # ``finalize`` can replace each parameter with a QuantTensor.
+ if getattr(qcfg, "moe", False):
+ for experts_module, experts_name in expert_owners:
+ param_qinfos: dict[str, QuantInfo] = {}
+ for pname, param in experts_module.named_parameters(recurse=False):
+ if param.dim() != 3:
+ continue
+ full_name = f"{experts_name}.{pname}" if experts_name else pname
+ if match_skip(full_name, skip_patterns):
+ continue
+ qargs = qcfg.get_qlinear_init_args(full_name)
+ param_qinfos[pname] = QuantInfo(quantizer=WeightQuantizer(**qargs))
+ new_qargs[full_name] = qargs
+ if param_qinfos:
+ experts_module.quant_info_3d = param_qinfos
+
# remove overrides for modules not being quantized
for name in list(qcfg.overrides or {}):
+ # ``re:`` keys aren't tied to a specific module, so leave them in place.
+ if name.startswith("re:"):
+ continue
if name not in new_qargs:
qcfg.overrides.pop(name)
@@ -159,6 +233,7 @@ def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module:
merged_qcfg_dict["overrides"][name] = override
merged_qcfg_dict["lm_head"] |= qcfg.lm_head
merged_qcfg_dict["embeds"] |= qcfg.embeds
+ merged_qcfg_dict["moe"] = merged_qcfg_dict.get("moe", False) or getattr(qcfg, "moe", False)
qcfg = OliveHfQuantizationConfig(**merged_qcfg_dict)
word_embeddings_eligible_for_tieing = (
@@ -188,6 +263,8 @@ def get_quant_config(model: HfModelHandler, config: type[BasePassConfig]) -> Oli
"group_size": config.group_size,
"lm_head": config.lm_head,
"embeds": getattr(config, "embeds", False),
+ "moe": getattr(config, "moe", False),
+ "modules_to_not_convert": getattr(config, "modules_to_not_convert", None) or [],
"overrides": config.overrides or {},
}
if mp_info := (model.model_attributes or {}).get("mixed_precision_info"):
@@ -457,6 +534,32 @@ def quantize_and_pack(module: torch.nn.Module, _: str) -> QuantLinear | QuantEmb
description="Quantizing and packing linear layers",
)
+ # Fused-3D MoE: experts modules carry a per-parameter ``quant_info_3d``
+ # dict. Replace each parameter with a 3D ``QuantTensor`` parameter.
+ # The host module's forward continues to find the parameter at the
+ # same attribute name; eager forwards may go through QuantTensor's
+ # ``__getitem__`` / matmul fallbacks. ONNX export of fused-3D MoE is
+ # deferred to Mobius — Olive only persists the checkpoint.
+ for sub in wrapper.model.modules():
+ info_3d: dict = getattr(sub, "quant_info_3d", None)
+ if not info_3d:
+ continue
+ for pname, qinfo in info_3d.items():
+ param = getattr(sub, pname)
+ if not isinstance(param, torch.nn.Parameter):
+ continue
+ with torch.no_grad():
+ quantizer = qinfo.quantizer
+ qt = QuantTensor.from_float(
+ param.detach().to(device).to(param.dtype),
+ bits=quantizer.bits,
+ symmetric=quantizer.symmetric,
+ group_size=quantizer.group_size,
+ ).to("cpu")
+ sub._parameters[pname] = torch.nn.Parameter(qt, requires_grad=False) # pylint: disable=W0212
+ # Tidy up the marker so downstream save/load doesn't trip over it.
+ delattr(sub, "quant_info_3d")
+
if retie_word_embeddings:
tie_quant_word_embeddings(wrapper.model)
quant_config.tie_word_embeddings = True
@@ -464,8 +567,55 @@ def quantize_and_pack(module: torch.nn.Module, _: str) -> QuantLinear | QuantEmb
wrapper.model.quantization_method = quant_config.quant_method
wrapper.model.config.quantization_config = quant_config
+ # Flatten any QuantTensor parameters into plain ``register_buffer`` entries so
+ # ``save_pretrained`` (which writes safetensors) can serialize them.
+ # Naming convention: ``_qweight`` / ``_scales`` / ``_qzeros``
+ # sit on the same parent module as the original parameter. This matches
+ # the suffix-style layout used by other prequantized formats and lets
+ # downstream loaders (Mobius) discover the buffers without a registry.
+ flatten_quant_tensor_params(wrapper.model)
+
# save the quantized model
wrapper.model.save_pretrained(output_model_path)
model.save_metadata(output_model_path)
return inherit_hf_from_hf(model, output_model_path, adapter_path=model.adapter_path)
+
+
+def flatten_quant_tensor_params(module: torch.nn.Module) -> None:
+ """Replace every ``QuantTensor`` ``nn.Parameter`` with plain registered buffers.
+
+ For each parameter ```` whose data is a :class:`QuantTensor`, the
+ parameter is deleted and the inner tensors are re-attached as
+ ``_qweight``, ``_scales`` (and ``_qzeros`` when
+ asymmetric) on the same parent module. Additional metadata
+ (``bits``, ``group_size``, ``symmetric``, ``shape``) is stored on
+ ``_quant_meta`` for round-trip loading.
+
+ This is used at save time so safetensors / pickle never sees a
+ tensor subclass.
+ """
+ for sub in module.modules():
+ names = [n for n, p in sub.named_parameters(recurse=False) if isinstance(p.data, QuantTensor)]
+ for name in names:
+ qt: QuantTensor = sub._parameters[name].data # type: ignore[assignment]
+ del sub._parameters[name]
+ sub.register_buffer(f"{name}_qweight", qt.qweight.detach().clone())
+ sub.register_buffer(f"{name}_scales", qt.scales.detach().clone())
+ if qt.qzeros is not None:
+ sub.register_buffer(f"{name}_qzeros", qt.qzeros.detach().clone())
+ # ``register_buffer`` requires Tensors, so the metadata lives
+ # outside the state_dict as a plain attribute. It is consumed
+ # by ``OliveHfQuantizer`` at load time via the
+ # ``quantization_config`` for shape/bits and is also encoded
+ # in the parameter's full name + scales shape if needed.
+ setattr(
+ sub,
+ f"{name}_quant_meta",
+ {
+ "bits": qt.bits,
+ "group_size": qt.group_size,
+ "symmetric": qt.symmetric,
+ "shape": tuple(qt.shape),
+ },
+ )
diff --git a/olive/passes/pytorch/rtn.py b/olive/passes/pytorch/rtn.py
index c7dbd462ee..1180e0a1dd 100644
--- a/olive/passes/pytorch/rtn.py
+++ b/olive/passes/pytorch/rtn.py
@@ -26,7 +26,7 @@ class Rtn(Pass):
@classmethod
def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]:
- return get_quantizer_config(allow_embeds=True)
+ return get_quantizer_config(allow_embeds=True, allow_moe=True)
@torch.no_grad()
def _run_for_config(
diff --git a/test/common/quant/test_hf_utils.py b/test/common/quant/test_hf_utils.py
index 92302fa5d0..c31fb3fe5f 100644
--- a/test/common/quant/test_hf_utils.py
+++ b/test/common/quant/test_hf_utils.py
@@ -429,3 +429,152 @@ def test_tie_word_embeddings(self):
# Now they should share buffers
assert model.embed.qweight is model.lm_head.qweight
assert model.embed.scales is model.lm_head.scales
+
+
+# -- MoE quantization fixtures and tests --------------------------------------
+
+
+class MoEConfig(PretrainedConfig):
+ model_type = "moe_simple"
+
+ def __init__(self, hidden_size=64, vocab_size=64, num_experts=4, num_layers=1, **kwargs):
+ super().__init__(**kwargs)
+ self.hidden_size = hidden_size
+ self.vocab_size = vocab_size
+ self.num_experts = num_experts
+ self.num_hidden_layers = num_layers
+ # required by ModelWrapper / LayerWrapper
+ self.num_attention_heads = 4
+ self.num_key_value_heads = 4
+ self.head_dim = hidden_size // 4
+ self.intermediate_size = hidden_size
+
+
+class _MoEExpert(nn.Module):
+ """Per-expert sub-module (ModuleList style — like Mixtral / PhiMoE)."""
+
+ def __init__(self, hidden_size: int):
+ super().__init__()
+ self.w1 = nn.Linear(hidden_size, hidden_size, bias=False)
+ self.w2 = nn.Linear(hidden_size, hidden_size, bias=False)
+
+
+class _MoELayer(nn.Module):
+ def __init__(self, hidden_size: int, num_experts: int):
+ super().__init__()
+ self.mlp = nn.Module()
+ self.mlp.gate = nn.Linear(hidden_size, num_experts, bias=False)
+ self.mlp.experts = nn.ModuleList([_MoEExpert(hidden_size) for _ in range(num_experts)])
+
+
+class MoESimpleModel(PreTrainedModel):
+ config_class = MoEConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = nn.Module()
+ self.model.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.model.layers = nn.ModuleList(
+ [_MoELayer(config.hidden_size, config.num_experts) for _ in range(config.num_hidden_layers)]
+ )
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def get_output_embeddings(self):
+ return self.lm_head
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_head = new_embeddings
+
+
+class TestOliveHfQuantizerMoE:
+ """MoE-specific behaviour of ``OliveHfQuantizer``.
+
+ The previous implementation silently quantized every per-expert
+ ``nn.Linear`` in ``ModuleList(Expert)`` blocks (Mixtral / PhiMoE /
+ Qwen2/3-MoE). With the new ``moe`` category flag (default ``False``),
+ every ``nn.Module`` under each experts subtree is added to the skip
+ set — fixing the silent quantization.
+ """
+
+ def test_module_list_experts_skipped_by_default(self):
+ """Regression: ModuleList(Expert) linears must stay as nn.Linear when moe=False."""
+ config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16, moe=False)
+ quantizer = OliveHfQuantizer(config)
+ model = MoESimpleModel(MoEConfig())
+
+ quantizer._process_model_before_weight_loading(model)
+
+ for expert in model.model.layers[0].mlp.experts:
+ assert isinstance(expert.w1, nn.Linear)
+ assert not isinstance(expert.w1, QuantLinear)
+ assert isinstance(expert.w2, nn.Linear)
+ assert not isinstance(expert.w2, QuantLinear)
+ # The router (gate) is also under the mlp but not under .experts;
+ # it should still be quantized.
+ assert isinstance(model.model.layers[0].mlp.gate, QuantLinear)
+
+ def test_module_list_experts_quantized_when_moe_true(self):
+ config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16, moe=True)
+ quantizer = OliveHfQuantizer(config)
+ model = MoESimpleModel(MoEConfig())
+
+ quantizer._process_model_before_weight_loading(model)
+
+ for expert in model.model.layers[0].mlp.experts:
+ assert isinstance(expert.w1, QuantLinear)
+ assert isinstance(expert.w2, QuantLinear)
+
+ def test_moe_default_is_false(self):
+ """``moe`` defaults to False — opt-in, matching lm_head / embeds."""
+ config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16)
+ assert config.moe is False
+
+ def test_regex_skip_pattern(self):
+ """``re:`` prefix opts into regex fullmatch for modules_to_not_convert."""
+ config = OliveHfQuantizationConfig(
+ bits=4,
+ symmetric=True,
+ group_size=16,
+ moe=True,
+ # Quantize experts but keep the router in fp.
+ modules_to_not_convert=["re:.*\\.mlp\\.gate"],
+ )
+ quantizer = OliveHfQuantizer(config)
+ model = MoESimpleModel(MoEConfig())
+
+ quantizer._process_model_before_weight_loading(model)
+
+ # gate is skipped via regex
+ assert isinstance(model.model.layers[0].mlp.gate, nn.Linear)
+ assert not isinstance(model.model.layers[0].mlp.gate, QuantLinear)
+ # experts are quantized
+ assert isinstance(model.model.layers[0].mlp.experts[0].w1, QuantLinear)
+
+
+class TestRegexOverrides:
+ def test_get_qlinear_init_args_with_regex_override(self):
+ overrides = {
+ "re:.*\\.mlp\\.experts\\..*\\.w1": {"bits": 8, "group_size": 32},
+ "model.lm_head": {"bits": 4},
+ }
+ config = OliveHfQuantizationConfig(
+ bits=4,
+ symmetric=True,
+ group_size=128,
+ overrides=overrides,
+ )
+ init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1")
+ assert init_args == {"bits": 8, "symmetric": True, "group_size": 32}
+
+ def test_literal_beats_unrelated_regex(self):
+ overrides = {
+ "re:.*\\.w1": {"bits": 8},
+ "model.layers.0.mlp.experts.0.w1": {"bits": 6},
+ }
+ config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128, overrides=overrides)
+ init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1")
+ # literal (longer) wins
+ assert init_args["bits"] == 6
diff --git a/test/common/quant/test_patterns.py b/test/common/quant/test_patterns.py
new file mode 100644
index 0000000000..1a66287983
--- /dev/null
+++ b/test/common/quant/test_patterns.py
@@ -0,0 +1,83 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import pytest
+
+from olive.common.quant.patterns import is_regex_pattern, match_override, match_skip
+
+
+class TestIsRegexPattern:
+ def test_returns_true_for_re_prefix(self):
+ assert is_regex_pattern("re:foo")
+
+ def test_returns_false_for_plain_string(self):
+ assert not is_regex_pattern("model.embed_tokens")
+
+ def test_returns_false_for_non_string(self):
+ assert not is_regex_pattern(None) # type: ignore[arg-type]
+
+
+class TestMatchOverride:
+ def test_returns_none_when_no_patterns(self):
+ assert match_override("foo", []) is None
+ assert match_override("foo", None) is None
+
+ def test_literal_equality(self):
+ assert match_override("model.embed_tokens", ["model.embed_tokens"]) == "model.embed_tokens"
+ assert match_override("model.embed", ["model.embed_tokens"]) is None
+
+ def test_regex_fullmatch(self):
+ assert match_override("layers.0.experts.gate_proj", ["re:layers\\.\\d+\\.experts\\..*"]) == (
+ "re:layers\\.\\d+\\.experts\\..*"
+ )
+ # not a fullmatch
+ assert match_override("prefix_layers.0.experts", ["re:layers\\.\\d+\\.experts"]) is None
+
+ def test_longest_pattern_wins(self):
+ # both match — the more specific (longer) wins
+ patterns = ["re:.*\\.experts\\..*", "re:layers\\.\\d+\\.experts\\.gate_proj"]
+ match = match_override("layers.0.experts.gate_proj", patterns)
+ assert match == "re:layers\\.\\d+\\.experts\\.gate_proj"
+
+ def test_tie_break_is_deterministic(self):
+ # same-length patterns: lexically smallest wins
+ patterns = ["re:foo.bar", "re:foo.baz"]
+ # only the first matches
+ assert match_override("foo.bar", patterns) == "re:foo.bar"
+
+ def test_literal_beats_shorter_regex_when_longer(self):
+ patterns = ["re:.*", "model.embed_tokens"]
+ # literal pattern (len 18) is longer than "re:.*" (len 5)
+ assert match_override("model.embed_tokens", patterns) == "model.embed_tokens"
+
+
+class TestMatchSkip:
+ def test_substring_match_for_plain_string(self):
+ # substring preserves HF semantics
+ assert match_skip("model.layers.0.experts.gate_proj", ["experts"])
+ assert match_skip("model.embed_tokens", ["embed_tokens"])
+ assert not match_skip("model.layers.0.attn.q_proj", ["experts"])
+
+ def test_regex_fullmatch_for_re_prefix(self):
+ assert match_skip("model.layers.0.experts.router.gate", ["re:.*\\.router\\.gate"])
+ assert not match_skip("model.layers.0.experts.router.gate.bias", ["re:.*\\.router\\.gate"])
+
+ def test_empty_pattern_does_not_match(self):
+ assert not match_skip("foo", [""])
+
+ def test_none_or_empty_patterns(self):
+ assert not match_skip("foo", None)
+ assert not match_skip("foo", [])
+
+ @pytest.mark.parametrize(
+ ("name", "patterns", "expected"),
+ [
+ ("model.layers.0.experts.0.w1", ["experts"], True),
+ ("model.layers.0.attn.q_proj", ["experts"], False),
+ ("router.gate", ["re:router\\.gate"], True),
+ ("router_gate", ["re:router\\.gate"], False),
+ ],
+ )
+ def test_parametrized(self, name, patterns, expected):
+ assert match_skip(name, patterns) is expected
diff --git a/test/common/quant/test_tensor.py b/test/common/quant/test_tensor.py
new file mode 100644
index 0000000000..3f36cf4354
--- /dev/null
+++ b/test/common/quant/test_tensor.py
@@ -0,0 +1,138 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from olive.common.quant.tensor import QuantTensor
+
+
+@pytest.fixture
+def w2d():
+ torch.manual_seed(0)
+ return torch.randn(64, 128, dtype=torch.float32)
+
+
+@pytest.fixture
+def w3d():
+ torch.manual_seed(0)
+ return torch.randn(4, 32, 128, dtype=torch.float32)
+
+
+class TestQuantTensor2D:
+ def test_shape_dtype_device_preserved(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ assert qt.shape == w2d.shape
+ assert qt.dtype == w2d.dtype
+ assert qt.device == w2d.device
+ assert qt.requires_grad is False
+
+ def test_inner_buffer_layout(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ # 4-bit packed → in_features / 2
+ assert qt.qweight.shape == (64, 64)
+ assert qt.qweight.dtype == torch.uint8
+ # groupwise scales: (out, num_groups)
+ assert qt.scales.shape == (64, 128 // 32)
+ # symmetric → no zero_points
+ assert qt.qzeros is None
+
+ def test_asymmetric_has_qzeros(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=False, group_size=32)
+ assert qt.qzeros is not None
+ assert qt.qzeros.dtype == torch.uint8
+
+ def test_to_dense_round_trip(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=False, group_size=32)
+ dense = qt.to_dense()
+ assert dense.shape == w2d.shape
+ # Round trip should be close (4-bit groupwise is reasonably accurate)
+ assert (dense - w2d).abs().mean().item() < 0.1
+
+ def test_dispatches_through_f_linear(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ x = torch.randn(2, 128)
+ out_quant = F.linear(x, qt)
+ out_dense = F.linear(x, qt.to_dense())
+ assert torch.allclose(out_quant, out_dense, atol=1e-5)
+
+ def test_nn_parameter_preserves_subclass(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ p = nn.Parameter(qt, requires_grad=False)
+ assert isinstance(p, QuantTensor)
+ assert isinstance(p.data, QuantTensor)
+
+ def test_nn_linear_forward_with_quant_tensor_weight(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ layer = nn.Linear(128, 64, bias=False)
+ layer.weight = nn.Parameter(qt, requires_grad=False)
+ x = torch.randn(2, 128)
+ out_layer = layer(x)
+ out_ref = F.linear(x, qt.to_dense())
+ assert torch.allclose(out_layer, out_ref, atol=1e-5)
+
+ def test_model_to_dtype_propagates(self, w2d):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=False, group_size=32)
+
+ class M(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.lin = nn.Linear(128, 64, bias=False)
+
+ m = M()
+ m.lin.weight = nn.Parameter(qt, requires_grad=False)
+ m = m.to(torch.float16)
+ # dtype follows the wrapper subclass; scales are floating-point
+ assert m.lin.weight.dtype == torch.float16
+ assert m.lin.weight.scales.dtype == torch.float16
+ # qweight is uint8 — non-floating-point, kept as-is
+ assert m.lin.weight.qweight.dtype == torch.uint8
+
+ def test_nn_embedding_forward(self, w2d):
+ # 64 embeddings of dim 128
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ emb = nn.Embedding(64, 128)
+ emb.weight = nn.Parameter(qt, requires_grad=False)
+ ids = torch.tensor([0, 5, 60])
+ out = emb(ids)
+ out_ref = F.embedding(ids, qt.to_dense())
+ assert torch.allclose(out, out_ref, atol=1e-5)
+
+
+class TestQuantTensor3D:
+ def test_3d_shape(self, w3d):
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ assert qt.shape == w3d.shape
+ assert qt.qweight.shape == (4, 32, 64)
+ assert qt.scales.shape == (4, 32, 4)
+ assert qt.qzeros is not None
+ assert qt.qzeros.shape == (4, 32, 2)
+
+ def test_3d_round_trip(self, w3d):
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ dense = qt.to_dense()
+ assert (dense - w3d).abs().mean().item() < 0.1
+
+ def test_slice_returns_2d_quant_tensor(self, w3d):
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ sliced = qt[2]
+ assert isinstance(sliced, QuantTensor)
+ assert sliced.shape == w3d[2].shape
+ # F.linear over the slice
+ x = torch.randn(2, 128)
+ out = F.linear(x, sliced)
+ out_ref = F.linear(x, qt.to_dense()[2])
+ assert torch.allclose(out, out_ref, atol=1e-5)
+
+
+class TestQuantTensorOnnxExportGuards:
+ def test_linear_raises_when_in_onnx_export(self, w2d, monkeypatch):
+ qt = QuantTensor.from_float(w2d, bits=4, symmetric=True, group_size=32)
+ x = torch.randn(1, 128)
+ # Simulate being inside ONNX export
+ monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True)
+ with pytest.raises(RuntimeError, match="QuantTensor cannot be traced"):
+ F.linear(x, qt)
diff --git a/test/passes/onnx/test_model_builder.py b/test/passes/onnx/test_model_builder.py
index be5b728c65..4f574063ec 100644
--- a/test/passes/onnx/test_model_builder.py
+++ b/test/passes/onnx/test_model_builder.py
@@ -196,3 +196,33 @@ def fake_create_model(
assert str(output_folder / "encoder.onnx.data") not in additional_files
assert str(output_folder / "decoder.onnx.data") not in additional_files
assert str(output_folder / "tokenizer.json") in additional_files
+
+
+def test_olive_quantized_model_raises_for_moe():
+ """ModelBuilder cannot consume Olive-quantized MoE checkpoints; it must error out
+ cleanly so the user reaches for Mobius (or re-runs RTN without moe=True).
+ """
+ from olive.passes.onnx.model_builder import OliveQuantizedModel
+
+ quant_attrs = {
+ "config": {
+ "bits": 4,
+ "group_size": 32,
+ "symmetric": True,
+ "embeds": False,
+ "lm_head": False,
+ "tie_word_embeddings": False,
+ "moe": True,
+ "overrides": {},
+ }
+ }
+ with pytest.raises(NotImplementedError, match="MoE"):
+ OliveQuantizedModel(
+ quant_type="olive",
+ input_path="/tmp/does_not_matter",
+ quant_attrs=quant_attrs,
+ q_size=64,
+ kv_size=64,
+ intermediate_size=64,
+ num_layers=1,
+ )
diff --git a/test/passes/pytorch/test_quant_utils.py b/test/passes/pytorch/test_quant_utils.py
new file mode 100644
index 0000000000..ca3f7a93dd
--- /dev/null
+++ b/test/passes/pytorch/test_quant_utils.py
@@ -0,0 +1,87 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Unit tests for olive.passes.pytorch.quant_utils helpers.
+
+The full ``prepare_model`` / ``finalize`` pipeline is exercised end-to-end
+by ``test_rtn.py`` against a real HF model. This file targets the helpers
+that the MoE path adds — primarily ``flatten_quant_tensor_params`` — and
+keeps the tests free of network access.
+"""
+
+import torch
+from torch import nn
+
+from olive.common.quant.tensor import QuantTensor
+from olive.passes.pytorch.quant_utils import flatten_quant_tensor_params
+
+
+class _ExpertsBlock(nn.Module):
+ """Fake fused-3D experts module (gpt-oss / Qwen3-MoE style).
+
+ Carries a single 3D ``nn.Parameter`` of shape ``(num_experts, out, in)``,
+ matching the layout that ``prepare_model`` annotates with
+ ``quant_info_3d`` and that ``finalize`` rewrites into a 3D
+ :class:`QuantTensor` parameter.
+ """
+
+ def __init__(self, num_experts: int = 4, out_features: int = 32, in_features: int = 16):
+ super().__init__()
+ self.gate_up_proj = nn.Parameter(
+ torch.randn(num_experts, out_features, in_features, dtype=torch.float32),
+ requires_grad=False,
+ )
+
+
+class TestFlattenQuantTensorParams:
+ def test_flatten_3d_quant_tensor_param(self):
+ block = _ExpertsBlock()
+ qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=True, group_size=16)
+ block._parameters["gate_up_proj"] = nn.Parameter(qt, requires_grad=False)
+
+ flatten_quant_tensor_params(block)
+
+ # The original parameter is gone, replaced by buffers.
+ assert "gate_up_proj" not in dict(block.named_parameters())
+ buffer_names = dict(block.named_buffers())
+ assert "gate_up_proj_qweight" in buffer_names
+ assert "gate_up_proj_scales" in buffer_names
+ # symmetric → no qzeros buffer
+ assert "gate_up_proj_qzeros" not in buffer_names
+ # metadata stays as a plain attribute (not in state_dict)
+ meta = block.gate_up_proj_quant_meta
+ assert meta["bits"] == 4
+ assert meta["group_size"] == 16
+ assert meta["symmetric"] is True
+ assert tuple(meta["shape"]) == (4, 32, 16)
+
+ def test_flatten_asymmetric_emits_qzeros(self):
+ block = _ExpertsBlock()
+ qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=False, group_size=16)
+ block._parameters["gate_up_proj"] = nn.Parameter(qt, requires_grad=False)
+
+ flatten_quant_tensor_params(block)
+
+ buffer_names = dict(block.named_buffers())
+ assert "gate_up_proj_qzeros" in buffer_names
+
+ def test_state_dict_serializable_after_flatten(self):
+ """After flattening, the module's state_dict must be plain Tensors."""
+ block = _ExpertsBlock()
+ qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=True, group_size=16)
+ block._parameters["gate_up_proj"] = nn.Parameter(qt, requires_grad=False)
+
+ flatten_quant_tensor_params(block)
+
+ sd = block.state_dict()
+ for key, value in sd.items():
+ assert type(value) is torch.Tensor, f"{key} is not a plain Tensor: {type(value)}"
+
+ def test_no_op_on_module_without_quant_tensor(self):
+ """Plain modules pass through unchanged."""
+ block = nn.Linear(8, 8)
+ before = {n: id(p) for n, p in block.named_parameters()}
+ flatten_quant_tensor_params(block)
+ after = {n: id(p) for n, p in block.named_parameters()}
+ assert before == after
From 7cc2e7ea56a7faab379aae6f7ce411f855b4bbe9 Mon Sep 17 00:00:00 2001
From: Jambay Kinley
Date: Fri, 15 May 2026 23:41:03 +0000
Subject: [PATCH 02/18] MoE quantization v4: QuantTensor + aliased buffers
(Design 3)
Switch Olive's native quantization representation to a single design:
every quantized weight is an nn.Parameter(QuantTensor) on the original
host nn.Linear / nn.Embedding / fused-3D experts module, with sibling
_qweight / _scales / _qzeros buffers aliasing the QuantTensor's
inner tensors.
Save: a state-dict hook drops the QuantTensor parameter entry; the
buffers already carry the data (plain Tensors, safetensors-friendly).
Load: HF's loader fills the buffers natively via dotted paths; a
post-load helper re-binds the QuantTensor inner refs to the freshly
loaded buffer storage.
QuantLinear / QuantEmbedding (olive/common/quant/nn.py) are kept only
as ONNX-exportable wrappers used by make_export_compatible_quant; they
are no longer the runtime representation.
* New olive/common/quant/state_dict.py with install_quant_tensor_param
and refresh_quant_tensor_refs helpers.
* OliveHfQuantizer rewritten for the new layout (placeholder install
before weight load + ref refresh after).
* finalize() in passes/pytorch/quant_utils.py installs QuantTensor
params via install_quant_tensor_param (replaces the old
flatten_quant_tensor_params helper).
* prepare_model skips modules whose weight is already a QuantTensor,
so composing multiple Rtn passes on top of a partially quantized
model works.
* make_export_compatible_quant detects nn.Linear / nn.Embedding whose
weight is a QuantTensor and swaps them with QuantLinear /
QuantEmbedding wrappers before any model dtype casting, preserving
the existing com.microsoft::MatMulNBits /
com.microsoft::GatherBlockQuantized symbolic export path.
* OliveQuantizedModel (model_builder.py) normalizes the new
.weight_qweight key layout back to the legacy
.qweight layout for the existing genai loader, and raises
NotImplementedError for moe=True checkpoints.
* Tests updated to assert against QuantTensor weight instead of
isinstance(module, QuantLinear); legacy tie_quant_modules tests
removed; new install_quant_tensor_param test suite added.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/quant.py | 71 +++++++
olive/common/quant/hf_utils.py | 236 +++++++++++++++---------
olive/common/quant/state_dict.py | 152 +++++++++++++++
olive/common/quant/tensor.py | 3 +-
olive/passes/onnx/model_builder.py | 14 ++
olive/passes/pytorch/quant_utils.py | 156 ++++++----------
test/common/quant/test_hf_utils.py | 99 ++++------
test/passes/pytorch/test_gptq.py | 23 ++-
test/passes/pytorch/test_quant_utils.py | 103 ++++++-----
test/passes/pytorch/test_rtn.py | 32 +++-
10 files changed, 565 insertions(+), 324 deletions(-)
create mode 100644 olive/common/quant/state_dict.py
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index 07d18b318a..06bb0b4c01 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -374,6 +374,12 @@ def make_auto_gptq_qlinearnbit(qlinear, dynamo):
def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch.nn.Module:
"""Make the model export compatible by replacing the quantized linear layers with 4-bit versions."""
+ # Olive-native path: replace nn.Linear / nn.Embedding whose weight is a
+ # ``QuantTensor`` parameter with the ONNX-exportable wrappers in
+ # ``olive.common.quant.nn``. Done first so subsequent passes that
+ # cast the model dtype don't dequantize through QuantTensor.
+ _replace_olive_quant_tensor_modules(model)
+
model, modified = _replace_qlinear_modules(
model, dynamo, EXPORT_QLINEAR_MAPPING, "Making export compatible quantized model"
)
@@ -381,3 +387,68 @@ def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch.
# set quantization method to None, gptq doesn't allow dtype casting
model.quantization_method = None
return model
+
+
+def _replace_olive_quant_tensor_modules(model: torch.nn.Module) -> None:
+ """Swap host ``nn.Linear`` / ``nn.Embedding`` (whose weight is a ``QuantTensor``)
+ with the ONNX-exportable ``QuantLinear`` / ``QuantEmbedding`` wrappers.
+
+ The on-disk buffer layout used by Olive's QuantTensor (packed-uint8
+ along the last dim, per-group scales, optional packed qzeros) is
+ bit-identical to the layout used by ``olive.common.quant.nn``
+ ``QuantLinear`` / ``QuantEmbedding`` — both follow the auto-gptq
+ style. So conversion is a direct buffer copy.
+ """
+ from olive.common.quant.nn import QuantEmbedding, QuantLinear
+ from olive.common.quant.tensor import QuantTensor
+
+ targets: list[tuple[str, torch.nn.Module]] = []
+ for name, module in model.named_modules():
+ if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
+ continue
+ weight = module._parameters.get("weight")
+ if weight is None or not isinstance(weight.data, QuantTensor):
+ continue
+ targets.append((name, module))
+
+ if not targets:
+ return
+
+ for name, module in targets:
+ qt: QuantTensor = module.weight.data # type: ignore[assignment]
+ if isinstance(module, torch.nn.Embedding):
+ new = QuantEmbedding(
+ num_embeddings=module.num_embeddings,
+ embedding_dim=module.embedding_dim,
+ bits=qt.bits,
+ symmetric=qt.symmetric,
+ group_size=qt.group_size,
+ device=qt.qweight.device,
+ dtype=qt.scales.dtype,
+ padding_idx=module.padding_idx,
+ )
+ else:
+ new = QuantLinear(
+ in_features=module.in_features,
+ out_features=module.out_features,
+ bias=module.bias is not None,
+ bits=qt.bits,
+ symmetric=qt.symmetric,
+ group_size=qt.group_size,
+ device=qt.qweight.device,
+ dtype=qt.scales.dtype,
+ )
+ if module.bias is not None:
+ new.bias = module.bias.detach().clone()
+ new.qweight = qt.qweight.detach().clone()
+ new.scales = qt.scales.detach().clone()
+ if qt.qzeros is not None:
+ new.qzeros = qt.qzeros.detach().clone()
+ set_attr(model, name, new)
+
+ # After all replacements, the original ``quantization_method`` /
+ # ``quantization_config`` set by ``OliveHfQuantizer`` no longer
+ # describes the runtime modules. Clear it so downstream passes
+ # (e.g. dtype casts) treat the model as plain torch.
+ if hasattr(model, "quantization_method"):
+ model.quantization_method = None
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index 67d82172f7..b6118ba6ce 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -4,23 +4,26 @@
# --------------------------------------------------------------------------
from __future__ import annotations
+import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable
+import torch
import torch.nn as nn
from transformers.quantizers.base import HfQuantizer
from transformers.utils.quantization_config import QuantizationConfigMixin
from olive.common.hf.wrapper import ModelWrapper
from olive.common.quant.patterns import match_override, match_skip
+from olive.common.quant.state_dict import buffer_names, install_quant_tensor_param, refresh_quant_tensor_refs
+from olive.common.quant.tensor import QuantTensor
+from olive.common.quant.utils import WeightQuantizer
from olive.common.utils import StrEnumBase
if TYPE_CHECKING:
from tqdm.auto import tqdm
from transformers import PreTrainedModel
- from olive.common.quant.nn import QuantModule
-
# older transformers expects a StrEnum and accesses .value
class OliveHfQuantizationMethod(StrEnumBase):
@@ -161,22 +164,32 @@ def sort_key(name: str) -> tuple:
class OliveHfQuantizer(HfQuantizer):
- """Olive quantizer."""
+ """Olive quantizer.
+
+ Layout (Design 3, see ``olive/common/quant/state_dict.py``):
+
+ * Each quantized weight is installed as ``nn.Parameter(QuantTensor)``
+ on the original host module (``nn.Linear``, ``nn.Embedding``, or an
+ experts module that owns a fused-3D parameter).
+ * Sibling buffers ``_qweight`` / ``_scales`` / ``_qzeros``
+ alias the QuantTensor's inner tensors. These are the only things
+ written to safetensors; HF's loader fills them via normal dotted
+ paths. After load we re-bind the QuantTensor inner refs to point
+ at the freshly-loaded buffer storage.
+ """
# only support load and inference, no on-the-fly quantization
requires_calibration = True
- modules_to_not_convert: list[str] | None = None
def _process_model_before_weight_loading(
self, model: PreTrainedModel, keep_in_fp32_modules: list[str] | None = None, **kwargs
):
- from olive.common.quant.nn import QuantEmbedding, QuantLinear
-
ids_to_skip: list[int] = []
if not self.quantization_config.lm_head:
ids_to_skip.append(id(model.get_output_embeddings()))
if not self.quantization_config.embeds:
ids_to_skip.append(id(model.get_input_embeddings()))
+ expert_modules: list[tuple[nn.Module, str]] = []
if not self.quantization_config.moe:
# Skip every nn.Module under each experts subtree — this both
# leaves fused-3D experts alone *and* fixes the previous silent
@@ -194,9 +207,22 @@ def _process_model_before_weight_loading(
# Not every model is wrappable (e.g., random tests). Falling
# back to the previous behaviour (no experts skip) is safe.
pass
+ else:
+ # collect (experts_module, dotted_name) so we can also install
+ # 3D placeholders on fused-experts modules below.
+ try:
+ wrapper = ModelWrapper.from_model(model)
+ module_to_name = {id(m): n for n, m in model.named_modules()}
+ for lw in wrapper.get_layer_wrappers():
+ experts = lw.get_experts(return_name=False)
+ if experts is None:
+ continue
+ expert_modules.append((experts, module_to_name.get(id(experts), "")))
+ except Exception: # pylint: disable=broad-except
+ pass
- modules_to_not_convert: list[str] = (
- [name for name, module in model.named_modules() if id(module) in ids_to_skip] if ids_to_skip else []
+ skip_literal_names: set[str] = (
+ {name for name, module in model.named_modules() if id(module) in ids_to_skip} if ids_to_skip else set()
)
# Pattern-aware skip list. Plain strings keep their HF substring
# semantics; ``re:`` opts into regex fullmatch.
@@ -205,47 +231,60 @@ def _process_model_before_weight_loading(
skip_patterns.extend(self.quantization_config.modules_to_not_convert)
if keep_in_fp32_modules:
skip_patterns.extend(keep_in_fp32_modules)
- self.modules_to_not_convert = modules_to_not_convert
+ self.modules_to_not_convert = sorted(skip_literal_names)
self._skip_patterns = skip_patterns
- def should_quantize(module: nn.Module, name: str) -> bool:
- """Check if a module should be quantized."""
+ def _should_quantize(module: nn.Module, name: str) -> bool:
if not isinstance(module, (nn.Linear, nn.Embedding)):
return False
- if any(literal == name for literal in self.modules_to_not_convert):
+ if name in skip_literal_names:
return False
- return not match_skip(name, self._skip_patterns)
-
- def create_quantized_module(module: nn.Linear | nn.Embedding, name: str) -> QuantLinear | QuantEmbedding:
- """Create a quantized version of a module."""
- common_kwargs = {
- **self.quantization_config.get_qlinear_init_args(name),
- "device": module.weight.device,
- "dtype": module.weight.dtype,
- }
- if isinstance(module, nn.Embedding):
- return QuantEmbedding(
- num_embeddings=module.num_embeddings,
- embedding_dim=module.embedding_dim,
- padding_idx=module.padding_idx,
- **common_kwargs,
- )
- return QuantLinear(
- in_features=module.in_features,
- out_features=module.out_features,
- bias=module.bias is not None,
- **common_kwargs,
+ return not match_skip(name, skip_patterns)
+
+ # 2D placeholders: nn.Linear / nn.Embedding
+ for name, module in list(model.named_modules()):
+ if not _should_quantize(module, name):
+ continue
+ qargs = self.quantization_config.get_qlinear_init_args(name)
+ qt = _build_placeholder_quant_tensor(
+ shape=tuple(module.weight.shape),
+ bits=qargs["bits"],
+ symmetric=qargs["symmetric"],
+ group_size=qargs["group_size"],
+ dtype=module.weight.dtype,
+ device=module.weight.device,
)
-
- replace_matching_submodules(model, should_quantize, create_quantized_module)
+ install_quant_tensor_param(module, "weight", qt)
+
+ # 3D placeholders: fused-experts modules
+ for experts_module, experts_name in expert_modules:
+ for pname, param in list(experts_module._parameters.items()):
+ if param is None or param.dim() != 3 or isinstance(param.data, QuantTensor):
+ continue
+ full_name = f"{experts_name}.{pname}" if experts_name else pname
+ if full_name in skip_literal_names or match_skip(full_name, skip_patterns):
+ continue
+ qargs = self.quantization_config.get_qlinear_init_args(full_name)
+ qt = _build_placeholder_quant_tensor(
+ shape=tuple(param.shape),
+ bits=qargs["bits"],
+ symmetric=qargs["symmetric"],
+ group_size=qargs["group_size"],
+ dtype=param.dtype,
+ device=param.device,
+ )
+ install_quant_tensor_param(experts_module, pname, qt)
if self.quantization_config.tie_word_embeddings:
# doing first time so that the weight load doesn't complain about missing weights
tie_quant_word_embeddings(model)
def _process_model_after_weight_loading(self, model: PreTrainedModel, **kwargs):
+ # HF's loader assigns freshly-loaded buffer tensors in place, so
+ # re-bind every QuantTensor parameter to point at the current
+ # ``_qweight`` / ``_scales`` / ``_qzeros`` buffer storages.
+ refresh_quant_tensor_refs(model)
if self.quantization_config.tie_word_embeddings:
- # doing again to ensure buffers are tied after loading weights
tie_quant_word_embeddings(model)
return model
@@ -254,11 +293,59 @@ def is_serializable(self, safe_serialization=None) -> bool:
@property
def is_trainable(self) -> bool:
- # TODO(jambayk): investigate what this means (peft, scale+bias, etc.?)
- # need to support peft, scale+bias comes for free since everything is in torch
return False
+def _build_placeholder_quant_tensor(
+ *,
+ shape: tuple[int, ...],
+ bits: int,
+ symmetric: bool,
+ group_size: int,
+ dtype: torch.dtype,
+ device: torch.device,
+) -> QuantTensor:
+ """Build a zero-filled ``QuantTensor`` with the correct buffer shapes.
+
+ This is used as a placeholder before ``from_pretrained`` fills in the
+ real values; the buffer shapes must match what was written at save
+ time so HF's loader assigns into them without complaint.
+ """
+ quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False)
+ packing_factor = 8 // bits
+ qparam_shape = quantizer.get_qparam_shape(shape)
+
+ if len(shape) == 2:
+ qweight_shape = (shape[0], math.ceil(shape[1] / packing_factor))
+ elif len(shape) == 3:
+ qweight_shape = (shape[0], shape[1], math.ceil(shape[2] / packing_factor))
+ else:
+ raise ValueError(f"QuantTensor placeholder only supports 2D/3D shapes, got {shape}")
+
+ qweight = torch.zeros(qweight_shape, dtype=torch.uint8, device=device)
+ scales = torch.zeros(qparam_shape, dtype=dtype, device=device)
+ qzeros: torch.Tensor | None
+ if symmetric:
+ qzeros = None
+ else:
+ if len(qparam_shape) == 2:
+ qz_shape = (qparam_shape[0], math.ceil(qparam_shape[1] / packing_factor))
+ else:
+ qz_shape = (qparam_shape[0], qparam_shape[1], math.ceil(qparam_shape[2] / packing_factor))
+ qzeros = torch.zeros(qz_shape, dtype=torch.uint8, device=device)
+
+ return QuantTensor.from_packed(
+ qweight=qweight,
+ scales=scales,
+ qzeros=qzeros,
+ bits=bits,
+ group_size=group_size,
+ symmetric=symmetric,
+ shape=shape,
+ dtype=dtype,
+ )
+
+
def replace_matching_submodules(
module: nn.Module,
condition: Callable[[nn.Module, str], bool],
@@ -309,56 +396,33 @@ def replace_matching_submodules(
return module
-def _repoint_buffer(src: nn.Module, dst: nn.Module, name: str):
- """Repoint a buffer from src to dst.
-
- Args:
- src: Source module.
- dst: Destination module.
- name: Name of the buffer to repoint.
+def tie_quant_word_embeddings(model: PreTrainedModel) -> None:
+ """Tie the input and output embeddings when they share a quantized weight.
+ Both modules' ``weight`` ``nn.Parameter`` is set to the **same**
+ ``nn.Parameter(QuantTensor)`` object, and the underlying
+ ``weight_qweight`` / ``weight_scales`` / ``weight_qzeros`` buffers
+ are tied (aliased to the input embedding's buffers). This preserves
+ the standard HF tied-weights semantics for the quantized layout.
"""
- src_buf = getattr(src, name, None)
-
- # ensure both are None or both exist
- if src_buf is None:
- assert getattr(dst, name, None) is None, f"Output embedding has {name} but input does not."
+ src = model.get_input_embeddings()
+ dst = model.get_output_embeddings()
+ if src is None or dst is None:
return
- # ensure both have the buffer shapes and types match
- dst_buf = getattr(dst, name, None)
- assert src_buf.shape == dst_buf.shape, (
- f"Cannot tie embeddings: input embedding {name} shape {src_buf.shape} "
- f"does not match output embedding shape {dst_buf.shape}."
- )
- assert src_buf.dtype == dst_buf.dtype, (
- f"Cannot tie embeddings: input embedding {name} dtype {src_buf.dtype} "
- f"does not match output embedding dtype {dst_buf.dtype}."
- )
-
- # tie the buffers
- # pylint: disable=W0212
- dst._buffers[name] = src_buf
- dst._non_persistent_buffers_set.add(name)
-
-
-def tie_quant_modules(src: QuantModule, dst: QuantModule):
- """Tie the quantization buffers of two QuantModules.
-
- Args:
- src: Source QuantModule.
- dst: Destination QuantModule.
-
- """
- for name in ["qweight", "scales", "qzeros"]:
- _repoint_buffer(src, dst, name)
-
-
-def tie_quant_word_embeddings(model: PreTrainedModel):
- """Tie the word embeddings and output embeddings if they have the same shape.
-
- Args:
- model: The HuggingFace model to tie embeddings for.
+ # The input embedding owns the canonical QuantTensor + buffers.
+ src_param = src._parameters.get("weight")
+ if src_param is None or not isinstance(src_param.data, QuantTensor):
+ return
- """
- tie_quant_modules(model.get_input_embeddings(), model.get_output_embeddings())
+ qname, sname, zname = buffer_names("weight")
+ # tie buffers
+ for n in (qname, sname, zname):
+ src_buf = src._buffers.get(n)
+ if src_buf is None:
+ continue
+ dst._buffers[n] = src_buf
+
+ # tie the QuantTensor parameter itself (same Python Parameter object,
+ # so both modules see the same .data and the same inner tensors).
+ dst._parameters["weight"] = src_param
diff --git a/olive/common/quant/state_dict.py b/olive/common/quant/state_dict.py
new file mode 100644
index 0000000000..fb09f6d9d2
--- /dev/null
+++ b/olive/common/quant/state_dict.py
@@ -0,0 +1,152 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""State-dict helpers for Olive's quantized weight representation.
+
+Olive's quantization layout (Design 3):
+
+* The quantized state of a weight named ```` (typically
+ ``"weight"`` for ``nn.Linear``/``nn.Embedding``, or e.g.
+ ``"gate_up_proj"`` for a fused-3D MoE expert tensor) is stored as
+ plain buffers on the host module:
+
+ * ``_qweight`` - packed uint8 tensor (always present)
+ * ``_scales`` - per-group scales (always present)
+ * ``_qzeros`` - per-group zero points (asymmetric only)
+
+ This matches the suffix convention so on-disk safetensors keys are
+ HF-loader friendly: ``model.layers.0.mlp.gate_proj.weight_qweight``,
+ ``...experts.gate_up_proj_qweight``, etc.
+
+* At runtime, the host module's original parameter
+ ``module._parameters[pname]`` becomes
+ ``nn.Parameter(QuantTensor(...), requires_grad=False)``. The
+ ``QuantTensor``'s inner ``qweight``/``scales``/``qzeros`` references
+ alias the same Python tensor objects as the buffers above, so
+ ``module.`` is a live view over the buffers and the original
+ forward (e.g. ``F.linear``) dispatches through
+ ``QuantTensor.__torch_function__``.
+
+To keep save/load simple we only need a single state-dict save hook
+(to suppress the QuantTensor parameter entry, since the buffers already
+carry the data) plus a post-load helper that refreshes the QuantTensor
+inner references after HF assigns freshly-loaded buffer tensors.
+"""
+
+from __future__ import annotations
+
+import torch
+
+_INSTALLED_FLAG = "_olive_quant_state_dict_hook_installed"
+
+QWEIGHT_SUFFIX = "_qweight"
+SCALES_SUFFIX = "_scales"
+QZEROS_SUFFIX = "_qzeros"
+
+
+def buffer_names(pname: str) -> tuple[str, str, str]:
+ """Return the ``(qweight, scales, qzeros)`` buffer names for parameter ``pname``."""
+ return f"{pname}{QWEIGHT_SUFFIX}", f"{pname}{SCALES_SUFFIX}", f"{pname}{QZEROS_SUFFIX}"
+
+
+def _save_hook(module: torch.nn.Module, state_dict: dict, prefix: str, local_metadata: dict) -> None:
+ """Drop ``QuantTensor`` parameter entries from ``state_dict``.
+
+ Inner ``qweight``/``scales``/``qzeros`` tensors are already exposed
+ as plain buffers on ``module`` and therefore appear in ``state_dict``
+ under their own keys; the QuantTensor parameter entry is a redundant
+ (and non-serialisable) duplicate.
+ """
+ # Local import to avoid a circular dependency at module-import time.
+ from olive.common.quant.tensor import QuantTensor
+
+ for pname in list(module._parameters):
+ full_key = f"{prefix}{pname}"
+ value = state_dict.get(full_key)
+ if isinstance(value, QuantTensor):
+ del state_dict[full_key]
+
+
+def install_state_dict_hooks(module: torch.nn.Module) -> None:
+ """Install Olive's state-dict save hook on ``module`` (idempotent)."""
+ if getattr(module, _INSTALLED_FLAG, False):
+ return
+ module._register_state_dict_hook(_save_hook)
+ setattr(module, _INSTALLED_FLAG, True)
+
+
+def install_quant_tensor_param(
+ module: torch.nn.Module,
+ pname: str,
+ qt, # QuantTensor
+) -> None:
+ """Install ``qt`` on ``module`` as ```` plus aliased sibling buffers.
+
+ Replaces ``module._parameters[pname]`` with ``nn.Parameter(qt)`` and
+ registers ``_qweight``/``_scales``/(optionally) ``_qzeros``
+ as buffers whose storage is the same as the QuantTensor's inner
+ tensors. The state-dict save hook is installed as a side effect.
+ """
+ from olive.common.quant.tensor import QuantTensor
+
+ if not isinstance(qt, QuantTensor):
+ raise TypeError(f"Expected QuantTensor, got {type(qt).__name__}")
+
+ qname, sname, zname = buffer_names(pname)
+
+ # Detach existing buffers/parameters with the same names first so
+ # ``register_buffer`` / parameter assignment is idempotent.
+ for n in (qname, sname, zname):
+ if n in module._buffers:
+ del module._buffers[n]
+
+ # ``nn.Parameter(qt, requires_grad=False)`` for a tensor subclass
+ # returns the underlying QuantTensor instance directly (after a
+ # ``detach()`` that goes through ``_apply_fn_to_data`` and produces
+ # view-aliased inner tensors). See ``torch.nn.Parameter.__new__`` —
+ # for a tensor subclass it returns ``data.detach().requires_grad_(...)``
+ # which is the QuantTensor itself, not a wrapping Parameter. We
+ # alias the host module's buffers to *that* instance's inner tensors
+ # so save / refresh paths consistently read the same storage.
+ param = torch.nn.Parameter(qt, requires_grad=False)
+ module._parameters[pname] = param
+
+ module.register_buffer(qname, param.qweight, persistent=True)
+ module.register_buffer(sname, param.scales, persistent=True)
+ if param.qzeros is not None:
+ module.register_buffer(zname, param.qzeros, persistent=True)
+
+ install_state_dict_hooks(module)
+
+
+def refresh_quant_tensor_refs(module: torch.nn.Module) -> None:
+ """Re-point each ``QuantTensor`` parameter at the module's current buffers.
+
+ HF's loader assigns freshly-loaded buffer tensors via
+ ``module.load_state_dict({name: tensor}, assign=True)``, which
+ replaces the buffer object in ``module._buffers``. Any
+ ``QuantTensor`` parameter installed earlier would still reference
+ the old (placeholder) storage, so we walk the parameters and re-bind
+ each one's inner tensors to the current buffers.
+
+ This is a no-op for modules with no QuantTensor parameters.
+ """
+ from olive.common.quant.tensor import QuantTensor
+
+ for sub_module in module.modules():
+ for pname, param in list(sub_module._parameters.items()):
+ # ``param`` itself is the QuantTensor instance stored on the
+ # module (``nn.Parameter(qt)`` for a tensor subclass returns
+ # the underlying QuantTensor — see ``torch.nn.Parameter.__new__``).
+ if param is None or not isinstance(param, QuantTensor):
+ continue
+ qname, sname, zname = buffer_names(pname)
+ qweight = sub_module._buffers.get(qname)
+ scales = sub_module._buffers.get(sname)
+ qzeros = sub_module._buffers.get(zname)
+ if qweight is None or scales is None:
+ continue
+ param.qweight = qweight
+ param.scales = scales
+ param.qzeros = qzeros
diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py
index 7cb181a7a5..b25c788cfd 100644
--- a/olive/common/quant/tensor.py
+++ b/olive/common/quant/tensor.py
@@ -241,6 +241,7 @@ def from_packed(
group_size: int,
symmetric: bool,
shape: tuple[int, ...],
+ dtype: torch.dtype | None = None,
) -> QuantTensor:
"""Reconstruct a ``QuantTensor`` from already-packed buffers."""
return cls(
@@ -251,7 +252,7 @@ def from_packed(
group_size=group_size,
symmetric=symmetric,
shape=shape,
- dtype=scales.dtype,
+ dtype=dtype if dtype is not None else scales.dtype,
)
# ------------------------------------------------------------------
diff --git a/olive/passes/onnx/model_builder.py b/olive/passes/onnx/model_builder.py
index 6f86201bd2..35a4f7d5cd 100644
--- a/olive/passes/onnx/model_builder.py
+++ b/olive/passes/onnx/model_builder.py
@@ -505,6 +505,20 @@ def set_tensor(module, tensor_name, tensor_value, local_bits, local_group_size):
if weight_file.suffix == ".safetensors":
weights = load_file(weight_file)
+ # Normalize Olive v4 key layout: ``.weight_qweight``
+ # (and ``_scales`` / ``_qzeros``) → ``.qweight`` (etc.)
+ # so the existing dotted-path resolver keeps working.
+ normalized: dict[str, "torch.Tensor"] = {}
+ for k, v in weights.items():
+ for suffix in ("_qweight", "_scales", "_qzeros"):
+ legacy_suffix = "." + suffix.lstrip("_")
+ new_suffix = ".weight" + suffix
+ if k.endswith(new_suffix):
+ k = k[: -len(new_suffix)] + legacy_suffix
+ break
+ normalized[k] = v
+ weights = normalized
+
# Map weights to modules
for name, tensor in weights.items():
if name.endswith("inv_freq"):
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index b4c112708a..d026144fc6 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -17,8 +17,8 @@
replace_matching_submodules,
tie_quant_word_embeddings,
)
-from olive.common.quant.nn import QuantEmbedding, QuantLinear
from olive.common.quant.patterns import match_skip
+from olive.common.quant.state_dict import install_quant_tensor_param
from olive.common.quant.tensor import QuantTensor
from olive.common.quant.utils import WeightQuantizer
from olive.common.utils import tensor_data_to_device
@@ -176,6 +176,11 @@ def prepare_model(
def should_quantize(module: torch.nn.Module, name: str) -> bool:
if module in excluded_attn_inputs:
return False
+ # already-quantized (QuantTensor weight) — leave it alone so the user
+ # can compose passes on top of a partially-quantized model.
+ weight = getattr(module, "weight", None)
+ if weight is not None and isinstance(weight, QuantTensor):
+ return False
if match_skip(name, skip_patterns):
return False
# category-flag skips (lm_head / embeds / moe) — first rule wins
@@ -497,68 +502,56 @@ def finalize(
device: str,
retie_word_embeddings: bool = False,
) -> HfModelHandler:
- """Finalize quantization by replacing linear and embedding layers with their quantized counterparts.
-
- Args:
- model: The HuggingFace model to finalize.
- output_model_path: Path to save the finalized quantized model.
- wrapper: ModelWrapper containing the model to finalize.
- quant_config: Quantization configuration to use.
- device: Device to perform quantization on.
- retie_word_embeddings: Whether to retie word embeddings if they were originally tied and have compatible quantization.
-
- Returns:
- HfModelHandler with the finalized quantized model.
-
+ """Finalize quantization by installing ``QuantTensor`` parameters on the host modules.
+
+ For every module marked with ``quant_info`` (2D linear / embedding) and
+ every fused-3D MoE experts module marked with ``quant_info_3d``, build a
+ ``QuantTensor`` from the float weight + computed qparams and install it
+ via :func:`install_quant_tensor_param` so that:
+
+ * ``module.`` is an ``nn.Parameter(QuantTensor)`` whose dispatch
+ still drives the original eager forward;
+ * ``module._qweight`` / ``_scales`` / ``_qzeros`` are plain
+ buffers (aliasing the QuantTensor's inner tensors), so the model
+ saves cleanly via ``save_pretrained`` / safetensors with **no
+ tensor subclass on disk**.
"""
-
- def should_quantize(module: torch.nn.Module, _: str) -> bool:
- return hasattr(module, "quant_info")
-
- def quantize_and_pack(module: torch.nn.Module, _: str) -> QuantLinear | QuantEmbedding:
- module.to(device)
- quant_cls = QuantEmbedding if isinstance(module, torch.nn.Embedding) else QuantLinear
- return quant_cls.from_module(
- module.to(device),
- bits=module.quant_info.quantizer.bits,
- symmetric=module.quant_info.quantizer.symmetric,
- group_size=module.quant_info.quantizer.group_size,
- scales=module.quant_info.scales,
- zero_points=module.quant_info.zero_points,
- ).to("cpu") # move the original module to CPU
-
- replace_matching_submodules(
- wrapper.model,
- should_quantize,
- quantize_and_pack,
- description="Quantizing and packing linear layers",
- )
-
- # Fused-3D MoE: experts modules carry a per-parameter ``quant_info_3d``
- # dict. Replace each parameter with a 3D ``QuantTensor`` parameter.
- # The host module's forward continues to find the parameter at the
- # same attribute name; eager forwards may go through QuantTensor's
- # ``__getitem__`` / matmul fallbacks. ONNX export of fused-3D MoE is
- # deferred to Mobius — Olive only persists the checkpoint.
- for sub in wrapper.model.modules():
- info_3d: dict = getattr(sub, "quant_info_3d", None)
- if not info_3d:
- continue
- for pname, qinfo in info_3d.items():
- param = getattr(sub, pname)
- if not isinstance(param, torch.nn.Parameter):
- continue
+ for sub_module in wrapper.model.modules():
+ # ---- 2D path: nn.Linear / nn.Embedding ----------------------
+ qinfo = getattr(sub_module, "quant_info", None)
+ if qinfo is not None and isinstance(sub_module, (torch.nn.Linear, torch.nn.Embedding)):
+ sub_module.to(device)
+ quantizer = qinfo.quantizer
with torch.no_grad():
- quantizer = qinfo.quantizer
qt = QuantTensor.from_float(
- param.detach().to(device).to(param.dtype),
+ sub_module.weight.detach(),
bits=quantizer.bits,
symmetric=quantizer.symmetric,
group_size=quantizer.group_size,
+ scales=qinfo.scales,
+ zero_points=qinfo.zero_points,
).to("cpu")
- sub._parameters[pname] = torch.nn.Parameter(qt, requires_grad=False) # pylint: disable=W0212
- # Tidy up the marker so downstream save/load doesn't trip over it.
- delattr(sub, "quant_info_3d")
+ sub_module.to("cpu")
+ install_quant_tensor_param(sub_module, "weight", qt)
+ del sub_module.quant_info
+
+ # ---- fused-3D MoE path -------------------------------------
+ info_3d: dict | None = getattr(sub_module, "quant_info_3d", None)
+ if info_3d:
+ for pname, qinfo_3d in info_3d.items():
+ param = sub_module._parameters.get(pname)
+ if param is None:
+ continue
+ quantizer = qinfo_3d.quantizer
+ with torch.no_grad():
+ qt = QuantTensor.from_float(
+ param.detach().to(device),
+ bits=quantizer.bits,
+ symmetric=quantizer.symmetric,
+ group_size=quantizer.group_size,
+ ).to("cpu")
+ install_quant_tensor_param(sub_module, pname, qt)
+ del sub_module.quant_info_3d
if retie_word_embeddings:
tie_quant_word_embeddings(wrapper.model)
@@ -567,55 +560,10 @@ def quantize_and_pack(module: torch.nn.Module, _: str) -> QuantLinear | QuantEmb
wrapper.model.quantization_method = quant_config.quant_method
wrapper.model.config.quantization_config = quant_config
- # Flatten any QuantTensor parameters into plain ``register_buffer`` entries so
- # ``save_pretrained`` (which writes safetensors) can serialize them.
- # Naming convention: ``_qweight`` / ``_scales`` / ``_qzeros``
- # sit on the same parent module as the original parameter. This matches
- # the suffix-style layout used by other prequantized formats and lets
- # downstream loaders (Mobius) discover the buffers without a registry.
- flatten_quant_tensor_params(wrapper.model)
-
- # save the quantized model
+ # save the quantized model — state_dict hooks drop QuantTensor entries;
+ # only plain ``_qweight`` / ``_scales`` / ``_qzeros`` buffers
+ # are written to safetensors.
wrapper.model.save_pretrained(output_model_path)
model.save_metadata(output_model_path)
return inherit_hf_from_hf(model, output_model_path, adapter_path=model.adapter_path)
-
-
-def flatten_quant_tensor_params(module: torch.nn.Module) -> None:
- """Replace every ``QuantTensor`` ``nn.Parameter`` with plain registered buffers.
-
- For each parameter ```` whose data is a :class:`QuantTensor`, the
- parameter is deleted and the inner tensors are re-attached as
- ``_qweight``, ``_scales`` (and ``_qzeros`` when
- asymmetric) on the same parent module. Additional metadata
- (``bits``, ``group_size``, ``symmetric``, ``shape``) is stored on
- ``_quant_meta`` for round-trip loading.
-
- This is used at save time so safetensors / pickle never sees a
- tensor subclass.
- """
- for sub in module.modules():
- names = [n for n, p in sub.named_parameters(recurse=False) if isinstance(p.data, QuantTensor)]
- for name in names:
- qt: QuantTensor = sub._parameters[name].data # type: ignore[assignment]
- del sub._parameters[name]
- sub.register_buffer(f"{name}_qweight", qt.qweight.detach().clone())
- sub.register_buffer(f"{name}_scales", qt.scales.detach().clone())
- if qt.qzeros is not None:
- sub.register_buffer(f"{name}_qzeros", qt.qzeros.detach().clone())
- # ``register_buffer`` requires Tensors, so the metadata lives
- # outside the state_dict as a plain attribute. It is consumed
- # by ``OliveHfQuantizer`` at load time via the
- # ``quantization_config`` for shape/bits and is also encoded
- # in the parameter's full name + scales shape if needed.
- setattr(
- sub,
- f"{name}_quant_meta",
- {
- "bits": qt.bits,
- "group_size": qt.group_size,
- "symmetric": qt.symmetric,
- "shape": tuple(qt.shape),
- },
- )
diff --git a/test/common/quant/test_hf_utils.py b/test/common/quant/test_hf_utils.py
index c31fb3fe5f..f54b3faa2b 100644
--- a/test/common/quant/test_hf_utils.py
+++ b/test/common/quant/test_hf_utils.py
@@ -3,7 +3,6 @@
# Licensed under the MIT License.
# --------------------------------------------------------------------------
import pytest
-import torch
import torch.nn as nn
from transformers import PretrainedConfig
from transformers.modeling_utils import PreTrainedModel
@@ -14,14 +13,21 @@
OliveHfQuantizationOverrideConfig,
OliveHfQuantizer,
replace_matching_submodules,
- tie_quant_modules,
tie_quant_word_embeddings,
)
-from olive.common.quant.nn import QuantEmbedding, QuantLinear
+from olive.common.quant.tensor import QuantTensor
# pylint: disable=W0212
+def _is_olive_quant(module: nn.Module) -> bool:
+ """Return True if ``module`` is an nn.Linear/nn.Embedding whose weight is a QuantTensor."""
+ if not isinstance(module, (nn.Linear, nn.Embedding)):
+ return False
+ weight = module._parameters.get("weight")
+ return weight is not None and isinstance(weight.data, QuantTensor)
+
+
class TestOliveHfQuantizationMethod:
def test_enum_value(self):
"""Test that the enum has the expected value."""
@@ -209,13 +215,13 @@ def test_process_model_before_weight_loading(self):
# Process model
quantizer._process_model_before_weight_loading(model)
- # Check that linear layers are replaced with QuantLinear
- assert isinstance(model.linear1, QuantLinear)
- assert isinstance(model.linear2, QuantLinear)
+ # Check that linear layers are replaced with QuantTensor weights
+ assert _is_olive_quant(model.linear1)
+ assert _is_olive_quant(model.linear2)
# Check that embeddings and lm_head are NOT quantized by default
- assert not isinstance(model.embed, QuantEmbedding)
- assert not isinstance(model.lm_head, QuantLinear)
+ assert not _is_olive_quant(model.embed)
+ assert not _is_olive_quant(model.lm_head)
def test_process_model_with_embeds_quantization(self):
"""Test that embeddings are quantized when embeds=True."""
@@ -229,7 +235,7 @@ def test_process_model_with_embeds_quantization(self):
quantizer._process_model_before_weight_loading(model)
# Check that embeddings are quantized
- assert isinstance(model.embed, QuantEmbedding)
+ assert _is_olive_quant(model.embed)
def test_process_model_with_lm_head_quantization(self):
"""Test that lm_head is quantized when lm_head=True."""
@@ -243,7 +249,7 @@ def test_process_model_with_lm_head_quantization(self):
quantizer._process_model_before_weight_loading(model)
# Check that lm_head is quantized
- assert isinstance(model.lm_head, QuantLinear)
+ assert _is_olive_quant(model.lm_head)
def test_process_model_with_modules_to_not_convert(self):
"""Test that specified modules are not converted."""
@@ -263,8 +269,9 @@ def test_process_model_with_modules_to_not_convert(self):
# Check that linear1 is NOT quantized
assert isinstance(model.linear1, nn.Linear)
+ assert not _is_olive_quant(model.linear1)
# But linear2 should be quantized
- assert isinstance(model.linear2, QuantLinear)
+ assert _is_olive_quant(model.linear2)
def test_is_serializable(self):
"""Test that quantizer is serializable."""
@@ -357,49 +364,6 @@ def transform(m, name):
assert isinstance(result.sub[2], nn.Identity)
-class TestTieQuantModules:
- def test_tie_quant_linear_modules(self):
- """Test tying two QuantLinear modules."""
- # Create two QuantLinear modules
- qlinear1 = QuantLinear(32, 10, bits=4, symmetric=True, group_size=16)
- qlinear2 = QuantLinear(32, 10, bits=4, symmetric=True, group_size=16)
-
- # Set some values in qlinear1
- qlinear1.qweight.fill_(1)
- qlinear1.scales.fill_(0.5)
-
- # Initially different
- assert not torch.all(qlinear1.qweight == qlinear2.qweight)
-
- # Tie them
- tie_quant_modules(qlinear1, qlinear2)
-
- # Now they should share buffers
- assert qlinear1.qweight is qlinear2.qweight
- assert qlinear1.scales is qlinear2.scales
-
- # Modifications to one should affect the other
- qlinear1.qweight.fill_(2)
- assert torch.all(qlinear2.qweight == 2)
-
- def test_tie_quant_embedding_modules(self):
- """Test tying two QuantEmbedding modules."""
- # Create two QuantEmbedding modules
- qembed1 = QuantEmbedding(100, 64, bits=4, symmetric=True, group_size=16)
- qembed2 = QuantEmbedding(100, 64, bits=4, symmetric=True, group_size=16)
-
- # Set some values in qembed1
- qembed1.qweight.fill_(1)
- qembed1.scales.fill_(0.5)
-
- # Tie them
- tie_quant_modules(qembed1, qembed2)
-
- # Now they should share buffers
- assert qembed1.qweight is qembed2.qweight
- assert qembed1.scales is qembed2.scales
-
-
class TestTieQuantWordEmbeddings:
def test_tie_word_embeddings(self):
"""Test tying word embeddings in a model."""
@@ -419,16 +383,17 @@ def test_tie_word_embeddings(self):
# Process model to quantize embeddings
quantizer._process_model_before_weight_loading(model)
- # Set different values
- model.embed.qweight.fill_(1)
- model.lm_head.qweight.fill_(2)
+ # Set different values on the aliased buffers
+ model.embed.weight_qweight.fill_(1)
+ model.lm_head.weight_qweight.fill_(2)
# Tie embeddings
tie_quant_word_embeddings(model)
- # Now they should share buffers
- assert model.embed.qweight is model.lm_head.qweight
- assert model.embed.scales is model.lm_head.scales
+ # Now they should share the underlying buffers and the Parameter
+ assert model.embed.weight_qweight is model.lm_head.weight_qweight
+ assert model.embed.weight_scales is model.lm_head.weight_scales
+ assert model.embed._parameters["weight"] is model.lm_head._parameters["weight"]
# -- MoE quantization fixtures and tests --------------------------------------
@@ -509,12 +474,12 @@ def test_module_list_experts_skipped_by_default(self):
for expert in model.model.layers[0].mlp.experts:
assert isinstance(expert.w1, nn.Linear)
- assert not isinstance(expert.w1, QuantLinear)
+ assert not _is_olive_quant(expert.w1)
assert isinstance(expert.w2, nn.Linear)
- assert not isinstance(expert.w2, QuantLinear)
+ assert not _is_olive_quant(expert.w2)
# The router (gate) is also under the mlp but not under .experts;
# it should still be quantized.
- assert isinstance(model.model.layers[0].mlp.gate, QuantLinear)
+ assert _is_olive_quant(model.model.layers[0].mlp.gate)
def test_module_list_experts_quantized_when_moe_true(self):
config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=16, moe=True)
@@ -524,8 +489,8 @@ def test_module_list_experts_quantized_when_moe_true(self):
quantizer._process_model_before_weight_loading(model)
for expert in model.model.layers[0].mlp.experts:
- assert isinstance(expert.w1, QuantLinear)
- assert isinstance(expert.w2, QuantLinear)
+ assert _is_olive_quant(expert.w1)
+ assert _is_olive_quant(expert.w2)
def test_moe_default_is_false(self):
"""``moe`` defaults to False — opt-in, matching lm_head / embeds."""
@@ -549,9 +514,9 @@ def test_regex_skip_pattern(self):
# gate is skipped via regex
assert isinstance(model.model.layers[0].mlp.gate, nn.Linear)
- assert not isinstance(model.model.layers[0].mlp.gate, QuantLinear)
+ assert not _is_olive_quant(model.model.layers[0].mlp.gate)
# experts are quantized
- assert isinstance(model.model.layers[0].mlp.experts[0].w1, QuantLinear)
+ assert _is_olive_quant(model.model.layers[0].mlp.experts[0].w1)
class TestRegexOverrides:
diff --git a/test/passes/pytorch/test_gptq.py b/test/passes/pytorch/test_gptq.py
index 0675255aa9..9ec4d24bc5 100644
--- a/test/passes/pytorch/test_gptq.py
+++ b/test/passes/pytorch/test_gptq.py
@@ -8,7 +8,7 @@
import torch
from olive.common.quant.hf_utils import OliveHfQuantizationConfig
-from olive.common.quant.nn import QuantLinear
+from olive.common.quant.tensor import QuantTensor
from olive.hardware.accelerator import AcceleratorSpec, Device
from olive.model import HfModelHandler
from olive.passes.olive_pass import create_pass_from_dict
@@ -16,6 +16,17 @@
from test.utils import get_tiny_phi3, make_local_tiny_llama
+def _is_quant(module: torch.nn.Module) -> bool:
+ if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
+ return False
+ weight = module._parameters.get("weight")
+ return weight is not None and isinstance(weight.data, QuantTensor)
+
+
+def _bits(module: torch.nn.Module) -> int:
+ return module.weight.data.bits
+
+
# running on CPU takes time so will only run a subset of tests when GPU is not available
@pytest.mark.parametrize(
("model_path", "expected_model_type"),
@@ -58,9 +69,9 @@ def test_gptq(tmp_path: Path, model_path: str, expected_model_type: str, group_s
assert hasattr(loaded_model.config, "quantization_config")
assert isinstance(loaded_model.config.quantization_config, OliveHfQuantizationConfig)
assert loaded_model.config.quantization_config.group_size == group_size
- assert not any(isinstance(m, torch.nn.Linear) for m in loaded_model.model.layers.modules())
- assert isinstance(loaded_model.model.layers[0].self_attn.o_proj, QuantLinear)
- assert loaded_model.model.layers[0].self_attn.o_proj.quantizer.bits == 8
- assert loaded_model.model.layers[0].mlp.down_proj.quantizer.bits == 4
+ assert not any(isinstance(m, torch.nn.Linear) and not _is_quant(m) for m in loaded_model.model.layers.modules())
+ assert _is_quant(loaded_model.model.layers[0].self_attn.o_proj)
+ assert _bits(loaded_model.model.layers[0].self_attn.o_proj) == 8
+ assert _bits(loaded_model.model.layers[0].mlp.down_proj) == 4
assert loaded_model.config.quantization_config.lm_head == lm_head
- assert isinstance(loaded_model.lm_head, QuantLinear) == lm_head
+ assert _is_quant(loaded_model.lm_head) == lm_head
diff --git a/test/passes/pytorch/test_quant_utils.py b/test/passes/pytorch/test_quant_utils.py
index ca3f7a93dd..402377dacb 100644
--- a/test/passes/pytorch/test_quant_utils.py
+++ b/test/passes/pytorch/test_quant_utils.py
@@ -2,29 +2,23 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
-"""Unit tests for olive.passes.pytorch.quant_utils helpers.
+"""Unit tests for olive quant state-dict helpers.
The full ``prepare_model`` / ``finalize`` pipeline is exercised end-to-end
-by ``test_rtn.py`` against a real HF model. This file targets the helpers
-that the MoE path adds — primarily ``flatten_quant_tensor_params`` — and
-keeps the tests free of network access.
+by ``test_rtn.py`` against a real HF model. This file targets the
+``install_quant_tensor_param`` helper used by both the 2D linear/embed
+path and the 3D fused-MoE path. The tests stay free of network access.
"""
import torch
from torch import nn
+from olive.common.quant.state_dict import install_quant_tensor_param
from olive.common.quant.tensor import QuantTensor
-from olive.passes.pytorch.quant_utils import flatten_quant_tensor_params
class _ExpertsBlock(nn.Module):
- """Fake fused-3D experts module (gpt-oss / Qwen3-MoE style).
-
- Carries a single 3D ``nn.Parameter`` of shape ``(num_experts, out, in)``,
- matching the layout that ``prepare_model`` annotates with
- ``quant_info_3d`` and that ``finalize`` rewrites into a 3D
- :class:`QuantTensor` parameter.
- """
+ """Fake fused-3D experts module (gpt-oss / Qwen3-MoE style)."""
def __init__(self, num_experts: int = 4, out_features: int = 32, in_features: int = 16):
super().__init__()
@@ -34,54 +28,63 @@ def __init__(self, num_experts: int = 4, out_features: int = 32, in_features: in
)
-class TestFlattenQuantTensorParams:
- def test_flatten_3d_quant_tensor_param(self):
+class TestInstallQuantTensorParam:
+ def test_install_3d_quant_tensor(self):
block = _ExpertsBlock()
qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=True, group_size=16)
- block._parameters["gate_up_proj"] = nn.Parameter(qt, requires_grad=False)
-
- flatten_quant_tensor_params(block)
-
- # The original parameter is gone, replaced by buffers.
- assert "gate_up_proj" not in dict(block.named_parameters())
- buffer_names = dict(block.named_buffers())
- assert "gate_up_proj_qweight" in buffer_names
- assert "gate_up_proj_scales" in buffer_names
- # symmetric → no qzeros buffer
- assert "gate_up_proj_qzeros" not in buffer_names
- # metadata stays as a plain attribute (not in state_dict)
- meta = block.gate_up_proj_quant_meta
- assert meta["bits"] == 4
- assert meta["group_size"] == 16
- assert meta["symmetric"] is True
- assert tuple(meta["shape"]) == (4, 32, 16)
-
- def test_flatten_asymmetric_emits_qzeros(self):
+
+ install_quant_tensor_param(block, "gate_up_proj", qt)
+
+ # Parameter is a QuantTensor (tensor subclass parameters return
+ # the subclass instance directly from nn.Parameter.__new__).
+ param = block._parameters["gate_up_proj"]
+ assert isinstance(param, QuantTensor)
+ # Sibling buffers are registered and share storage with the QuantTensor.
+ buffers = dict(block.named_buffers())
+ assert "gate_up_proj_qweight" in buffers
+ assert "gate_up_proj_scales" in buffers
+ assert "gate_up_proj_qzeros" not in buffers # symmetric → no qzeros
+ assert buffers["gate_up_proj_qweight"] is param.qweight
+ assert buffers["gate_up_proj_scales"] is param.scales
+
+ def test_install_asymmetric_emits_qzeros(self):
block = _ExpertsBlock()
qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=False, group_size=16)
- block._parameters["gate_up_proj"] = nn.Parameter(qt, requires_grad=False)
- flatten_quant_tensor_params(block)
+ install_quant_tensor_param(block, "gate_up_proj", qt)
- buffer_names = dict(block.named_buffers())
- assert "gate_up_proj_qzeros" in buffer_names
+ buffers = dict(block.named_buffers())
+ assert "gate_up_proj_qzeros" in buffers
+ assert buffers["gate_up_proj_qzeros"] is block._parameters["gate_up_proj"].qzeros
- def test_state_dict_serializable_after_flatten(self):
- """After flattening, the module's state_dict must be plain Tensors."""
+ def test_state_dict_drops_quant_tensor_entry(self):
+ """After install, state_dict must contain only plain Tensors (no QuantTensor entry)."""
block = _ExpertsBlock()
qt = QuantTensor.from_float(block.gate_up_proj.detach(), bits=4, symmetric=True, group_size=16)
- block._parameters["gate_up_proj"] = nn.Parameter(qt, requires_grad=False)
- flatten_quant_tensor_params(block)
+ install_quant_tensor_param(block, "gate_up_proj", qt)
sd = block.state_dict()
+ # No QuantTensor instance should appear in the state_dict.
for key, value in sd.items():
- assert type(value) is torch.Tensor, f"{key} is not a plain Tensor: {type(value)}"
-
- def test_no_op_on_module_without_quant_tensor(self):
- """Plain modules pass through unchanged."""
- block = nn.Linear(8, 8)
- before = {n: id(p) for n, p in block.named_parameters()}
- flatten_quant_tensor_params(block)
- after = {n: id(p) for n, p in block.named_parameters()}
- assert before == after
+ assert not isinstance(value, QuantTensor), f"{key} should not be a QuantTensor"
+ # The plain ``gate_up_proj`` key (the QuantTensor parameter) is dropped;
+ # the buffers carry the on-disk representation.
+ assert "gate_up_proj" not in sd
+ assert "gate_up_proj_qweight" in sd
+ assert "gate_up_proj_scales" in sd
+
+ def test_install_on_linear_module(self):
+ """Smoke test on a normal nn.Linear so that F.linear forward still works."""
+ linear = nn.Linear(16, 32, bias=False)
+ weight = linear.weight.detach().clone()
+ qt = QuantTensor.from_float(weight, bits=4, symmetric=True, group_size=16)
+
+ install_quant_tensor_param(linear, "weight", qt)
+
+ assert isinstance(linear.weight, QuantTensor)
+ # forward dispatches through QuantTensor.__torch_function__ and returns a plain Tensor
+ x = torch.randn(2, 16)
+ y = linear(x)
+ assert y.shape == (2, 32)
+ assert not isinstance(y, QuantTensor)
diff --git a/test/passes/pytorch/test_rtn.py b/test/passes/pytorch/test_rtn.py
index 934978ca58..f05da40cc8 100644
--- a/test/passes/pytorch/test_rtn.py
+++ b/test/passes/pytorch/test_rtn.py
@@ -8,7 +8,7 @@
import torch
from olive.common.quant.hf_utils import OliveHfQuantizationConfig
-from olive.common.quant.nn import QuantEmbedding, QuantLinear
+from olive.common.quant.tensor import QuantTensor
from olive.hardware.accelerator import AcceleratorSpec, Device
from olive.model import HfModelHandler
from olive.passes.olive_pass import create_pass_from_dict
@@ -16,6 +16,17 @@
from test.utils import get_tiny_phi3
+def _is_quant(module: torch.nn.Module) -> bool:
+ if not isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
+ return False
+ weight = module._parameters.get("weight")
+ return weight is not None and isinstance(weight.data, QuantTensor)
+
+
+def _bits(module: torch.nn.Module) -> int:
+ return module.weight.data.bits
+
+
@pytest.mark.parametrize("group_size", [-1, 16])
@pytest.mark.parametrize("sym", [True, False])
@pytest.mark.parametrize("lm_head", [True, False])
@@ -48,13 +59,14 @@ def test_gptq(tmp_path: Path, group_size: int, sym: bool, lm_head: bool):
assert hasattr(loaded_model.config, "quantization_config")
assert isinstance(loaded_model.config.quantization_config, OliveHfQuantizationConfig)
assert loaded_model.config.quantization_config.group_size == group_size
- assert not any(isinstance(m, torch.nn.Linear) for m in loaded_model.model.layers.modules())
- assert isinstance(loaded_model.model.layers[0].self_attn.o_proj, QuantLinear)
- assert loaded_model.model.layers[0].self_attn.o_proj.quantizer.bits == 8
- assert loaded_model.model.layers[0].mlp.down_proj.quantizer.bits == 4
+ assert not any(isinstance(m, torch.nn.Linear) and not _is_quant(m) for m in loaded_model.model.layers.modules())
+ assert _is_quant(loaded_model.model.layers[0].self_attn.o_proj)
+ assert _bits(loaded_model.model.layers[0].self_attn.o_proj) == 8
+ assert _bits(loaded_model.model.layers[0].mlp.down_proj) == 4
assert loaded_model.config.quantization_config.lm_head == lm_head
- assert isinstance(loaded_model.lm_head, QuantLinear) == lm_head
+ assert _is_quant(loaded_model.lm_head) == lm_head
assert isinstance(loaded_model.model.embed_tokens, torch.nn.Embedding)
+ assert not _is_quant(loaded_model.model.embed_tokens)
# compose another rtn pass on top of the partially quantized model
p2 = create_pass_from_dict(
@@ -76,8 +88,8 @@ def test_gptq(tmp_path: Path, group_size: int, sym: bool, lm_head: bool):
assert isinstance(out2, HfModelHandler)
loaded_model_2 = out2.load_model()
# check that the embed tokens layer is quantized to 8 bits
- assert isinstance(loaded_model_2.model.embed_tokens, QuantEmbedding)
- assert loaded_model_2.model.embed_tokens.quantizer.bits == 8
+ assert _is_quant(loaded_model_2.model.embed_tokens)
+ assert _bits(loaded_model_2.model.embed_tokens) == 8
# check that the lm head is quantized to 8 bits if it was not quantized before
- assert isinstance(loaded_model_2.lm_head, QuantLinear)
- assert loaded_model_2.lm_head.quantizer.bits == 4 if lm_head else 8
+ assert _is_quant(loaded_model_2.lm_head)
+ assert _bits(loaded_model_2.lm_head) == (4 if lm_head else 8)
From 8438093ad63199598166fc07d172a196848ebebe Mon Sep 17 00:00:00 2001
From: Jambay Kinley
Date: Fri, 15 May 2026 23:56:37 +0000
Subject: [PATCH 03/18] quant: drop QuantLinear/QuantEmbedding nn modules;
generalize quantizer to N-D
* Remove olive/common/quant/nn.py (QuantModule, QuantLinear,
QuantEmbedding) entirely. The only purpose of those modules was
ONNX export, which is now handled by reusing the existing
QuantLinearNbit from olive/common/hf/quant.py and a new
parallel QuantEmbeddingNbit (com.microsoft::GatherBlockQuantized
symbolic) in the same file.
* Add QuantLinearNbit.from_quant_tensor / QuantEmbeddingNbit.from_quant_tensor
factories so make_export_compatible_quant can swap any nn.Linear /
nn.Embedding whose weight is a QuantTensor into the export wrappers.
* Generalize WeightQuantizer (get_num_groups, get_qparam_shape,
find_qparams, quantize, dequantize, _reshape_tensor) and
pack_to_uint8 / unpack_from_uint8 to operate on any N-D tensor;
quantization is always along the last dim, leading dims are
preserved.
* Drop quantize_along_leading_dim / pack_to_uint8_along_last /
unpack_from_uint8_along_last and the explicit 3D leading-dim loops
in QuantTensor.from_float and _dequantize.
* Delete test/common/quant/test_nn.py; add N-D tests for the
generalized quantizer + pack helpers.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/quant.py | 265 ++++++++++++---
olive/common/quant/nn.py | 550 ------------------------------
olive/common/quant/tensor.py | 89 ++---
olive/common/quant/utils.py | 214 +++++-------
test/common/quant/test_nn.py | 571 --------------------------------
test/common/quant/test_utils.py | 48 ++-
6 files changed, 384 insertions(+), 1353 deletions(-)
delete mode 100644 olive/common/quant/nn.py
delete mode 100644 test/common/quant/test_nn.py
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index 06bb0b4c01..4eb8799f56 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -319,6 +319,218 @@ def from_tensors(
new_qlinear.bias = bias.clone()
return new_qlinear
+ @classmethod
+ def from_quant_tensor(
+ cls,
+ qt,
+ bias: torch.Tensor | None = None,
+ dynamo: bool = False,
+ ) -> QuantLinearNbit:
+ """Create a ``QuantLinearNbit`` from an Olive 2D ``QuantTensor``.
+
+ The QuantTensor stores ``qweight`` packed-along-last (uint8) in
+ ``(out_features, ceil(in_features / pack_factor))`` layout, with
+ ``scales`` of shape ``(out_features, n_groups)`` and optional
+ ``qzeros`` of shape ``(out_features, ceil(n_groups / pack_factor))``.
+ We unpack to dense uint8 / int8 tensors then re-pack into the
+ ``MatMulNBits`` layout via :meth:`pack`.
+ """
+ from olive.common.quant.tensor import QuantTensor
+ from olive.common.quant.utils import unpack_from_uint8
+
+ if not isinstance(qt, QuantTensor) or qt.dim() != 2:
+ raise ValueError("QuantLinearNbit.from_quant_tensor requires a 2D QuantTensor")
+
+ out_features, in_features = qt.shape
+ # unpack to (out, in) uint8 [0, 2^bits - 1]
+ iweight = unpack_from_uint8(qt.qweight, qt.bits, (out_features, in_features))
+ # unpack zero points; symmetric => midq
+ if qt.qzeros is not None:
+ izeros = unpack_from_uint8(qt.qzeros, qt.bits, tuple(qt.scales.shape))
+ else:
+ izeros = torch.full(
+ tuple(qt.scales.shape),
+ 1 << (qt.bits - 1),
+ dtype=torch.uint8,
+ device=qt.qweight.device,
+ )
+
+ # ``from_tensors`` expects iweight as (in, out) and scales/izeros as (n_groups, out)
+ return cls.from_tensors(
+ iweight=iweight.t().contiguous(),
+ scales=qt.scales.t().contiguous(),
+ izeros=izeros.t().contiguous(),
+ group_size=qt.group_size if qt.group_size > 0 else in_features,
+ bits=qt.bits,
+ bias=bias,
+ dynamo=dynamo,
+ )
+
+
+class QuantEmbeddingTorchFunction(torch.autograd.Function):
+ """Export a quantized embedding lookup as ``com.microsoft::GatherBlockQuantized``."""
+
+ # pylint: disable=W0223,W0221
+ @staticmethod
+ def symbolic(g, x, qweight, scales, qzeros, bits, group_size, embedding_dim):
+ tensor_args = [qweight, x, scales]
+ if qzeros is not None:
+ tensor_args.append(qzeros)
+ attrs = {"bits_i": bits, "block_size_i": group_size}
+
+ output = g.op(
+ "com.microsoft::GatherBlockQuantized",
+ *tensor_args,
+ outputs=1,
+ **attrs,
+ )
+ input_shape = x.type().varyingSizes()
+ if input_shape is not None and hasattr(x.type(), "with_sizes"):
+ output_type = scales.type().with_sizes([*input_shape, embedding_dim])
+ output.setType(output_type)
+ return output
+
+ @staticmethod
+ def forward(
+ ctx,
+ x: torch.Tensor,
+ qweight: torch.Tensor,
+ scales: torch.Tensor,
+ qzeros: torch.Tensor,
+ bits: int,
+ group_size: int,
+ embedding_dim: int,
+ dynamo: bool = False,
+ ):
+ if torch.onnx.is_in_onnx_export():
+ if dynamo and hasattr(torch.onnx, "ops"):
+ tensor_args = [qweight, x, scales]
+ if qzeros is not None:
+ tensor_args.append(qzeros)
+ attrs = {"bits": bits, "block_size": group_size}
+ return torch.onnx.ops.symbolic(
+ "com.microsoft::GatherBlockQuantized",
+ tensor_args,
+ attrs=attrs,
+ dtype=scales.dtype,
+ shape=[*x.shape, embedding_dim],
+ version=1,
+ )
+ if dynamo:
+ raise NotImplementedError(
+ "torch dynamo export for quantized embedding requires torch 2.8 or higher."
+ )
+ return torch.zeros((*x.shape, embedding_dim), dtype=scales.dtype, device=x.device)
+ raise NotImplementedError("QuantEmbeddingTorchFunction forward is only implemented for onnx export")
+
+
+class QuantEmbeddingNbit(torch.nn.Module):
+ """Quantized embedding layer exported as ``com.microsoft::GatherBlockQuantized``.
+
+ Buffer layout matches Olive's :class:`QuantTensor` 2D representation
+ exactly (packed uint8 along the last dim), so swapping a host
+ ``nn.Embedding`` with a ``QuantTensor`` weight for this module is a
+ direct buffer move.
+ """
+
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ group_size: int,
+ bits: int = 4,
+ symmetric: bool = True,
+ padding_idx: int | None = None,
+ dtype: torch.dtype = torch.float32,
+ device: torch.device | None = None,
+ dynamo: bool = False,
+ ):
+ super().__init__()
+ if bits not in (2, 4, 8):
+ raise ValueError(f"QuantEmbeddingNbit only supports 2/4/8 bits, got {bits}")
+ self.num_embeddings = num_embeddings
+ self.embedding_dim = embedding_dim
+ self.bits = bits
+ self.symmetric = symmetric
+ self.padding_idx = padding_idx
+ self.dynamo = dynamo
+
+ self.group_size = group_size if group_size > 0 else embedding_dim
+ pack_factor = 8 // bits
+ n_groups = math.ceil(embedding_dim / self.group_size)
+
+ self.register_buffer(
+ "qweight",
+ torch.zeros(
+ (num_embeddings, math.ceil(embedding_dim / pack_factor)),
+ dtype=torch.uint8,
+ device=device,
+ ),
+ )
+ self.register_buffer(
+ "scales",
+ torch.zeros((num_embeddings, n_groups), dtype=dtype, device=device),
+ )
+ if symmetric:
+ self.qzeros = None
+ else:
+ self.register_buffer(
+ "qzeros",
+ torch.zeros(
+ (num_embeddings, math.ceil(n_groups / pack_factor)),
+ dtype=torch.uint8,
+ device=device,
+ ),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return QuantEmbeddingTorchFunction.apply(
+ x,
+ self.qweight,
+ self.scales,
+ self.qzeros,
+ self.bits,
+ self.group_size,
+ self.embedding_dim,
+ self.dynamo,
+ )
+
+ @classmethod
+ def from_quant_tensor(
+ cls,
+ qt,
+ padding_idx: int | None = None,
+ dynamo: bool = False,
+ ) -> QuantEmbeddingNbit:
+ """Create a ``QuantEmbeddingNbit`` from an Olive 2D ``QuantTensor``.
+
+ The QuantTensor 2D buffer layout is bit-identical to this
+ module's layout, so the buffers are reused directly via
+ ``detach().clone()``.
+ """
+ from olive.common.quant.tensor import QuantTensor
+
+ if not isinstance(qt, QuantTensor) or qt.dim() != 2:
+ raise ValueError("QuantEmbeddingNbit.from_quant_tensor requires a 2D QuantTensor")
+
+ num_embeddings, embedding_dim = qt.shape
+ new = cls(
+ num_embeddings=num_embeddings,
+ embedding_dim=embedding_dim,
+ group_size=qt.group_size if qt.group_size > 0 else embedding_dim,
+ bits=qt.bits,
+ symmetric=qt.symmetric,
+ padding_idx=padding_idx,
+ dtype=qt.scales.dtype,
+ device=qt.qweight.device,
+ dynamo=dynamo,
+ )
+ new.qweight = qt.qweight.detach().clone()
+ new.scales = qt.scales.detach().clone()
+ if qt.qzeros is not None:
+ new.qzeros = qt.qzeros.detach().clone()
+ return new
+
def make_auto_awq_qlinearnbit(qlinear, dynamo):
if qlinear.w_bit not in [4, 8]:
@@ -375,10 +587,11 @@ def make_auto_gptq_qlinearnbit(qlinear, dynamo):
def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch.nn.Module:
"""Make the model export compatible by replacing the quantized linear layers with 4-bit versions."""
# Olive-native path: replace nn.Linear / nn.Embedding whose weight is a
- # ``QuantTensor`` parameter with the ONNX-exportable wrappers in
- # ``olive.common.quant.nn``. Done first so subsequent passes that
- # cast the model dtype don't dequantize through QuantTensor.
- _replace_olive_quant_tensor_modules(model)
+ # ``QuantTensor`` parameter with the ONNX-exportable wrappers
+ # ``QuantLinearNbit`` / ``QuantEmbeddingNbit`` defined in this module.
+ # Done first so subsequent passes that cast the model dtype don't
+ # dequantize through QuantTensor.
+ _replace_olive_quant_tensor_modules(model, dynamo)
model, modified = _replace_qlinear_modules(
model, dynamo, EXPORT_QLINEAR_MAPPING, "Making export compatible quantized model"
@@ -389,17 +602,11 @@ def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch.
return model
-def _replace_olive_quant_tensor_modules(model: torch.nn.Module) -> None:
- """Swap host ``nn.Linear`` / ``nn.Embedding`` (whose weight is a ``QuantTensor``)
- with the ONNX-exportable ``QuantLinear`` / ``QuantEmbedding`` wrappers.
-
- The on-disk buffer layout used by Olive's QuantTensor (packed-uint8
- along the last dim, per-group scales, optional packed qzeros) is
- bit-identical to the layout used by ``olive.common.quant.nn``
- ``QuantLinear`` / ``QuantEmbedding`` — both follow the auto-gptq
- style. So conversion is a direct buffer copy.
+def _replace_olive_quant_tensor_modules(model: torch.nn.Module, dynamo: bool) -> None:
+ """Swap host ``nn.Linear`` / ``nn.Embedding`` whose weight is a ``QuantTensor``
+ with the ONNX-exportable :class:`QuantLinearNbit` / :class:`QuantEmbeddingNbit`
+ wrappers.
"""
- from olive.common.quant.nn import QuantEmbedding, QuantLinear
from olive.common.quant.tensor import QuantTensor
targets: list[tuple[str, torch.nn.Module]] = []
@@ -417,33 +624,17 @@ def _replace_olive_quant_tensor_modules(model: torch.nn.Module) -> None:
for name, module in targets:
qt: QuantTensor = module.weight.data # type: ignore[assignment]
if isinstance(module, torch.nn.Embedding):
- new = QuantEmbedding(
- num_embeddings=module.num_embeddings,
- embedding_dim=module.embedding_dim,
- bits=qt.bits,
- symmetric=qt.symmetric,
- group_size=qt.group_size,
- device=qt.qweight.device,
- dtype=qt.scales.dtype,
+ new = QuantEmbeddingNbit.from_quant_tensor(
+ qt,
padding_idx=module.padding_idx,
+ dynamo=dynamo,
)
else:
- new = QuantLinear(
- in_features=module.in_features,
- out_features=module.out_features,
- bias=module.bias is not None,
- bits=qt.bits,
- symmetric=qt.symmetric,
- group_size=qt.group_size,
- device=qt.qweight.device,
- dtype=qt.scales.dtype,
+ new = QuantLinearNbit.from_quant_tensor(
+ qt,
+ bias=module.bias.detach().clone() if module.bias is not None else None,
+ dynamo=dynamo,
)
- if module.bias is not None:
- new.bias = module.bias.detach().clone()
- new.qweight = qt.qweight.detach().clone()
- new.scales = qt.scales.detach().clone()
- if qt.qzeros is not None:
- new.qzeros = qt.qzeros.detach().clone()
set_attr(model, name, new)
# After all replacements, the original ``quantization_method`` /
diff --git a/olive/common/quant/nn.py b/olive/common/quant/nn.py
deleted file mode 100644
index 78c804d052..0000000000
--- a/olive/common/quant/nn.py
+++ /dev/null
@@ -1,550 +0,0 @@
-# -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT License.
-# --------------------------------------------------------------------------
-from __future__ import annotations
-
-import math
-from abc import abstractmethod
-
-import torch
-import torch.nn as nn
-
-from olive.common.quant.utils import WeightQuantizer, pack_to_uint8, unpack_from_uint8
-
-
-class QuantModule(nn.Module):
- """A base class for quantized modules.
-
- Only supports 2D weight tensors for now. Quantization axis is assumed to be the last dimension.
- - 4-bit and 8-bit quantization
- - Symmetric and asymmetric quantization
- - Per-channel and groupwise quantization
-
- For blockwise quantization, it has the following restrictions:
- - The size of the last dimension must be divisible by the group size. This is to avoid needing to pad the weights.
- Padding is easy but it complicates compatibility between contrib ops and QDQ representations, and weight tieing.
- - group size must be >= 16 and power of 2. For compatibility with the contrib ops.
- # TODO(jambayk): extend to support more general cases if needed.
- """
-
- def __init__(
- self,
- rows: int,
- cols: int,
- bits: int = 4,
- symmetric: bool = True,
- group_size: int = -1,
- device: torch.device | None = None,
- dtype: torch.dtype = torch.float32,
- ):
- """Initialize QuantLinear layer.
-
- Args:
- rows: Number of rows in the weight matrix
- cols: Number of columns in the weight matrix
- bits: Number of bits for quantization (4 or 8)
- symmetric: Whether to use symmetric quantization
- group_size: Quantization group size (-1: per-channel, >0: groupwise)
- device: Device to place tensors on
- dtype: Data type for scales and bias
-
- """
- super().__init__()
- if bits not in [2, 4, 8]:
- raise ValueError(f"Only 2-bit, 4-bit and 8-bit quantization supported, got {bits}")
- if group_size != -1 and (group_size < 16 or (group_size & (group_size - 1)) != 0):
- raise ValueError("For blockwise quantization, group_size must be >= 16 and power of 2")
- if group_size != -1 and cols % group_size != 0:
- raise ValueError(
- f"For blockwise quantization, cols ({cols}) must be divisible by group_size ({group_size})"
- )
-
- self.rows = rows
- self.cols = cols
- self.quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False)
- self.device = device
- self.dtype = dtype
-
- packing_factor = 8 // bits
-
- # using the same layout and packing as auto-gptq
- # TODO(jambayk): consider other packing schemes
- self.register_buffer(
- "qweight",
- torch.zeros(
- # rows X cols, packed as uint8 along last dim
- (self.rows, math.ceil(self.cols / packing_factor)),
- dtype=torch.uint8,
- device=device,
- ),
- )
- scale_shape = self.quantizer.get_qparam_shape((self.rows, self.cols))
- self.register_buffer(
- "scales",
- torch.zeros(scale_shape, dtype=dtype, device=device),
- )
- if symmetric:
- self.qzeros = None
- else:
- self.register_buffer(
- "qzeros",
- torch.zeros(
- (scale_shape[0], math.ceil(scale_shape[1] / packing_factor)),
- dtype=torch.uint8,
- device=device,
- ),
- )
-
- @classmethod
- def from_tensors(
- cls,
- weight: torch.Tensor,
- bits: int = 4,
- symmetric: bool = True,
- group_size: int = -1,
- scales: torch.Tensor | None = None,
- zero_points: torch.Tensor | None = None,
- quantized: bool = False,
- **kwargs,
- ) -> QuantModule:
- """Create a QuantLinear layer from an existing nn.Linear layer.
-
- Args:
- weight: The weight tensor. Expected to be 2D (unsigned, in range [0, 2^bits - 1] if quantized is True).
- bits: Number of bits for quantization (4 or 8)
- symmetric: Whether to use symmetric quantization
- group_size: Quantization group size (-1: per-channel, >0: groupwise)
- scales: Optional precomputed scales for quantization
- zero_points: Optional precomputed zero points for quantization (unsigned, in range [0, 2^bits - 1]).
- quantized: Whether the provided weight is already quantized
- kwargs: Additional keyword arguments to pass to the QuantModule constructor
-
- Returns:
- A QuantModule instance with quantized weight
-
- """
- # pylint: disable=W0201
- if quantized:
- if scales is None or zero_points is None:
- raise ValueError("scales and/or zero_points missing for quantized weight")
- qweight = weight
- else:
- quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False)
-
- # compute quantization parameters if not provided
- if scales is None:
- scales, zero_points = quantizer.find_qparams(weight)
- else:
- scales = scales.to(weight.device).to(weight.dtype)
- zero_points = zero_points.to(weight.device).to(torch.int32)
-
- # quantize weights
- qweight = quantizer.quantize(weight, scales, zero_points)
-
- qmodule = cls(
- bits=bits,
- symmetric=symmetric,
- group_size=group_size,
- device=qweight.device,
- dtype=scales.dtype,
- **kwargs,
- )
- qmodule.qweight = pack_to_uint8(qweight, bits).contiguous()
- scale_shape = qmodule.quantizer.get_qparam_shape(qweight.shape)
- qmodule.scales = scales.reshape(scale_shape).contiguous()
-
- # enforce symmetric quantization constraints
- if symmetric and not torch.all(zero_points == qmodule.quantizer.midq):
- raise ValueError("Zero points must be equal to midq for symmetric quantization")
-
- if not symmetric:
- qmodule.qzeros = pack_to_uint8(zero_points.reshape(scale_shape), bits).contiguous()
-
- return qmodule
-
- @torch.no_grad()
- def unpack_and_dequantize(
- self, qweight: torch.Tensor, scales: torch.Tensor, zero_points: torch.Tensor | None
- ) -> torch.Tensor:
- """Unpack and dequantize the given quantized weight tensor.
-
- Returns:
- The dequantized weight tensor.
-
- """
- qweight = unpack_from_uint8(qweight, self.quantizer.bits, (qweight.shape[0], self.cols))
- if zero_points is not None:
- zero_points = unpack_from_uint8(zero_points, self.quantizer.bits, scales.shape)
- else:
- zero_points = torch.full_like(
- scales,
- self.quantizer.midq,
- dtype=torch.int32,
- )
- return self.quantizer.dequantize(qweight, scales, zero_points)
-
- @abstractmethod
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- """Forward pass of the quantized module.
-
- Args:
- x: The input tensor.
-
- Returns:
- The output tensor.
-
- """
-
-
-class QuantLinear(QuantModule):
- """Quantized Linear layer."""
-
- def __init__(
- self,
- in_features: int,
- out_features: int,
- bits: int = 4,
- symmetric: bool = True,
- group_size: int = -1,
- bias: bool = True,
- device: torch.device | None = None,
- dtype: torch.dtype = torch.float32,
- ):
- """Initialize QuantLinear layer.
-
- Args:
- in_features: Size of input features
- out_features: Size of output features
- bits: Number of bits for quantization (4 or 8)
- symmetric: Whether to use symmetric quantization
- group_size: Quantization group size (-1: per-channel, >0: groupwise)
- bias: Whether to include bias
- device: Device to place tensors on
- dtype: Data type for scales and bias
-
- """
- super().__init__(
- rows=out_features,
- cols=in_features,
- bits=bits,
- symmetric=symmetric,
- group_size=group_size,
- device=device,
- dtype=dtype,
- )
-
- if bias:
- self.register_buffer(
- "bias",
- torch.zeros(out_features, dtype=dtype, device=device),
- )
- else:
- self.bias = None
-
- @classmethod
- def from_module(
- cls,
- linear: nn.Linear,
- bits: int = 4,
- symmetric: bool = True,
- group_size: int = -1,
- scales: torch.Tensor | None = None,
- zero_points: torch.Tensor | None = None,
- ) -> QuantLinear:
- """Create a QuantLinear layer from an existing nn.Linear layer.
-
- Args:
- linear: The nn.Linear layer to convert
- bits: Number of bits for quantization (4 or 8)
- symmetric: Whether to use symmetric quantization
- group_size: Quantization group size (-1: per-channel, 0: per-tensor, >0: groupwise)
- scales: Optional precomputed scales for quantization
- zero_points: Optional precomputed zero points for quantization (unsigned, in range [0, 2^bits - 1]).
-
- Returns:
- A QuantLinear instance with quantized weights and scales
-
- """
- qlinear = cls.from_tensors(
- in_features=linear.in_features,
- out_features=linear.out_features,
- weight=linear.weight,
- bits=bits,
- symmetric=symmetric,
- group_size=group_size,
- scales=scales,
- zero_points=zero_points,
- bias=linear.bias is not None,
- )
- if linear.bias is not None:
- qlinear.bias = linear.bias.clone()
- return qlinear
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- """Forward pass for the quantized linear layer.
-
- Args:
- x: Input tensor of shape (..., in_features)
-
- Returns:
- Output tensor of shape (..., out_features)
-
- """
- assert x.shape[-1] == self.cols, f"Input shape {x.shape} does not match in_features {self.cols}"
-
- if torch.onnx.is_in_onnx_export():
- out = QuantLinearFunction.apply(
- x,
- self.qweight.reshape([*self.scales.shape, self.qweight.shape[-1] // self.scales.shape[-1]]),
- self.scales,
- self.qzeros,
- self.quantizer.bits,
- self.quantizer.group_size if self.quantizer.group_size > 0 else self.cols,
- self.cols,
- self.rows,
- )
- return out + self.bias if self.bias is not None else out
-
- x_dtype = x.dtype
-
- # unpack weights and zero points
- weight = self.unpack_and_dequantize(self.qweight, self.scales, self.qzeros)
- return nn.functional.linear(x, weight, bias=self.bias).to(x_dtype) # pylint: disable=not-callable
-
- def extra_repr(self) -> str:
- return (
- f"in_features={self.cols}, out_features={self.rows}, bits={self.quantizer.bits},"
- f" symmetric={self.quantizer.symmetric}, group_size={self.quantizer.group_size},"
- f" bias={self.bias is not None}"
- )
-
-
-# this is required for torchscript onnx export
-class QuantLinearFunction(torch.autograd.Function):
- # pylint: disable=W0223,W0221
- @staticmethod
- def symbolic(g, x, qweight, scales, qzeros, bits, group_size, in_features, out_features):
- tensor_args = [x, qweight, scales]
- if qzeros is not None:
- tensor_args.append(qzeros)
- attrs = {
- "K_i": in_features,
- "N_i": out_features,
- "bits_i": bits,
- "block_size_i": group_size,
- "accuracy_level_i": 4,
- }
-
- output = g.op(
- "com.microsoft::MatMulNBits",
- *tensor_args,
- outputs=1,
- **attrs,
- )
- input_shape = x.type().varyingSizes()
- if input_shape is not None and hasattr(x.type(), "with_sizes"):
- output_type = x.type().with_sizes([*input_shape[:-1], out_features])
- output.setType(output_type)
-
- return output
-
- @staticmethod
- def forward(
- ctx,
- x: torch.Tensor,
- qweight: torch.Tensor,
- scales: torch.Tensor,
- qzeros: torch.Tensor,
- bits: int,
- group_size: int,
- in_features: int,
- out_features: int,
- ):
- # there is no clean way to differentiate between torchscript and dynamo export
- # using FakeTensor as a proxy for dynamo export
- if hasattr(torch.onnx, "ops") and isinstance(x, torch._subclasses.FakeTensor): # pylint: disable=W0212
- tensor_args = [x, qweight, scales]
- if qzeros is not None:
- tensor_args.append(qzeros)
- attrs = {
- "K": in_features,
- "N": out_features,
- "bits": bits,
- "block_size": group_size,
- "accuracy_level": 4,
- }
- return torch.onnx.ops.symbolic(
- "com.microsoft::MatMulNBits",
- tensor_args,
- attrs=attrs,
- dtype=x.dtype,
- shape=[*x.shape[:-1], out_features],
- version=1,
- )
- return torch.zeros([*x.shape[:-1], out_features], dtype=x.dtype, device=x.device)
-
-
-class QuantEmbedding(QuantModule):
- """Quantized Embedding layer."""
-
- def __init__(
- self,
- num_embeddings: int,
- embedding_dim: int,
- bits: int = 4,
- symmetric: bool = True,
- group_size: int = -1,
- device: torch.device | None = None,
- dtype: torch.dtype = torch.float32,
- padding_idx: int | None = None,
- ):
- """Initialize QuantEmbedding layer.
-
- Args:
- num_embeddings: Number of embeddings
- embedding_dim: Dimension of each embedding
- padding_idx: Index of the padding token
- bits: Number of bits for quantization (4 or 8)
- symmetric: Whether to use symmetric quantization
- group_size: Quantization group size (-1: per-channel, >0: groupwise)
- device: Device to place tensors on
- dtype: Data type for scales
-
- """
- super().__init__(
- rows=num_embeddings,
- cols=embedding_dim,
- bits=bits,
- symmetric=symmetric,
- group_size=group_size,
- device=device,
- dtype=dtype,
- )
- self.padding_idx = padding_idx
-
- @classmethod
- def from_module(
- cls,
- embedding: nn.Embedding,
- bits: int = 4,
- symmetric: bool = True,
- group_size: int = -1,
- scales: torch.Tensor | None = None,
- zero_points: torch.Tensor | None = None,
- ) -> QuantEmbedding:
- """Create a QuantEmbedding layer from an existing nn.Embedding layer.
-
- Args:
- embedding: The nn.Embedding layer to convert
- bits: Number of bits for quantization (4 or 8)
- symmetric: Whether to use symmetric quantization
- group_size: Quantization group size (-1: per-channel, >0: groupwise)
- scales: Optional precomputed scales for quantization
- zero_points: Optional precomputed zero points for quantization (unsigned, in range [0, 2^bits - 1]).
-
- Returns:
- A QuantEmbedding instance with quantized weights and scales
-
- """
- return cls.from_tensors(
- num_embeddings=embedding.num_embeddings,
- embedding_dim=embedding.embedding_dim,
- weight=embedding.weight,
- bits=bits,
- symmetric=symmetric,
- group_size=group_size,
- scales=scales,
- zero_points=zero_points,
- padding_idx=embedding.padding_idx,
- )
-
- def forward(self, x: torch.Tensor) -> torch.Tensor:
- """Forward pass for the quantized embedding layer.
-
- Args:
- x: Input tensor of shape (...,)
-
- """
- if torch.onnx.is_in_onnx_export():
- return QuantEmbeddingFunction.apply(
- x,
- self.qweight,
- self.scales,
- self.qzeros,
- self.quantizer.bits,
- self.quantizer.group_size if self.quantizer.group_size > 0 else self.cols,
- self.cols,
- )
-
- # reshape input to 1D so that look up results are 2D
- x_shape = x.shape
- x = x.reshape(-1)
-
- # look up quantized weights and qparams
- qweight = nn.functional.embedding(x, self.qweight, padding_idx=self.padding_idx)
- scales = nn.functional.embedding(x, self.scales, padding_idx=self.padding_idx)
- if self.qzeros is not None:
- qzeros = nn.functional.embedding(x, self.qzeros, padding_idx=self.padding_idx)
- else:
- qzeros = None
-
- # unpack and dequantize
- return self.unpack_and_dequantize(qweight, scales, qzeros).reshape(*x_shape, -1)
-
- def extra_repr(self) -> str:
- return (
- f"{self.rows}, {self.cols}, bits={self.quantizer.bits},"
- f" symmetric={self.quantizer.symmetric}, group_size={self.quantizer.group_size},"
- f" padding_idx={self.padding_idx}"
- )
-
-
-# this is required for torchscript onnx export
-class QuantEmbeddingFunction(torch.autograd.Function):
- # pylint: disable=W0223,W0221
- @staticmethod
- def symbolic(g, x, qweight, scales, qzeros, bits, group_size, embedding_dim):
- tensor_args = [qweight, x, scales]
- if qzeros is not None:
- tensor_args.append(qzeros)
- attrs = {"bits_i": bits, "block_size_i": group_size}
-
- output = g.op(
- "com.microsoft::GatherBlockQuantized",
- *tensor_args,
- outputs=1,
- **attrs,
- )
- input_shape = x.type().varyingSizes()
- if input_shape is not None and hasattr(x.type(), "with_sizes"):
- output_type = scales.type().with_sizes([*input_shape, embedding_dim])
- output.setType(output_type)
-
- return output
-
- @staticmethod
- def forward(
- ctx,
- x: torch.Tensor,
- qweight: torch.Tensor,
- scales: torch.Tensor,
- qzeros: torch.Tensor,
- bits: int,
- group_size: int,
- embedding_dim: int,
- ):
- if hasattr(torch.onnx, "ops") and isinstance(x, torch._subclasses.FakeTensor): # pylint: disable=W0212
- tensor_args = [qweight, x, scales]
- if qzeros is not None:
- tensor_args.append(qzeros)
- attrs = {"bits": bits, "block_size": group_size}
- return torch.onnx.ops.symbolic(
- "com.microsoft::GatherBlockQuantized",
- tensor_args,
- attrs=attrs,
- dtype=scales.dtype,
- shape=[*x.shape, embedding_dim],
- version=1,
- )
- return torch.zeros((*x.shape, embedding_dim), dtype=scales.dtype, device=x.device)
diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py
index b25c788cfd..6c4628df7f 100644
--- a/olive/common/quant/tensor.py
+++ b/olive/common/quant/tensor.py
@@ -6,8 +6,7 @@
quantization buffers (``qweight``, ``scales``, ``qzeros``) but presents the
shape / dtype / device of the dequantized full-precision weight.
-Design notes (see ``/datadisks/jambaykinley/archive/qmoe/research.md`` and
-``plan.md``):
+Design notes:
* The class is a **wrapper** subclass (``_make_wrapper_subclass``) — it
carries no real storage of its own, so the dense FP weight is never
@@ -19,9 +18,8 @@
conversion pass swaps any ``nn.Linear`` / ``nn.Embedding`` whose
weight is a ``QuantTensor`` for the existing exportable
``QuantLinearNbit`` / ``QuantEmbeddingNbit`` ``nn.Module``s (see
- ``olive/common/quant/nn.py`` and ``olive/common/hf/quant.py``)
- *before* the tracer ever inspects the parameter. This keeps the
- legacy ``com.microsoft::MatMulNBits`` /
+ ``olive/common/hf/quant.py``) *before* the tracer ever inspects
+ the parameter. This keeps the legacy ``com.microsoft::MatMulNBits`` /
``com.microsoft::GatherBlockQuantized`` symbolic emission intact.
* All other ops (including ``model.to(dtype/device)``, ``.detach()``,
``.contiguous()``, ``.clone()``) are routed through ``_apply_fn_to_data``
@@ -44,7 +42,6 @@
WeightQuantizer,
pack_to_uint8,
unpack_from_uint8,
- unpack_from_uint8_along_last,
)
__all__ = ["QuantTensor", "implements"]
@@ -71,28 +68,21 @@ def _midq(bits: int) -> int:
def _zero_points_or_default(weight: QuantTensor) -> torch.Tensor:
"""Unpack zero_points or return a tensor full of the symmetric mid-q value."""
if weight.qzeros is not None:
- return unpack_from_uint8_along_last(weight.qzeros, weight.bits, tuple(weight.scales.shape)).to(torch.int32)
+ return unpack_from_uint8(weight.qzeros, weight.bits, tuple(weight.scales.shape)).to(torch.int32)
return torch.full(weight.scales.shape, _midq(weight.bits), dtype=torch.int32, device=weight.scales.device)
def _dequantize(weight: QuantTensor) -> torch.Tensor:
"""Unpack + dequantize ``weight`` into a dense tensor of ``weight.dtype``."""
+ if weight.dim() not in (2, 3):
+ raise NotImplementedError(f"QuantTensor only supports 2D / 3D layouts, got {weight.dim()}D")
+
quantizer = WeightQuantizer(
bits=weight.bits, symmetric=weight.symmetric, group_size=weight.group_size, signed=False
)
- if weight.dim() == 2:
- qw = unpack_from_uint8(weight.qweight, weight.bits, tuple(weight.shape))
- zp = _zero_points_or_default(weight)
- return quantizer.dequantize(qw, weight.scales, zp).to(weight.dtype)
- if weight.dim() == 3:
- qw = unpack_from_uint8_along_last(weight.qweight, weight.bits, tuple(weight.shape))
- zp = _zero_points_or_default(weight)
- leading = weight.shape[0]
- out = torch.empty(weight.shape, dtype=weight.dtype, device=weight.qweight.device)
- for i in range(leading):
- out[i] = quantizer.dequantize(qw[i], weight.scales[i], zp[i]).to(weight.dtype)
- return out
- raise NotImplementedError(f"QuantTensor only supports 2D / 3D layouts, got {weight.dim()}D")
+ qw = unpack_from_uint8(weight.qweight, weight.bits, tuple(weight.shape))
+ zp = _zero_points_or_default(weight)
+ return quantizer.dequantize(qw, weight.scales, zp).to(weight.dtype)
class QuantTensor(torch.Tensor):
@@ -174,51 +164,32 @@ def from_float(
) -> QuantTensor:
"""Quantize a 2D or 3D FP weight tensor and produce a ``QuantTensor``.
- For 3D weights the leading dim is iterated and each slice is
- quantized independently — the resulting qweight retains the
- leading dim and ``scales`` / ``qzeros`` gain one as well.
+ Quantization is along the last dim — for 3D fused MoE weights of
+ shape ``(num_experts, out, in)`` each ``(out, in)`` slice gets
+ its own per-group scales / zero-points along ``in``, with no
+ explicit leading-dim loop.
"""
if weight.dim() not in (2, 3):
raise ValueError(f"QuantTensor only supports 2D and 3D weights, got shape {tuple(weight.shape)}")
quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False)
- if weight.dim() == 2:
- if scales is None or zero_points is None:
- scales, zero_points = quantizer.find_qparams(weight)
- else:
- scales = scales.to(weight.device).to(weight.dtype)
- zero_points = zero_points.to(weight.device).to(torch.int32)
- qweight_int = quantizer.quantize(weight, scales, zero_points)
- qweight_packed = pack_to_uint8(qweight_int, bits).contiguous()
- scales_packed = scales.reshape(quantizer.get_qparam_shape(weight.shape)).contiguous()
- if symmetric:
- qzeros_packed = None
- if not torch.all(zero_points == quantizer.midq):
- raise ValueError("Zero points must equal midq for symmetric quantization")
- else:
- qzeros_packed = pack_to_uint8(
- zero_points.reshape(quantizer.get_qparam_shape(weight.shape)), bits
- ).contiguous()
+ qparam_shape = quantizer.get_qparam_shape(tuple(weight.shape))
+
+ if scales is None or zero_points is None:
+ scales, zero_points = quantizer.find_qparams(weight)
+ else:
+ scales = scales.to(weight.device).to(weight.dtype).reshape(qparam_shape)
+ zero_points = zero_points.to(weight.device).to(torch.int32).reshape(qparam_shape)
+
+ qweight_int = quantizer.quantize(weight, scales, zero_points)
+ qweight_packed = pack_to_uint8(qweight_int, bits).contiguous()
+ scales_packed = scales.reshape(qparam_shape).contiguous()
+ if symmetric:
+ if not torch.all(zero_points == quantizer.midq):
+ raise ValueError("Zero points must equal midq for symmetric quantization")
+ qzeros_packed = None
else:
- # 3D — iterate the leading dim
- if scales is not None or zero_points is not None:
- raise NotImplementedError("Pre-computed scales/zero_points are not yet supported for 3D weights")
- qweight_list, scales_list, zp_list = [], [], []
- for i in range(weight.shape[0]):
- s, zp = quantizer.find_qparams(weight[i])
- qw = quantizer.quantize(weight[i], s, zp)
- qweight_list.append(pack_to_uint8(qw, bits))
- scales_list.append(s.reshape(quantizer.get_qparam_shape(weight[i].shape)))
- zp_list.append(zp.reshape(quantizer.get_qparam_shape(weight[i].shape)))
- qweight_packed = torch.stack(qweight_list, dim=0).contiguous()
- scales_packed = torch.stack(scales_list, dim=0).contiguous()
- if symmetric:
- qzeros_packed = None
- for zp in zp_list:
- if not torch.all(zp == quantizer.midq):
- raise ValueError("Zero points must equal midq for symmetric quantization")
- else:
- qzeros_packed = torch.stack([pack_to_uint8(zp, bits) for zp in zp_list], dim=0).contiguous()
+ qzeros_packed = pack_to_uint8(zero_points.reshape(qparam_shape), bits).contiguous()
return cls(
qweight=qweight_packed,
diff --git a/olive/common/quant/utils.py b/olive/common/quant/utils.py
index 2aac15477f..7b478ad906 100644
--- a/olive/common/quant/utils.py
+++ b/olive/common/quant/utils.py
@@ -11,11 +11,19 @@
class WeightQuantizer:
"""Class to quantize weight tensors.
- Operates on 2D tensors of shape ``(out_features, in_features)`` with
- the last dimension quantized (groupwise / per-channel / per-tensor).
- Use :func:`quantize_along_leading_dim` to handle 3D tensors (e.g.
- fused MoE experts of shape ``(num_experts, out, in)``) by iterating
- the leading dimension.
+ Operates on N-D tensors and always quantizes along the **last**
+ dimension. For a tensor with shape ``(*leading_dims, last)``:
+
+ * per-tensor (``group_size=0``): a single scalar scale shared across
+ all elements; ``scales`` has shape ``(1,) * ndim``.
+ * per-channel (``group_size=-1``): one scale per leading-index;
+ ``scales`` has shape ``(*leading_dims, 1)``.
+ * groupwise (``group_size>0``): ``last`` must be divisible by
+ ``group_size``; ``scales`` has shape ``(*leading_dims, last // group_size)``.
+
+ The same code path handles the 2D ``(out, in)`` case used by
+ ``nn.Linear``/``nn.Embedding`` and the 3D fused-MoE
+ ``(num_experts, out, in)`` case — no leading-dim loop required.
"""
def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, signed: bool = False):
@@ -37,63 +45,64 @@ def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, s
self.maxq, self.minq = get_maxq_minq(self.bits, self.signed)
self.midq = (self.maxq + self.minq + 1) // 2
- def get_num_groups(self, shape: tuple[int, int]) -> int:
- """Get the number of groups for quantization based on the input shape and group_size.
+ def get_num_groups(self, shape: tuple[int, ...]) -> int:
+ """Get the number of groups along the last dim.
Args:
- shape: The shape (out_features, in_features) of the tensor to quantize
+ shape: The shape of the tensor to quantize.
Returns:
- The number of groups for quantization
+ The number of groups along the last dim.
"""
if self.group_size == 0:
raise ValueError("group_size must be greater than 0 for groupwise quantization")
- group_size = self.group_size if self.group_size > 0 else shape[1]
- assert shape[1] % group_size == 0, f"in_features {shape[1]} must be divisible by group_size {group_size}"
- return shape[1] // group_size
+ last = shape[-1]
+ group_size = self.group_size if self.group_size > 0 else last
+ assert last % group_size == 0, f"last dim {last} must be divisible by group_size {group_size}"
+ return last // group_size
- def get_qparam_shape(self, shape: tuple[int, int]) -> tuple[int, ...]:
- """Get the shapes for quantization parameters based on the input shape and group_size.
+ def get_qparam_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
+ """Get the shape for scales / zero-points given the input shape.
Args:
- shape: The shape (out_features, in_features) of the tensor to quantize
+ shape: The shape of the tensor to quantize.
Returns:
- A tuple of shapes for scales and zero points
+ The shape of the scales / zero-points tensor.
"""
if self.group_size == 0:
- return (1, 1)
- return (shape[0], self.get_num_groups(shape))
+ return (1,) * len(shape)
+ return (*shape[:-1], self.get_num_groups(shape))
@torch.no_grad()
def find_qparams(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Find quantization parameters (scale and zero point) for the given tensor.
Args:
- tensor: The tensor to quantize. Expected to be 2D with shape (out_features, in_features).
+ tensor: The N-D tensor to quantize. Quantization is along the last dim.
Returns:
- A tuple of (scale, zero_point)
+ A tuple of (scale, zero_point) with shape :meth:`get_qparam_shape`.
"""
tensor, _ = self._reshape_tensor(tensor)
- # calculate min and max
- tmp = torch.zeros(tensor.shape[0:-1], device=tensor.device, dtype=tensor.dtype)
- min_val = torch.minimum(tensor.min(-1)[0], tmp)
- max_val = torch.maximum(tensor.max(-1)[0], tmp)
+ # calculate min and max along the last (group) dim
+ zero = torch.zeros(tensor.shape[:-1], device=tensor.device, dtype=tensor.dtype)
+ min_val = torch.minimum(tensor.min(-1)[0], zero)
+ max_val = torch.maximum(tensor.max(-1)[0], zero)
if self.symmetric:
max_val = torch.maximum(abs(min_val), max_val)
- tmp = min_val < 0
- if torch.any(tmp):
- min_val[tmp] = -max_val[tmp]
+ neg = min_val < 0
+ if torch.any(neg):
+ min_val[neg] = -max_val[neg]
- tmp = (min_val == 0) & (max_val == 0)
- min_val[tmp] = -1
- max_val[tmp] = 1
+ zero_pair = (min_val == 0) & (max_val == 0)
+ min_val[zero_pair] = -1
+ max_val[zero_pair] = 1
scales = (max_val - min_val) / (self.maxq - self.minq)
if self.symmetric:
@@ -109,12 +118,12 @@ def quantize(self, tensor: torch.Tensor, scales: torch.Tensor, zero_points: torc
"""Quantize the given tensor using the provided scales and zero points.
Args:
- tensor: The tensor to quantize. Expected to be 2D with shape (out_features, in_features).
- scales: The scales for quantization.
- zero_points: The zero points for quantization.
+ tensor: The N-D tensor to quantize.
+ scales: The scales with shape :meth:`get_qparam_shape`.
+ zero_points: The zero points with shape :meth:`get_qparam_shape`.
Returns:
- The quantized tensor.
+ The quantized tensor with the original input shape.
"""
tensor, shape = self._reshape_tensor(tensor)
@@ -127,20 +136,19 @@ def quantize(self, tensor: torch.Tensor, scales: torch.Tensor, zero_points: torc
@torch.no_grad()
def dequantize(self, q_tensor: torch.Tensor, scales: torch.Tensor, zero_points: torch.Tensor) -> torch.Tensor:
- """Dequantize the given quantized tensor using the provided scales and zero points.
+ """Dequantize the given quantized tensor.
Args:
- q_tensor: The quantized tensor to dequantize. Expected to be 2D with shape (out_features, in_features).
- scales: The scales for dequantization.
- zero_points: The zero points for dequantization.
+ q_tensor: The quantized N-D tensor.
+ scales: The scales with shape :meth:`get_qparam_shape`.
+ zero_points: The zero points with shape :meth:`get_qparam_shape`.
Returns:
- The dequantized tensor.
+ The dequantized tensor with the original input shape.
"""
q_tensor, shape = self._reshape_tensor(q_tensor)
- # apply dequantization
# both q_tensor and zero_points should be int32, so no need to worry about overflow
tensor = (q_tensor - zero_points.unsqueeze(-1)) * scales.unsqueeze(-1)
return tensor.reshape(shape)
@@ -151,7 +159,7 @@ def fake_quantize(
"""Fake quantize the given tensor using the provided scales and zero points.
Args:
- tensor: The tensor to quantize. Expected to be 2D with shape (out_features, in_features).
+ tensor: The N-D tensor to quantize.
scales: The scales for quantization. If None, scales will be computed from the tensor.
zero_points: The zero points for quantization. If None, zero points will be computed from the tensor.
@@ -165,20 +173,16 @@ def fake_quantize(
return self.dequantize(q_tensor, scales, zero_points)
def _reshape_tensor(self, tensor: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]:
- """Reshape the tensor based on the group size.
-
- Args:
- tensor: The tensor to reshape.
-
- Returns:
- The reshaped tensor and its original shape.
+ """Reshape so the last dim becomes ``(num_groups, group_size)``.
+ Returns the reshaped tensor and its original shape.
"""
- shape = tensor.shape
+ shape = tuple(tensor.shape)
if self.group_size == 0:
- tensor = tensor.reshape(1, 1, -1)
+ # per-tensor: collapse everything into a single (1, ..., 1, prod) tensor
+ tensor = tensor.reshape(*((1,) * (len(shape) - 1)), 1, -1)
else:
- tensor = tensor.reshape(shape[0], self.get_num_groups(shape), -1)
+ tensor = tensor.reshape(*shape[:-1], self.get_num_groups(shape), -1)
return tensor, shape
@@ -205,14 +209,17 @@ def get_maxq_minq(bits: int, signed: bool) -> tuple[int, int]:
@torch.no_grad()
def pack_to_uint8(tensor: torch.Tensor, bits: int) -> torch.Tensor:
- """Pack 2/4/8 bit tensor into uint8 tensor on the last dimension.
+ """Pack 2/4/8 bit values into uint8 along the last dimension.
+
+ Works for tensors of any rank — only the last dim is packed; all
+ leading dims are preserved.
Args:
- tensor: The 2D input tensor. Values are expected to be unsigned in the range [0, 2^bits - 1]
+ tensor: The input tensor. Values are expected to be unsigned in the range ``[0, 2^bits - 1]``.
bits: Number of bits (2, 4 or 8)
Returns:
- A tensor of uint8 values with packed data
+ A tensor of uint8 values with packed data along the last dim.
"""
assert bits in [2, 4, 8], "Only 2-bit, 4-bit and 8-bit quantization supported"
@@ -223,110 +230,47 @@ def pack_to_uint8(tensor: torch.Tensor, bits: int) -> torch.Tensor:
packing_factor = 8 // bits
- # padd if necessary to ensure tensor shape is divisible by packing_factor
+ # pad the last dim so it is divisible by ``packing_factor``
num_padding = (-tensor.shape[-1]) % packing_factor
if num_padding > 0:
- pad = (0, num_padding, 0, 0)
- tensor = torch.nn.functional.pad(tensor, pad, mode="constant", value=0)
+ tensor = torch.nn.functional.pad(tensor, (0, num_padding), mode="constant", value=0)
packed_size = tensor.shape[-1] // packing_factor
tensor = tensor.to(torch.uint8)
packed_tensor = torch.zeros(
- (tensor.shape[0], packed_size),
+ (*tensor.shape[:-1], packed_size),
dtype=torch.uint8,
device=tensor.device,
)
for i in range(packing_factor):
- packed_tensor |= tensor[:, i::packing_factor] << bits * i
+ packed_tensor |= tensor[..., i::packing_factor] << bits * i
return packed_tensor
@torch.no_grad()
-def unpack_from_uint8(packed_tensor: torch.Tensor, bits: int, shape: tuple[int, int]) -> torch.Tensor:
- """Unpack uint8 tensor into 2/4/8 bit tensor on the last dimension.
+def unpack_from_uint8(packed_tensor: torch.Tensor, bits: int, shape: tuple[int, ...]) -> torch.Tensor:
+ """Unpack a uint8-packed tensor into 2/4/8 bit values along the last dimension.
+
+ Works for tensors of any rank — only the last dim is unpacked; all
+ leading dims must match ``shape[:-1]``.
Args:
- packed_tensor: The 2D input tensor with packed uint8 values
+ packed_tensor: The packed uint8 tensor.
bits: Number of bits (2, 4 or 8)
- shape: The original shape of the tensor before packing
+ shape: The original shape of the tensor before packing.
Returns:
- A tensor of int32 values with unpacked data
+ A tensor of int32 values with unpacked data.
"""
assert packed_tensor.dtype == torch.uint8, "Input tensor must be of dtype uint8"
maxq, _ = get_maxq_minq(bits, signed=False)
- wf = torch.arange(0, 8, bits, device=packed_tensor.device, dtype=torch.uint8).unsqueeze(0)
-
- unpacked_tensor = torch.bitwise_right_shift(packed_tensor.unsqueeze(2), wf.unsqueeze(0))
- unpacked_tensor = unpacked_tensor.reshape(shape[0], -1)
- unpacked_tensor = unpacked_tensor[:, : shape[1]]
+ wf = torch.arange(0, 8, bits, device=packed_tensor.device, dtype=torch.uint8)
+ # (..., packed_size, packing_factor)
+ unpacked_tensor = torch.bitwise_right_shift(packed_tensor.unsqueeze(-1), wf)
+ # collapse last two dims and trim padding back to original last-dim size
+ unpacked_tensor = unpacked_tensor.reshape(*packed_tensor.shape[:-1], -1)
+ unpacked_tensor = unpacked_tensor[..., : shape[-1]]
return torch.bitwise_and(unpacked_tensor, maxq).to(torch.int32)
-
-
-@torch.no_grad()
-def quantize_along_leading_dim(
- quantizer: WeightQuantizer, tensor: torch.Tensor
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- """Quantize a 2D or 3D tensor with a :class:`WeightQuantizer`.
-
- For a 2D input of shape ``(out, in)`` this is equivalent to
- ``quantizer.find_qparams`` + ``quantizer.quantize``.
-
- For a 3D input of shape ``(leading, out, in)`` the leading dimension
- is iterated and each slice is quantized independently. The returned
- ``qweight`` shares the same shape as the input; ``scales`` /
- ``zero_points`` carry an additional leading dimension matching the
- input's leading dim.
- """
- if tensor.dim() == 2:
- scales, zero_points = quantizer.find_qparams(tensor)
- qweight = quantizer.quantize(tensor, scales, zero_points)
- return qweight, scales, zero_points
- if tensor.dim() != 3:
- raise ValueError(f"Only 2D and 3D tensors are supported, got shape {tuple(tensor.shape)}")
-
- leading = tensor.shape[0]
- qweights, scales_list, zp_list = [], [], []
- for i in range(leading):
- s, zp = quantizer.find_qparams(tensor[i])
- qweights.append(quantizer.quantize(tensor[i], s, zp))
- scales_list.append(s)
- zp_list.append(zp)
- return (
- torch.stack(qweights, dim=0),
- torch.stack(scales_list, dim=0),
- torch.stack(zp_list, dim=0),
- )
-
-
-@torch.no_grad()
-def pack_to_uint8_along_last(tensor: torch.Tensor, bits: int) -> torch.Tensor:
- """Pack 2/4/8 bit values into uint8 along the last dim, supporting 2D or 3D.
-
- For 3D inputs (e.g. fused MoE experts) the leading dim is preserved.
- """
- if tensor.dim() == 2:
- return pack_to_uint8(tensor, bits)
- if tensor.dim() != 3:
- raise ValueError(f"Only 2D and 3D tensors are supported, got shape {tuple(tensor.shape)}")
-
- leading = tensor.shape[0]
- return torch.stack([pack_to_uint8(tensor[i], bits) for i in range(leading)], dim=0)
-
-
-@torch.no_grad()
-def unpack_from_uint8_along_last(packed_tensor: torch.Tensor, bits: int, shape: tuple[int, ...]) -> torch.Tensor:
- """Unpack uint8-packed values into 2/4/8 bit ints, supporting 2D or 3D."""
- if len(shape) == 2:
- return unpack_from_uint8(packed_tensor, bits, shape)
- if len(shape) != 3:
- raise ValueError(f"Only 2D and 3D tensors are supported, got shape {shape}")
-
- leading = shape[0]
- return torch.stack(
- [unpack_from_uint8(packed_tensor[i], bits, shape[1:]) for i in range(leading)],
- dim=0,
- )
diff --git a/test/common/quant/test_nn.py b/test/common/quant/test_nn.py
deleted file mode 100644
index c496dd21da..0000000000
--- a/test/common/quant/test_nn.py
+++ /dev/null
@@ -1,571 +0,0 @@
-# -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT License.
-# --------------------------------------------------------------------------
-import math
-
-import pytest
-import torch
-import torch.nn as nn
-
-from olive.common.quant.nn import QuantEmbedding, QuantLinear
-from olive.common.quant.utils import WeightQuantizer
-
-
-class TestQuantModule:
- @pytest.mark.parametrize("bits", [2, 4, 8])
- @pytest.mark.parametrize("symmetric", [True, False])
- @pytest.mark.parametrize("group_size", [-1, 16, 32])
- def test_initialization(self, bits, symmetric, group_size):
- """Test QuantModule initialization with various parameters."""
- rows, cols = 64, 128
-
- # Since QuantModule is abstract, we'll use QuantLinear
- qmodule = QuantLinear(
- in_features=cols,
- out_features=rows,
- bits=bits,
- symmetric=symmetric,
- group_size=group_size,
- )
-
- assert qmodule.rows == rows
- assert qmodule.cols == cols
- assert qmodule.quantizer.bits == bits
- assert qmodule.quantizer.symmetric == symmetric
- assert qmodule.quantizer.group_size == group_size
-
- def test_invalid_bits(self):
- """Test that invalid bits raise ValueError."""
- with pytest.raises(ValueError, match="Only 2-bit, 4-bit and 8-bit quantization supported"):
- QuantLinear(10, 20, bits=16, symmetric=True, group_size=-1)
-
- def test_invalid_group_size(self):
- """Test that invalid group size raises ValueError."""
- with pytest.raises(ValueError, match="group_size must be >= 16 and power of 2"):
- QuantLinear(10, 20, bits=4, symmetric=True, group_size=15)
-
- with pytest.raises(ValueError, match="group_size must be >= 16 and power of 2"):
- QuantLinear(10, 20, bits=4, symmetric=True, group_size=24)
-
- def test_invalid_in_features_for_group_size(self):
- """Test that in_features must be divisible by group_size."""
- # in_features=100 is not divisible by group_size=32
- with pytest.raises(ValueError, match=r"cols .* must be divisible by group_size"):
- QuantLinear(in_features=100, out_features=20, bits=4, symmetric=True, group_size=32)
-
- # in_features=50 is not divisible by group_size=16
- with pytest.raises(ValueError, match=r"cols .* must be divisible by group_size"):
- QuantLinear(in_features=50, out_features=20, bits=4, symmetric=True, group_size=16)
-
- # This should work: in_features=64 is divisible by group_size=32
- qlinear = QuantLinear(in_features=64, out_features=20, bits=4, symmetric=True, group_size=32)
- assert qlinear.cols == 64
-
- def test_buffer_shapes(self):
- """Test that buffers have correct shapes."""
- rows, cols = 64, 128
- bits = 4
- packing_factor = 8 // bits
-
- qmodule = QuantLinear(
- in_features=cols,
- out_features=rows,
- bits=bits,
- symmetric=True,
- group_size=32,
- )
-
- # Check qweight shape
- assert qmodule.qweight.shape == (rows, math.ceil(cols / packing_factor))
-
- # Check scales shape
- quantizer = WeightQuantizer(bits=bits, symmetric=True, group_size=32, signed=False)
- expected_scale_shape = quantizer.get_qparam_shape((rows, cols))
- assert qmodule.scales.shape == expected_scale_shape
-
- def test_symmetric_no_qzeros(self):
- """Test that symmetric quantization has no qzeros."""
- qmodule = QuantLinear(16, 20, bits=4, symmetric=True, group_size=16)
- assert qmodule.qzeros is None
-
- def test_asymmetric_has_qzeros(self):
- """Test that asymmetric quantization has qzeros."""
- qmodule = QuantLinear(16, 20, bits=4, symmetric=False, group_size=16)
- assert qmodule.qzeros is not None
-
- packing_factor = 8 // 4
- quantizer = WeightQuantizer(bits=4, symmetric=False, group_size=16, signed=False)
- scale_shape = quantizer.get_qparam_shape((20, 16))
- expected_qzeros_shape = (scale_shape[0], math.ceil(scale_shape[1] / packing_factor))
- assert qmodule.qzeros.shape == expected_qzeros_shape
-
-
-class TestQuantLinear:
- def test_basic_initialization(self):
- """Test basic QuantLinear initialization."""
- qlinear = QuantLinear(in_features=128, out_features=256, bits=4, symmetric=True, group_size=32)
- assert qlinear.cols == 128
- assert qlinear.rows == 256
- assert qlinear.bias is not None
-
- def test_initialization_without_bias(self):
- """Test QuantLinear initialization without bias."""
- qlinear = QuantLinear(
- in_features=128,
- out_features=256,
- bits=4,
- symmetric=True,
- group_size=32,
- bias=False,
- )
- assert qlinear.bias is None
-
- def test_from_module_basic(self):
- """Test creating QuantLinear from nn.Linear."""
- linear = nn.Linear(128, 256)
- linear.weight.data.normal_(0, 0.02)
-
- qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32)
-
- assert qlinear.cols == 128
- assert qlinear.rows == 256
- assert qlinear.bias is not None
-
- # Check that weights are quantized
- assert qlinear.qweight.dtype == torch.uint8
- assert qlinear.scales.dtype == linear.weight.dtype
-
- def test_from_module_without_bias(self):
- """Test creating QuantLinear from nn.Linear without bias."""
- linear = nn.Linear(128, 256, bias=False)
- linear.weight.data.normal_(0, 0.02)
-
- qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32)
-
- assert qlinear.bias is None
-
- def test_from_module_preserves_bias(self):
- """Test that bias is preserved when converting."""
- linear = nn.Linear(128, 256)
- linear.weight.data.normal_(0, 0.02)
- linear.bias.data.fill_(0.5)
-
- qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32)
-
- assert torch.all(qlinear.bias == 0.5)
-
- @pytest.mark.parametrize("bits", [2, 4, 8])
- @pytest.mark.parametrize("symmetric", [True, False])
- def test_from_module_quantization_accuracy(self, bits, symmetric):
- """Test that quantization/dequantization is reasonably accurate."""
- linear = nn.Linear(64, 32, bias=False)
- linear.weight.data.normal_(0, 0.02)
-
- qlinear = QuantLinear.from_module(linear, bits=bits, symmetric=symmetric, group_size=16)
-
- # Dequantize and compare
- dequantized = qlinear.unpack_and_dequantize(qlinear.qweight, qlinear.scales, qlinear.qzeros)
-
- # Check shapes match
- assert dequantized.shape == linear.weight.shape
-
- # Check that dequantized weights are close to original (within quantization error)
- # The tolerance depends on the bit width and data distribution
- max_diff = torch.max(torch.abs(dequantized - linear.weight))
- # For 4-bit, we expect larger errors than 8-bit
- tolerance = 0.1 if bits == 4 else 0.05
- assert max_diff < tolerance, f"Max difference {max_diff} exceeds tolerance {tolerance}"
-
- def test_forward_shape(self):
- """Test that forward pass produces correct output shape."""
- linear = nn.Linear(128, 256)
- linear.weight.data.normal_(0, 0.02)
-
- qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=32)
-
- x = torch.randn(16, 128)
- output = qlinear(x)
-
- assert output.shape == (16, 256)
-
- def test_forward_batch_shape(self):
- """Test forward pass with batched input."""
- linear = nn.Linear(64, 128)
- linear.weight.data.normal_(0, 0.02)
-
- qlinear = QuantLinear.from_module(linear, bits=4, symmetric=True, group_size=16)
-
- x = torch.randn(8, 32, 64)
- output = qlinear(x)
-
- assert output.shape == (8, 32, 128)
-
- def test_forward_invalid_shape(self):
- """Test that forward raises error for invalid input shape."""
- qlinear = QuantLinear(in_features=128, out_features=256, bits=4, symmetric=True, group_size=32)
-
- x = torch.randn(16, 64) # Wrong input features
-
- with pytest.raises(AssertionError, match=r"Input shape .* does not match in_features"):
- qlinear(x)
-
- def test_from_tensors_with_precomputed_params(self):
- """Test creating QuantLinear with precomputed scales and zero points."""
- weight = torch.randn(256, 128)
-
- # Compute quantization parameters
- quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False)
- scales, zero_points = quantizer.find_qparams(weight)
-
- # Create QuantLinear with precomputed params
- qlinear = QuantLinear.from_tensors(
- in_features=128,
- out_features=256,
- weight=weight,
- bits=4,
- symmetric=True,
- group_size=32,
- scales=scales,
- zero_points=zero_points,
- bias=False,
- )
-
- assert qlinear.scales.shape == scales.shape
- assert torch.all(qlinear.scales == scales)
-
- def test_from_tensors_quantized_weights(self):
- """Test creating QuantLinear from already quantized weights."""
- weight = torch.randn(256, 128)
-
- # Quantize weights
- quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False)
- scales, zero_points = quantizer.find_qparams(weight)
- qweight = quantizer.quantize(weight, scales, zero_points)
-
- # Create QuantLinear from quantized weights
- qlinear = QuantLinear.from_tensors(
- in_features=128,
- out_features=256,
- weight=qweight,
- bits=4,
- symmetric=True,
- group_size=32,
- scales=scales,
- zero_points=zero_points,
- quantized=True,
- bias=False,
- )
-
- assert qlinear.cols == 128
- assert qlinear.rows == 256
-
- def test_from_tensors_missing_params_for_quantized(self):
- """Test that error is raised when params are missing for quantized weights."""
- qweight = torch.randint(0, 16, (256, 128))
-
- with pytest.raises(ValueError, match="scales and/or zero_points missing"):
- QuantLinear.from_tensors(
- in_features=128,
- out_features=256,
- weight=qweight,
- bits=4,
- symmetric=True,
- group_size=32,
- quantized=True,
- bias=False,
- )
-
- def test_from_tensors_invalid_symmetric_zero_points(self):
- """Test that error is raised for invalid zero points in symmetric quantization."""
- weight = torch.randn(256, 128)
-
- quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False)
- scales, zero_points = quantizer.find_qparams(weight)
-
- # Modify zero points to be invalid for symmetric quantization
- zero_points.fill_(0)
-
- with pytest.raises(ValueError, match="Zero points must be equal to midq for symmetric quantization"):
- QuantLinear.from_tensors(
- in_features=128,
- out_features=256,
- weight=weight,
- bits=4,
- symmetric=True,
- group_size=32,
- scales=scales,
- zero_points=zero_points,
- bias=False,
- )
-
- def test_extra_repr(self):
- """Test string representation of QuantLinear."""
- qlinear = QuantLinear(
- in_features=128,
- out_features=256,
- bits=4,
- symmetric=True,
- group_size=32,
- bias=True,
- )
-
- repr_str = qlinear.extra_repr()
- assert "in_features=128" in repr_str
- assert "out_features=256" in repr_str
- assert "bits=4" in repr_str
- assert "symmetric=True" in repr_str
- assert "group_size=32" in repr_str
- assert "bias=True" in repr_str
-
-
-class TestQuantEmbedding:
- def test_basic_initialization(self):
- """Test basic QuantEmbedding initialization."""
- qembed = QuantEmbedding(
- num_embeddings=1000,
- embedding_dim=256,
- bits=4,
- symmetric=True,
- group_size=32,
- )
- assert qembed.rows == 1000
- assert qembed.cols == 256
- assert qembed.padding_idx is None
-
- def test_initialization_with_padding_idx(self):
- """Test QuantEmbedding initialization with padding_idx."""
- qembed = QuantEmbedding(
- num_embeddings=1000,
- embedding_dim=256,
- bits=4,
- symmetric=True,
- group_size=32,
- padding_idx=0,
- )
- assert qembed.padding_idx == 0
-
- def test_from_module_basic(self):
- """Test creating QuantEmbedding from nn.Embedding."""
- embedding = nn.Embedding(1000, 256)
- embedding.weight.data.normal_(0, 0.02)
-
- qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32)
-
- assert qembed.rows == 1000
- assert qembed.cols == 256
- assert qembed.padding_idx is None
-
- def test_from_module_with_padding_idx(self):
- """Test creating QuantEmbedding with padding_idx."""
- embedding = nn.Embedding(1000, 256, padding_idx=0)
- embedding.weight.data.normal_(0, 0.02)
-
- qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32)
-
- assert qembed.padding_idx == 0
-
- @pytest.mark.parametrize("bits", [2, 4, 8])
- @pytest.mark.parametrize("symmetric", [True, False])
- def test_from_module_quantization_accuracy(self, bits, symmetric):
- """Test that quantization/dequantization is reasonably accurate."""
- embedding = nn.Embedding(100, 64)
- embedding.weight.data.normal_(0, 0.02)
-
- qembed = QuantEmbedding.from_module(embedding, bits=bits, symmetric=symmetric, group_size=16)
-
- # Dequantize and compare
- dequantized = qembed.unpack_and_dequantize(qembed.qweight, qembed.scales, qembed.qzeros)
-
- # Check shapes match
- assert dequantized.shape == embedding.weight.shape
-
- # Check that dequantized weights are close to original
- max_diff = torch.max(torch.abs(dequantized - embedding.weight))
- tolerance = 0.1 if bits == 4 else 0.05
- assert max_diff < tolerance, f"Max difference {max_diff} exceeds tolerance {tolerance}"
-
- def test_forward_shape(self):
- """Test that forward pass produces correct output shape."""
- embedding = nn.Embedding(1000, 256)
- embedding.weight.data.normal_(0, 0.02)
-
- qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32)
-
- x = torch.randint(0, 1000, (16,))
- output = qembed(x)
-
- assert output.shape == (16, 256)
-
- def test_forward_2d_input(self):
- """Test forward pass with 2D input."""
- embedding = nn.Embedding(1000, 256)
- embedding.weight.data.normal_(0, 0.02)
-
- qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32)
-
- x = torch.randint(0, 1000, (8, 32))
- output = qembed(x)
-
- assert output.shape == (8, 32, 256)
-
- def test_forward_3d_input(self):
- """Test forward pass with 3D input."""
- embedding = nn.Embedding(1000, 256)
- embedding.weight.data.normal_(0, 0.02)
-
- qembed = QuantEmbedding.from_module(embedding, bits=4, symmetric=True, group_size=32)
-
- x = torch.randint(0, 1000, (4, 8, 32))
- output = qembed(x)
-
- assert output.shape == (4, 8, 32, 256)
-
- def test_from_tensors_with_precomputed_params(self):
- """Test creating QuantEmbedding with precomputed scales and zero points."""
- weight = torch.randn(1000, 256)
-
- # Compute quantization parameters
- quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False)
- scales, zero_points = quantizer.find_qparams(weight)
-
- # Create QuantEmbedding with precomputed params
- qembed = QuantEmbedding.from_tensors(
- num_embeddings=1000,
- embedding_dim=256,
- weight=weight,
- bits=4,
- symmetric=True,
- group_size=32,
- scales=scales,
- zero_points=zero_points,
- )
-
- assert qembed.scales.shape == scales.shape
- assert torch.all(qembed.scales == scales)
-
- def test_from_tensors_quantized_weights(self):
- """Test creating QuantEmbedding from already quantized weights."""
- weight = torch.randn(1000, 256)
-
- # Quantize weights
- quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32, signed=False)
- scales, zero_points = quantizer.find_qparams(weight)
- qweight = quantizer.quantize(weight, scales, zero_points)
-
- # Create QuantEmbedding from quantized weights
- qembed = QuantEmbedding.from_tensors(
- num_embeddings=1000,
- embedding_dim=256,
- weight=qweight,
- bits=4,
- symmetric=True,
- group_size=32,
- scales=scales,
- zero_points=zero_points,
- quantized=True,
- )
-
- assert qembed.rows == 1000
- assert qembed.cols == 256
-
- def test_extra_repr(self):
- """Test string representation of QuantEmbedding."""
- qembed = QuantEmbedding(
- num_embeddings=1000,
- embedding_dim=256,
- bits=4,
- symmetric=True,
- group_size=32,
- padding_idx=0,
- )
-
- repr_str = qembed.extra_repr()
- assert "1000" in repr_str
- assert "256" in repr_str
- assert "bits=4" in repr_str
- assert "symmetric=True" in repr_str
- assert "group_size=32" in repr_str
- assert "padding_idx=0" in repr_str
-
-
-class TestQuantModuleIntegration:
- def test_quantlinear_forward_backward_compatibility(self):
- """Test that QuantLinear produces similar results to nn.Linear."""
- # Create a linear layer
- linear = nn.Linear(128, 64, bias=True)
- linear.weight.data.normal_(0, 0.02)
- linear.bias.data.zero_()
-
- # Create quantized version
- qlinear = QuantLinear.from_module(linear, bits=8, symmetric=True, group_size=32)
-
- # Test with same input
- x = torch.randn(16, 128)
-
- # Forward pass
- linear_out = linear(x)
- qlinear_out = qlinear(x)
-
- # Outputs should be close (not exact due to quantization)
- # For 8-bit, we expect decent accuracy
- max_diff = torch.max(torch.abs(linear_out - qlinear_out))
- relative_error = max_diff / torch.max(torch.abs(linear_out))
- assert relative_error < 0.1, f"Relative error {relative_error} too large"
-
- def test_quantembedding_forward_backward_compatibility(self):
- """Test that QuantEmbedding produces similar results to nn.Embedding."""
- # Create an embedding layer
- embedding = nn.Embedding(100, 64)
- embedding.weight.data.normal_(0, 0.02)
-
- # Create quantized version
- qembed = QuantEmbedding.from_module(embedding, bits=8, symmetric=True, group_size=16)
-
- # Test with same input
- x = torch.randint(0, 100, (16,))
-
- # Forward pass
- embed_out = embedding(x)
- qembed_out = qembed(x)
-
- # Outputs should be close (not exact due to quantization)
- max_diff = torch.max(torch.abs(embed_out - qembed_out))
- relative_error = max_diff / torch.max(torch.abs(embed_out))
- assert relative_error < 0.1, f"Relative error {relative_error} too large"
-
- def test_device_placement(self):
- """Test that QuantLinear respects device placement."""
- if not torch.cuda.is_available():
- pytest.skip("CUDA not available")
-
- device = torch.device("cuda")
- qlinear = QuantLinear(
- in_features=128,
- out_features=256,
- bits=4,
- symmetric=True,
- group_size=32,
- device=device,
- )
-
- assert qlinear.qweight.device.type == device.type
- assert qlinear.scales.device.type == device.type
- if qlinear.bias is not None:
- assert qlinear.bias.device.type == device.type
-
- def test_dtype_consistency(self):
- """Test that QuantLinear maintains dtype consistency."""
- dtype = torch.float16
- qlinear = QuantLinear(
- in_features=128,
- out_features=256,
- bits=4,
- symmetric=True,
- group_size=32,
- dtype=dtype,
- )
-
- assert qlinear.scales.dtype == dtype
- if qlinear.bias is not None:
- assert qlinear.bias.dtype == dtype
diff --git a/test/common/quant/test_utils.py b/test/common/quant/test_utils.py
index 802eab746d..a69b1d71f9 100644
--- a/test/common/quant/test_utils.py
+++ b/test/common/quant/test_utils.py
@@ -98,7 +98,7 @@ def test_get_num_groups_invalid_divisibility(self):
quantizer = WeightQuantizer(bits=4, symmetric=True, group_size=32)
shape = (64, 100) # 100 is not divisible by 32
- with pytest.raises(AssertionError, match=r"in_features .* must be divisible by group_size"):
+ with pytest.raises(AssertionError, match=r"last dim .* must be divisible by group_size"):
quantizer.get_num_groups(shape)
@pytest.mark.parametrize("group_size", [0, 16, -1])
@@ -399,3 +399,49 @@ def test_pack_unpack_device_consistency(self, device):
unpacked = unpack_from_uint8(packed, bits=4, shape=tensor.shape)
assert unpacked.device.type == device.type
assert torch.all(unpacked == tensor.to(torch.int32))
+
+
+class TestNDimensional:
+ """Verify that the quantizer / pack helpers operate identically on N-D tensors,
+ always quantizing along the last dim, without an explicit leading-dim loop.
+ """
+
+ @pytest.mark.parametrize("bits", [2, 4, 8])
+ @pytest.mark.parametrize("group_size", [-1, 16, 32])
+ def test_quantizer_3d_matches_2d_per_slice(self, bits, group_size):
+ """A 3D quantize must match independently quantizing each ``[i]`` slice."""
+ torch.manual_seed(0)
+ weight = torch.randn(3, 8, 64)
+ quantizer = WeightQuantizer(bits=bits, symmetric=False, group_size=group_size)
+
+ scales, zero_points = quantizer.find_qparams(weight)
+ q = quantizer.quantize(weight, scales, zero_points)
+ dq = quantizer.dequantize(q, scales, zero_points)
+
+ # quantization parameters carry shape (num_experts, out, num_groups)
+ num_groups = 1 if group_size == -1 else 64 // group_size
+ assert scales.shape == (3, 8, num_groups)
+ assert zero_points.shape == (3, 8, num_groups)
+ assert q.shape == weight.shape
+ assert dq.shape == weight.shape
+
+ for i in range(weight.shape[0]):
+ s_i, zp_i = quantizer.find_qparams(weight[i])
+ assert torch.equal(scales[i], s_i)
+ assert torch.equal(zero_points[i], zp_i)
+ assert torch.equal(q[i], quantizer.quantize(weight[i], s_i, zp_i))
+
+ @pytest.mark.parametrize("bits", [2, 4, 8])
+ def test_pack_unpack_3d_round_trip(self, bits):
+ """``pack_to_uint8`` / ``unpack_from_uint8`` round-trip on 3D inputs."""
+ torch.manual_seed(0)
+ shape = (3, 8, 64)
+ tensor = torch.randint(0, 2**bits, shape, dtype=torch.uint8)
+
+ packed = pack_to_uint8(tensor, bits)
+ packing_factor = 8 // bits
+ assert packed.shape == (3, 8, 64 // packing_factor)
+
+ unpacked = unpack_from_uint8(packed, bits, shape)
+ assert unpacked.shape == shape
+ assert torch.all(unpacked == tensor.to(torch.int32))
From 455ac53b1fe63eff8ea621d8daa012582aeb63df Mon Sep 17 00:00:00 2001
From: Jambay Kinley
Date: Fri, 15 May 2026 23:59:23 +0000
Subject: [PATCH 04/18] model_builder: fold Olive quant-suffix rename into
set_tensor
Instead of pre-walking the safetensors dict to rewrite
``.weight_qweight`` -> ``.qweight``, derive the
destination attribute name inside ``set_tensor`` once we already
know ``submodule`` is a ``QuantizedTensorModule``. Strip any of the
known Olive buffer suffixes (``QWEIGHT_SUFFIX``, ``SCALES_SUFFIX``,
``QZEROS_SUFFIX`` from ``olive.common.quant.state_dict``) from the
last path component to produce the bare ``qweight`` / ``scales`` /
``qzeros`` attribute that the genai ``QuantizedTensorModule``
expects.
Also drops internal dev-iteration version labels from comments and
docstrings in olive/common/quant and olive/passes/onnx.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/quant/hf_utils.py | 2 +-
olive/common/quant/state_dict.py | 2 +-
olive/passes/onnx/model_builder.py | 31 +++++++++++++++---------------
3 files changed, 17 insertions(+), 18 deletions(-)
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index b6118ba6ce..dca1934b51 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -166,7 +166,7 @@ def sort_key(name: str) -> tuple:
class OliveHfQuantizer(HfQuantizer):
"""Olive quantizer.
- Layout (Design 3, see ``olive/common/quant/state_dict.py``):
+ Layout (see ``olive/common/quant/state_dict.py``):
* Each quantized weight is installed as ``nn.Parameter(QuantTensor)``
on the original host module (``nn.Linear``, ``nn.Embedding``, or an
diff --git a/olive/common/quant/state_dict.py b/olive/common/quant/state_dict.py
index fb09f6d9d2..c3a8aff187 100644
--- a/olive/common/quant/state_dict.py
+++ b/olive/common/quant/state_dict.py
@@ -4,7 +4,7 @@
# --------------------------------------------------------------------------
"""State-dict helpers for Olive's quantized weight representation.
-Olive's quantization layout (Design 3):
+Olive's quantization layout:
* The quantized state of a weight named ```` (typically
``"weight"`` for ``nn.Linear``/``nn.Embedding``, or e.g.
diff --git a/olive/passes/onnx/model_builder.py b/olive/passes/onnx/model_builder.py
index 35a4f7d5cd..083957c147 100644
--- a/olive/passes/onnx/model_builder.py
+++ b/olive/passes/onnx/model_builder.py
@@ -435,6 +435,8 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
from onnxruntime_genai.models.quantized_model import QuantizedDecoderLayer, QuantizedTensorModule, TensorModule
from safetensors.torch import load_file
+ from olive.common.quant.state_dict import QWEIGHT_SUFFIX, QZEROS_SUFFIX, SCALES_SUFFIX
+
config = quant_attrs["config"]
if config.get("moe"):
@@ -487,38 +489,35 @@ def set_tensor(module, tensor_name, tensor_value, local_bits, local_group_size):
child = QuantizedTensorModule()
setattr(submodule, sub_name, child)
submodule = child
+
+ attr_name = tensor_name.split(".")[-1]
if isinstance(submodule, QuantizedTensorModule):
+ # Olive's native quantized checkpoints store buffers as
+ # ``_qweight`` / ``_scales`` / ``_qzeros`` (typically
+ # ``weight_*``); ``QuantizedTensorModule`` expects bare
+ # ``qweight`` / ``scales`` / ``qzeros`` attributes.
+ for suffix in (QWEIGHT_SUFFIX, SCALES_SUFFIX, QZEROS_SUFFIX):
+ if attr_name.endswith(suffix):
+ attr_name = suffix.lstrip("_")
+ break
+
for q_attr, q_value in [("bits", local_bits), ("_group_size", local_group_size)]:
setattr(submodule, q_attr, q_value)
# in_features is always a multiple of group_size, group_size is a power of 2
# assumes no padding
- if tensor_name.endswith("qweight"):
+ if attr_name == "qweight":
out_features, in_features_packed = tensor_value.shape
in_features = in_features_packed * 8 // local_bits
submodule.in_features = in_features
submodule.out_features = out_features
num_blocks = in_features // local_group_size if local_group_size != -1 else 1
tensor_value = tensor_value.reshape(out_features, num_blocks, -1)
- setattr(submodule, tensor_name.split(".")[-1], tensor_value)
+ setattr(submodule, attr_name, tensor_value)
for weight_file in Path(input_path).iterdir():
if weight_file.suffix == ".safetensors":
weights = load_file(weight_file)
- # Normalize Olive v4 key layout: ``.weight_qweight``
- # (and ``_scales`` / ``_qzeros``) → ``.qweight`` (etc.)
- # so the existing dotted-path resolver keeps working.
- normalized: dict[str, "torch.Tensor"] = {}
- for k, v in weights.items():
- for suffix in ("_qweight", "_scales", "_qzeros"):
- legacy_suffix = "." + suffix.lstrip("_")
- new_suffix = ".weight" + suffix
- if k.endswith(new_suffix):
- k = k[: -len(new_suffix)] + legacy_suffix
- break
- normalized[k] = v
- weights = normalized
-
# Map weights to modules
for name, tensor in weights.items():
if name.endswith("inv_freq"):
From 27458c06653bd425c8b1e5d739b1ecff003258c9 Mon Sep 17 00:00:00 2001
From: Jambay Kinley
Date: Sat, 16 May 2026 00:05:35 +0000
Subject: [PATCH 05/18] QuantLinearNbit.from_quant_tensor: reshape-only, no
unpack/repack
Both QuantTensor's 2D layout and QuantLinearNbit's MatMulNBits
buffer layout pack the quantization axis as uint8 with the same
in-byte order (low nibble = elem[2j], high nibble = elem[2j+1] for
4-bit, etc.). They differ only in qweight rank: QuantTensor uses
(out, in / pack_factor), QuantLinearNbit uses
(out, n_blocks, blob_size) where n_blocks * blob_size ==
in / pack_factor. So the conversion is a pure reshape; the previous
unpack -> .t() -> from_tensors round-trip is unnecessary.
scales and qzeros buffer shapes also match exactly between the two
layouts, so they are copied as-is. For symmetric weights
(QuantTensor.qzeros is None) we fill the QuantLinearNbit.qzeros
buffer with the packed midq pattern that the contrib op expects.
Verified numerically: F.linear via QuantTensor and the
dequantize-from-buffers path through QuantLinearNbit produce
bit-identical outputs across {4,8} bits, {symmetric, asymmetric},
{groupwise, per-channel}.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/quant.py | 65 +++++++++++++++++++++++-----------------
1 file changed, 38 insertions(+), 27 deletions(-)
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index 4eb8799f56..ac64d5d082 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -328,44 +328,55 @@ def from_quant_tensor(
) -> QuantLinearNbit:
"""Create a ``QuantLinearNbit`` from an Olive 2D ``QuantTensor``.
- The QuantTensor stores ``qweight`` packed-along-last (uint8) in
- ``(out_features, ceil(in_features / pack_factor))`` layout, with
- ``scales`` of shape ``(out_features, n_groups)`` and optional
- ``qzeros`` of shape ``(out_features, ceil(n_groups / pack_factor))``.
- We unpack to dense uint8 / int8 tensors then re-pack into the
- ``MatMulNBits`` layout via :meth:`pack`.
+ Both layouts pack the quantization axis (the input dim) as uint8
+ with the same in-byte order, so the buffer **bytes** are
+ bit-identical and we only need to reshape:
+
+ * ``QuantTensor.qweight`` ``(out, in / pack_factor)`` ->
+ ``QuantLinearNbit.qweight`` ``(out, n_blocks, blob_size)`` where
+ ``n_blocks * blob_size == in / pack_factor``.
+ * ``scales`` already match in shape (``(out, n_blocks)``).
+ * ``qzeros`` already match in shape; for symmetric weights
+ ``QuantTensor.qzeros`` is ``None`` so we fill the buffer with
+ the packed midq pattern that the contrib op expects.
"""
from olive.common.quant.tensor import QuantTensor
- from olive.common.quant.utils import unpack_from_uint8
if not isinstance(qt, QuantTensor) or qt.dim() != 2:
raise ValueError("QuantLinearNbit.from_quant_tensor requires a 2D QuantTensor")
out_features, in_features = qt.shape
- # unpack to (out, in) uint8 [0, 2^bits - 1]
- iweight = unpack_from_uint8(qt.qweight, qt.bits, (out_features, in_features))
- # unpack zero points; symmetric => midq
- if qt.qzeros is not None:
- izeros = unpack_from_uint8(qt.qzeros, qt.bits, tuple(qt.scales.shape))
- else:
- izeros = torch.full(
- tuple(qt.scales.shape),
- 1 << (qt.bits - 1),
- dtype=torch.uint8,
- device=qt.qweight.device,
- )
-
- # ``from_tensors`` expects iweight as (in, out) and scales/izeros as (n_groups, out)
- return cls.from_tensors(
- iweight=iweight.t().contiguous(),
- scales=qt.scales.t().contiguous(),
- izeros=izeros.t().contiguous(),
- group_size=qt.group_size if qt.group_size > 0 else in_features,
+ group_size = qt.group_size if qt.group_size > 0 else in_features
+ new = cls(
+ group_size=group_size,
+ in_features=in_features,
+ out_features=out_features,
+ g_idx=False,
+ bias=bias is not None,
bits=qt.bits,
- bias=bias,
+ dtype=qt.scales.dtype,
dynamo=dynamo,
)
+ # bit-identical layouts -> reshape (no unpack/repack)
+ new.qweight = qt.qweight.detach().clone().reshape(new.qweight.shape).contiguous()
+ new.scales = qt.scales.detach().clone().reshape(new.scales.shape).contiguous()
+ if qt.qzeros is not None:
+ new.qzeros = qt.qzeros.detach().clone().reshape(new.qzeros.shape).contiguous()
+ else:
+ # symmetric: fill the qzeros buffer with packed midq so the
+ # contrib op sees zp=midq everywhere.
+ midq = 1 << (qt.bits - 1)
+ packing_factor = 8 // qt.bits
+ packed_midq = 0
+ for i in range(packing_factor):
+ packed_midq |= midq << (qt.bits * i)
+ new.qzeros = torch.full_like(new.qzeros, packed_midq)
+
+ if bias is not None:
+ new.bias = bias.detach().clone()
+ return new
+
class QuantEmbeddingTorchFunction(torch.autograd.Function):
"""Export a quantized embedding lookup as ``com.microsoft::GatherBlockQuantized``."""
From 6cbacadc5eb52f7674a8d2d1d394a4775388ce55 Mon Sep 17 00:00:00 2001
From: Jambay Kinley
Date: Sat, 16 May 2026 00:23:15 +0000
Subject: [PATCH 06/18] Make qzeros optional on QuantLinearNbit /
QuantEmbeddingNbit
The ORT contrib MatMulNBits / GatherBlockQuantized ops treat a missing
zero-points input as midq for unsigned quantization, matching Olive's
symmetric-quantization convention. Drop the synthetic packed-midq
buffer that was previously emitted for symmetric weights and instead
omit the input entirely:
* QuantLinearNbit gains a has_qzeros flag (default True for back-compat);
pack/from_tensors/from_quant_tensor pass through None as needed.
* QuantLinearTorchFunction (TorchScript + dynamo) skips the qzeros input
when None, inserting an empty placeholder only when g_idx must be
positionally aligned.
* QuantEmbeddingTorchFunction.symbolic gains the missing dynamo arg
exposed by the new symmetric-embedding export path.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/quant.py | 80 ++++++++++++++++++++++++++--------------
1 file changed, 52 insertions(+), 28 deletions(-)
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index ac64d5d082..b6ea918cce 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -118,7 +118,13 @@ class QuantLinearTorchFunction(torch.autograd.Function):
# pylint: disable=W0223,W0221
@staticmethod
def symbolic(g, x, qweight, scales, qzeros, g_idx, bits, group_size, in_features, out_features, dynamo):
- tensor_args = [x, qweight, scales, qzeros]
+ tensor_args = [x, qweight, scales]
+ if qzeros is not None:
+ tensor_args.append(qzeros)
+ elif g_idx is not None:
+ # MatMulNBits inputs are positional; insert a missing-input
+ # placeholder for qzeros so g_idx lands at the correct slot.
+ tensor_args.append(g.op("Constant", value_t=torch.tensor([], dtype=torch.uint8)))
if g_idx is not None:
tensor_args.append(g_idx)
attrs = {
@@ -160,7 +166,11 @@ def forward(
if torch.onnx.is_in_onnx_export():
if dynamo and hasattr(torch.onnx, "ops"):
# torch.onnx.ops was introduced in 2.8
- tensor_args = [x, qweight, scales, qzeros]
+ tensor_args = [x, qweight, scales]
+ if qzeros is not None:
+ tensor_args.append(qzeros)
+ elif g_idx is not None:
+ tensor_args.append(torch.zeros(0, dtype=torch.uint8))
if g_idx is not None:
tensor_args.append(g_idx)
attrs = {
@@ -202,6 +212,7 @@ def __init__(
bits: int = 4,
dtype: torch.dtype = torch.float32,
dynamo: bool = False,
+ has_qzeros: bool = True,
):
super().__init__()
self.in_features = in_features
@@ -220,10 +231,16 @@ def __init__(
dtype=torch.uint8,
),
)
- self.register_buffer(
- "qzeros",
- torch.zeros(out_features, math.ceil(n_blocks_per_col * self.bits / 8), dtype=torch.uint8),
- )
+ if has_qzeros:
+ self.register_buffer(
+ "qzeros",
+ torch.zeros(out_features, math.ceil(n_blocks_per_col * self.bits / 8), dtype=torch.uint8),
+ )
+ else:
+ # When qzeros is omitted the contrib op interprets the
+ # zero point as the unsigned mid value (e.g. 8 for 4-bit),
+ # which matches Olive's symmetric quantization convention.
+ self.qzeros = None
self.register_buffer("scales", torch.zeros((out_features, n_blocks_per_col), dtype=dtype))
if g_idx:
self.register_buffer(
@@ -255,8 +272,11 @@ def forward(self, x):
def pack(self, iweight, izeros, scales, g_idx=None):
"""Pack int8 weight and zeros to int4 and int8 respectively.
- iweight, izeros and scales must have out_features as the first dimension
- iweight, izeros must be uint8 tensors with each holding 4/8 bit values
+ ``iweight`` and ``scales`` must have ``out_features`` as the
+ first dimension. ``iweight`` must be a uint8 tensor with each
+ value holding ``bits`` bits. ``izeros`` may be ``None`` for
+ symmetric quantization (the contrib op then treats the zero
+ point as the unsigned mid value).
"""
# pylint: disable=W0201
# shapes for packing
@@ -277,12 +297,15 @@ def pack(self, iweight, izeros, scales, g_idx=None):
iweight = (iweight[:, 0::2] & 0xF) | ((iweight[:, 1::2] & 0xF) << 4)
self.qweight = iweight.reshape(n, k_blocks, blob_size).contiguous()
- # pad to make the K dimension even
- izeros = torch.nn.functional.pad(izeros, (0, izeros.shape[-1] & 1), value=0)
- # pack the zeros
- if bits == 4:
- izeros = (izeros[:, 0::2] & 0xF) | ((izeros[:, 1::2] & 0xF) << 4)
- self.qzeros = izeros.contiguous()
+ if izeros is not None:
+ # pad to make the K dimension even
+ izeros = torch.nn.functional.pad(izeros, (0, izeros.shape[-1] & 1), value=0)
+ # pack the zeros
+ if bits == 4:
+ izeros = (izeros[:, 0::2] & 0xF) | ((izeros[:, 1::2] & 0xF) << 4)
+ self.qzeros = izeros.contiguous()
+ else:
+ self.qzeros = None
self.scales = scales.contiguous()
@@ -293,7 +316,7 @@ def from_tensors(
cls,
iweight: torch.Tensor,
scales: torch.Tensor,
- izeros: torch.Tensor,
+ izeros: torch.Tensor | None,
group_size: int,
bits: int = 4,
g_idx: torch.Tensor | None = None,
@@ -303,6 +326,8 @@ def from_tensors(
"""Create a QuantLinearNbit instance from the given tensors.
Weight is expected to be in in_features x out_features layout and unsigned.
+ ``izeros`` may be ``None`` for symmetric quantization — the
+ contrib op then assumes a midq zero point.
"""
new_qlinear = cls(
group_size,
@@ -313,8 +338,14 @@ def from_tensors(
bits=bits,
dtype=scales.dtype,
dynamo=dynamo,
+ has_qzeros=izeros is not None,
+ )
+ new_qlinear.pack(
+ iweight.to(torch.uint8).t(),
+ izeros.to(torch.uint8).t() if izeros is not None else None,
+ scales.t(),
+ g_idx,
)
- new_qlinear.pack(iweight.to(torch.uint8).t(), izeros.to(torch.uint8).t(), scales.t(), g_idx)
if bias is not None:
new_qlinear.bias = bias.clone()
return new_qlinear
@@ -337,8 +368,9 @@ def from_quant_tensor(
``n_blocks * blob_size == in / pack_factor``.
* ``scales`` already match in shape (``(out, n_blocks)``).
* ``qzeros`` already match in shape; for symmetric weights
- ``QuantTensor.qzeros`` is ``None`` so we fill the buffer with
- the packed midq pattern that the contrib op expects.
+ ``QuantTensor.qzeros`` is ``None`` and the contrib op
+ interprets the missing input as a midq zero point, so the
+ export module simply omits the buffer.
"""
from olive.common.quant.tensor import QuantTensor
@@ -356,6 +388,7 @@ def from_quant_tensor(
bits=qt.bits,
dtype=qt.scales.dtype,
dynamo=dynamo,
+ has_qzeros=qt.qzeros is not None,
)
# bit-identical layouts -> reshape (no unpack/repack)
@@ -363,15 +396,6 @@ def from_quant_tensor(
new.scales = qt.scales.detach().clone().reshape(new.scales.shape).contiguous()
if qt.qzeros is not None:
new.qzeros = qt.qzeros.detach().clone().reshape(new.qzeros.shape).contiguous()
- else:
- # symmetric: fill the qzeros buffer with packed midq so the
- # contrib op sees zp=midq everywhere.
- midq = 1 << (qt.bits - 1)
- packing_factor = 8 // qt.bits
- packed_midq = 0
- for i in range(packing_factor):
- packed_midq |= midq << (qt.bits * i)
- new.qzeros = torch.full_like(new.qzeros, packed_midq)
if bias is not None:
new.bias = bias.detach().clone()
@@ -383,7 +407,7 @@ class QuantEmbeddingTorchFunction(torch.autograd.Function):
# pylint: disable=W0223,W0221
@staticmethod
- def symbolic(g, x, qweight, scales, qzeros, bits, group_size, embedding_dim):
+ def symbolic(g, x, qweight, scales, qzeros, bits, group_size, embedding_dim, dynamo):
tensor_args = [qweight, x, scales]
if qzeros is not None:
tensor_args.append(qzeros)
From d758ef3fd9aec34a5ac8002c58bb14fa34f9064f Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 00:38:16 +0000
Subject: [PATCH 07/18] Drop placeholder Constant for qzeros in MatMulNBits
symbolic
g_idx alongside a missing qzeros is not a real combination in Olive
(GPTQ always produces qzeros), so skip the empty-tensor placeholder
and just omit qzeros from the input list entirely.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/quant.py | 10 +---------
1 file changed, 1 insertion(+), 9 deletions(-)
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index b6ea918cce..b429367484 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -121,10 +121,6 @@ def symbolic(g, x, qweight, scales, qzeros, g_idx, bits, group_size, in_features
tensor_args = [x, qweight, scales]
if qzeros is not None:
tensor_args.append(qzeros)
- elif g_idx is not None:
- # MatMulNBits inputs are positional; insert a missing-input
- # placeholder for qzeros so g_idx lands at the correct slot.
- tensor_args.append(g.op("Constant", value_t=torch.tensor([], dtype=torch.uint8)))
if g_idx is not None:
tensor_args.append(g_idx)
attrs = {
@@ -169,8 +165,6 @@ def forward(
tensor_args = [x, qweight, scales]
if qzeros is not None:
tensor_args.append(qzeros)
- elif g_idx is not None:
- tensor_args.append(torch.zeros(0, dtype=torch.uint8))
if g_idx is not None:
tensor_args.append(g_idx)
attrs = {
@@ -452,9 +446,7 @@ def forward(
version=1,
)
if dynamo:
- raise NotImplementedError(
- "torch dynamo export for quantized embedding requires torch 2.8 or higher."
- )
+ raise NotImplementedError("torch dynamo export for quantized embedding requires torch 2.8 or higher.")
return torch.zeros((*x.shape, embedding_dim), dtype=scales.dtype, device=x.device)
raise NotImplementedError("QuantEmbeddingTorchFunction forward is only implemented for onnx export")
From 46c8c7a64aeae0449920def6d90e4ea80963460e Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 00:52:30 +0000
Subject: [PATCH 08/18] Address PR feedback: factor quant target selection
- Replace duplicated model walks in hf_utils._process_model_before_weight_loading
and quant_utils.prepare_model with a shared iter_quant_targets helper that
returns a list of (module, dotted_name, param_name, shape, dtype, device,
kind) entries. Selection rules (lm_head/embeds/moe category flags, skip
patterns, extra_skip_modules, already-quantized) live in one place.
- QuantLinearNbit/QuantEmbeddingNbit: raise instead of synthesising a
placeholder when g_idx is supplied alongside symmetric quantization.
- tie_quant_word_embeddings: require both input and output embeddings to
already be QuantTensor-backed with matching shape/dtype before tying.
- Fix CodeQL mismatched-assignment false positives in QuantTensor dispatch
(index args directly), fix ruff D205/D401/PLW0108/A002 warnings.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/hf/quant.py | 13 +-
olive/common/quant/hf_utils.py | 125 +++++----------
olive/common/quant/selection.py | 212 +++++++++++++++++++++++++
olive/common/quant/tensor.py | 20 +--
olive/passes/pytorch/quant_utils.py | 95 ++++-------
test/common/quant/test_selection.py | 167 +++++++++++++++++++
test/common/quant/test_utils.py | 4 +-
test/passes/onnx/test_model_builder.py | 6 +-
8 files changed, 477 insertions(+), 165 deletions(-)
create mode 100644 olive/common/quant/selection.py
create mode 100644 test/common/quant/test_selection.py
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index b429367484..8d9ceb0048 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -121,6 +121,8 @@ def symbolic(g, x, qweight, scales, qzeros, g_idx, bits, group_size, in_features
tensor_args = [x, qweight, scales]
if qzeros is not None:
tensor_args.append(qzeros)
+ elif g_idx is not None:
+ raise ValueError("MatMulNBits with g_idx requires zero_points; symmetric quantization is not supported.")
if g_idx is not None:
tensor_args.append(g_idx)
attrs = {
@@ -165,6 +167,10 @@ def forward(
tensor_args = [x, qweight, scales]
if qzeros is not None:
tensor_args.append(qzeros)
+ elif g_idx is not None:
+ raise ValueError(
+ "MatMulNBits with g_idx requires zero_points; symmetric quantization is not supported."
+ )
if g_idx is not None:
tensor_args.append(g_idx)
attrs = {
@@ -630,9 +636,10 @@ def make_export_compatible_quant(model: torch.nn.Module, dynamo: bool) -> torch.
def _replace_olive_quant_tensor_modules(model: torch.nn.Module, dynamo: bool) -> None:
- """Swap host ``nn.Linear`` / ``nn.Embedding`` whose weight is a ``QuantTensor``
- with the ONNX-exportable :class:`QuantLinearNbit` / :class:`QuantEmbeddingNbit`
- wrappers.
+ """Swap host modules whose weight is a ``QuantTensor`` with export-compatible wrappers.
+
+ Targets ``nn.Linear`` / ``nn.Embedding`` and replaces them with
+ :class:`QuantLinearNbit` / :class:`QuantEmbeddingNbit`.
"""
from olive.common.quant.tensor import QuantTensor
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index dca1934b51..58f9866c78 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -13,8 +13,7 @@
from transformers.quantizers.base import HfQuantizer
from transformers.utils.quantization_config import QuantizationConfigMixin
-from olive.common.hf.wrapper import ModelWrapper
-from olive.common.quant.patterns import match_override, match_skip
+from olive.common.quant.patterns import match_override
from olive.common.quant.state_dict import buffer_names, install_quant_tensor_param, refresh_quant_tensor_refs
from olive.common.quant.tensor import QuantTensor
from olive.common.quant.utils import WeightQuantizer
@@ -184,96 +183,45 @@ class OliveHfQuantizer(HfQuantizer):
def _process_model_before_weight_loading(
self, model: PreTrainedModel, keep_in_fp32_modules: list[str] | None = None, **kwargs
):
- ids_to_skip: list[int] = []
- if not self.quantization_config.lm_head:
- ids_to_skip.append(id(model.get_output_embeddings()))
- if not self.quantization_config.embeds:
- ids_to_skip.append(id(model.get_input_embeddings()))
- expert_modules: list[tuple[nn.Module, str]] = []
- if not self.quantization_config.moe:
- # Skip every nn.Module under each experts subtree — this both
- # leaves fused-3D experts alone *and* fixes the previous silent
- # quantization of per-expert nn.Linears inside
- # ModuleList(Expert) blocks (Mixtral, PhiMoE, Qwen2/3-MoE).
- try:
- wrapper = ModelWrapper.from_model(model)
- for lw in wrapper.get_layer_wrappers():
- experts = lw.get_experts(return_name=False)
- if experts is None:
- continue
- for sub in experts.modules():
- ids_to_skip.append(id(sub))
- except Exception: # pylint: disable=broad-except
- # Not every model is wrappable (e.g., random tests). Falling
- # back to the previous behaviour (no experts skip) is safe.
- pass
- else:
- # collect (experts_module, dotted_name) so we can also install
- # 3D placeholders on fused-experts modules below.
- try:
- wrapper = ModelWrapper.from_model(model)
- module_to_name = {id(m): n for n, m in model.named_modules()}
- for lw in wrapper.get_layer_wrappers():
- experts = lw.get_experts(return_name=False)
- if experts is None:
- continue
- expert_modules.append((experts, module_to_name.get(id(experts), "")))
- except Exception: # pylint: disable=broad-except
- pass
-
- skip_literal_names: set[str] = (
- {name for name, module in model.named_modules() if id(module) in ids_to_skip} if ids_to_skip else set()
- )
- # Pattern-aware skip list. Plain strings keep their HF substring
- # semantics; ``re:`` opts into regex fullmatch.
+ from olive.common.quant.selection import iter_quant_targets
+
skip_patterns: list[str] = []
if self.quantization_config.modules_to_not_convert:
skip_patterns.extend(self.quantization_config.modules_to_not_convert)
if keep_in_fp32_modules:
skip_patterns.extend(keep_in_fp32_modules)
- self.modules_to_not_convert = sorted(skip_literal_names)
self._skip_patterns = skip_patterns
- def _should_quantize(module: nn.Module, name: str) -> bool:
- if not isinstance(module, (nn.Linear, nn.Embedding)):
- return False
- if name in skip_literal_names:
- return False
- return not match_skip(name, skip_patterns)
-
- # 2D placeholders: nn.Linear / nn.Embedding
- for name, module in list(model.named_modules()):
- if not _should_quantize(module, name):
- continue
- qargs = self.quantization_config.get_qlinear_init_args(name)
+ skip_literal_names: set[str] = set()
+ for target in iter_quant_targets(
+ model,
+ quantize_lm_head=self.quantization_config.lm_head,
+ quantize_embeds=self.quantization_config.embeds,
+ quantize_moe=self.quantization_config.moe,
+ skip_patterns=skip_patterns,
+ ):
+ qargs = self.quantization_config.get_qlinear_init_args(target.full_name)
qt = _build_placeholder_quant_tensor(
- shape=tuple(module.weight.shape),
+ shape=target.shape,
bits=qargs["bits"],
symmetric=qargs["symmetric"],
group_size=qargs["group_size"],
- dtype=module.weight.dtype,
- device=module.weight.device,
+ dtype=target.dtype,
+ device=target.device,
)
- install_quant_tensor_param(module, "weight", qt)
-
- # 3D placeholders: fused-experts modules
- for experts_module, experts_name in expert_modules:
- for pname, param in list(experts_module._parameters.items()):
- if param is None or param.dim() != 3 or isinstance(param.data, QuantTensor):
- continue
- full_name = f"{experts_name}.{pname}" if experts_name else pname
- if full_name in skip_literal_names or match_skip(full_name, skip_patterns):
- continue
- qargs = self.quantization_config.get_qlinear_init_args(full_name)
- qt = _build_placeholder_quant_tensor(
- shape=tuple(param.shape),
- bits=qargs["bits"],
- symmetric=qargs["symmetric"],
- group_size=qargs["group_size"],
- dtype=param.dtype,
- device=param.device,
- )
- install_quant_tensor_param(experts_module, pname, qt)
+ install_quant_tensor_param(target.module, target.param_name, qt)
+
+ # Record literal names of every nn.Linear/Embedding that was
+ # *not* converted so ``modules_to_not_convert`` reflects reality
+ # (mirrors HF quantizer conventions).
+ for name, module in model.named_modules():
+ if not isinstance(module, (nn.Linear, nn.Embedding)):
+ continue
+ if not isinstance(module._parameters.get("weight"), nn.Parameter) or not isinstance(
+ module._parameters["weight"].data, QuantTensor
+ ):
+ skip_literal_names.add(name)
+ self.modules_to_not_convert = sorted(skip_literal_names)
if self.quantization_config.tie_word_embeddings:
# doing first time so that the weight load doesn't complain about missing weights
@@ -397,22 +345,33 @@ def replace_matching_submodules(
def tie_quant_word_embeddings(model: PreTrainedModel) -> None:
- """Tie the input and output embeddings when they share a quantized weight.
+ """Tie the input and output embeddings when both share a quantized weight.
Both modules' ``weight`` ``nn.Parameter`` is set to the **same**
``nn.Parameter(QuantTensor)`` object, and the underlying
``weight_qweight`` / ``weight_scales`` / ``weight_qzeros`` buffers
are tied (aliased to the input embedding's buffers). This preserves
the standard HF tied-weights semantics for the quantized layout.
+
+ Tying is a no-op unless **both** the input and output embeddings
+ are already backed by compatible ``QuantTensor`` weights with
+ matching shape and dtype.
"""
src = model.get_input_embeddings()
dst = model.get_output_embeddings()
if src is None or dst is None:
return
- # The input embedding owns the canonical QuantTensor + buffers.
src_param = src._parameters.get("weight")
- if src_param is None or not isinstance(src_param.data, QuantTensor):
+ dst_param = dst._parameters.get("weight")
+ if (
+ src_param is None
+ or dst_param is None
+ or not isinstance(src_param.data, QuantTensor)
+ or not isinstance(dst_param.data, QuantTensor)
+ ):
+ return
+ if src_param.shape != dst_param.shape or src_param.dtype != dst_param.dtype:
return
qname, sname, zname = buffer_names("weight")
diff --git a/olive/common/quant/selection.py b/olive/common/quant/selection.py
new file mode 100644
index 0000000000..f376d33054
--- /dev/null
+++ b/olive/common/quant/selection.py
@@ -0,0 +1,212 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Quantization target selection.
+
+Centralises the logic that walks a model once and decides which
+parameters to quantize. Both Olive's HF quantizer (which installs
+:class:`QuantTensor` placeholders before weight loading) and the
+PyTorch RTN/GPTQ passes (which attach calibration metadata) consume
+the same set of targets — only the per-target action differs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+
+from olive.common.quant.patterns import match_skip
+
+if TYPE_CHECKING:
+ from collections.abc import Iterable, Iterator
+ from typing import Callable
+
+ from olive.common.hf.wrapper import ModelWrapper
+
+
+@dataclass
+class QuantTarget:
+ """A single parameter selected for quantization.
+
+ Attributes:
+ module: The owning ``nn.Module``.
+ module_name: Dotted name of ``module`` relative to the model root.
+ param_name: Name of the parameter on ``module`` (e.g. ``"weight"``
+ for ``nn.Linear``/``nn.Embedding``; an arbitrary 3D
+ parameter name for a fused-experts module).
+ full_name: ``f"{module_name}.{param_name}"`` (with an exception
+ for ``param_name == "weight"`` on ``nn.Linear``/``nn.Embedding``,
+ in which case ``full_name == module_name`` — matching how
+ overrides and skip patterns have always been keyed).
+ kind: ``"linear"`` for ``nn.Linear``, ``"embedding"`` for
+ ``nn.Embedding``, ``"fused_experts"`` for a 3D parameter on
+ an experts module.
+ shape: Logical (dequantized) shape of the parameter.
+ dtype: Original parameter dtype.
+ device: Original parameter device.
+
+ """
+
+ module: nn.Module
+ module_name: str
+ param_name: str
+ full_name: str
+ kind: str
+ shape: tuple[int, ...]
+ dtype: torch.dtype
+ device: torch.device
+
+
+def _collect_experts(
+ model: nn.Module,
+ wrapper: ModelWrapper | None,
+) -> list[tuple[nn.Module, str]]:
+ """Return ``(experts_module, dotted_name)`` for every MoE layer."""
+ if wrapper is None:
+ return []
+ out: list[tuple[nn.Module, str]] = []
+ for lw in wrapper.get_layer_wrappers():
+ experts, name = lw.get_experts(return_name=True)
+ if experts is not None:
+ out.append((experts, name))
+ return out
+
+
+def iter_quant_targets(
+ model: nn.Module,
+ *,
+ quantize_lm_head: bool,
+ quantize_embeds: bool,
+ quantize_moe: bool,
+ skip_patterns: Iterable[str] = (),
+ extra_skip_modules: Iterable[nn.Module] = (),
+ skip_already_quantized: bool = True,
+ consider_linears: bool = True,
+ consider_embeddings: bool = True,
+) -> Iterator[QuantTarget]:
+ """Walk ``model`` once and yield every parameter selected for quantization.
+
+ Selection rules (first matching skip wins):
+
+ * ``extra_skip_modules`` (caller-supplied set, e.g. attention inputs
+ excluded by GPTQ) skips the module by identity.
+ * ``quantize_lm_head=False`` skips the output embedding module.
+ * ``quantize_embeds=False`` skips the input embedding module.
+ * ``quantize_moe=False`` skips every ``nn.Module`` under any experts
+ subtree (this both leaves fused-3D experts alone and prevents
+ silently quantizing per-expert ``nn.Linear``s inside
+ ``ModuleList(Expert)`` blocks).
+ * ``skip_patterns`` matches the parameter's ``full_name`` via the
+ shared HF-style substring / ``re:``-prefixed regex matcher.
+ * When ``skip_already_quantized=True``, weights that are already a
+ :class:`QuantTensor` are skipped (idempotent re-runs).
+
+ When ``quantize_moe=True`` and an experts module exposes a 3D
+ ``nn.Parameter`` (fused experts), that parameter is yielded as a
+ ``"fused_experts"`` target. Per-expert 2D ``nn.Linear``s inside an
+ ``nn.ModuleList`` experts wrapper continue to come through the
+ regular linear walk.
+ """
+ from olive.common.hf.wrapper import ModelWrapper
+ from olive.common.quant.tensor import QuantTensor
+
+ try:
+ wrapper = ModelWrapper.from_model(model)
+ except Exception: # pylint: disable=broad-except
+ # Not every model is wrappable (e.g., random sklearn-like
+ # test fixtures). Without the wrapper we cannot honour MoE /
+ # lm_head / embeds category flags; fall back to the unfiltered
+ # 2D walk.
+ wrapper = None
+
+ lm_head_module: nn.Module | None = None
+ embed_module: nn.Module | None = None
+ if hasattr(model, "get_output_embeddings"):
+ lm_head_module = model.get_output_embeddings()
+ if hasattr(model, "get_input_embeddings"):
+ embed_module = model.get_input_embeddings()
+
+ expert_modules = _collect_experts(model, wrapper)
+
+ # ID-based skip set for fast identity checks during the named_modules walk.
+ skip_ids: set[int] = {id(m) for m in extra_skip_modules}
+ if not quantize_lm_head and lm_head_module is not None:
+ skip_ids.add(id(lm_head_module))
+ if not quantize_embeds and embed_module is not None:
+ skip_ids.add(id(embed_module))
+ if not quantize_moe:
+ for experts, _ in expert_modules:
+ for sub in experts.modules():
+ skip_ids.add(id(sub))
+
+ patterns = list(skip_patterns or ())
+
+ def _is_skipped(module: nn.Module, full_name: str) -> bool:
+ if id(module) in skip_ids:
+ return True
+ return bool(patterns) and match_skip(full_name, patterns)
+
+ # 2D pass: every nn.Linear / nn.Embedding under the model.
+ for name, module in model.named_modules():
+ if isinstance(module, nn.Linear) and consider_linears:
+ kind = "linear"
+ elif isinstance(module, nn.Embedding) and consider_embeddings:
+ kind = "embedding"
+ else:
+ continue
+ if _is_skipped(module, name):
+ continue
+ weight = module.weight
+ if skip_already_quantized and isinstance(weight, QuantTensor):
+ continue
+ if weight is None:
+ continue
+ yield QuantTarget(
+ module=module,
+ module_name=name,
+ param_name="weight",
+ full_name=name,
+ kind=kind,
+ shape=tuple(weight.shape),
+ dtype=weight.dtype,
+ device=weight.device,
+ )
+
+ # 3D pass: fused-experts modules only when MoE quantization is requested.
+ if not quantize_moe:
+ return
+ for experts_module, experts_name in expert_modules:
+ for pname, param in experts_module.named_parameters(recurse=False):
+ if param is None or param.dim() != 3:
+ continue
+ if skip_already_quantized and isinstance(param.data, QuantTensor):
+ continue
+ full_name = f"{experts_name}.{pname}" if experts_name else pname
+ if _is_skipped(experts_module, full_name):
+ continue
+ yield QuantTarget(
+ module=experts_module,
+ module_name=experts_name,
+ param_name=pname,
+ full_name=full_name,
+ kind="fused_experts",
+ shape=tuple(param.shape),
+ dtype=param.dtype,
+ device=param.device,
+ )
+
+
+def for_each_target(
+ model: nn.Module,
+ handler: Callable[[QuantTarget], None],
+ **selection_kwargs,
+) -> list[QuantTarget]:
+ """Run ``handler`` once per target and return the materialised list."""
+ targets = list(iter_quant_targets(model, **selection_kwargs))
+ for t in targets:
+ handler(t)
+ return targets
diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py
index 6c4628df7f..b81e32d7f3 100644
--- a/olive/common/quant/tensor.py
+++ b/olive/common/quant/tensor.py
@@ -2,8 +2,9 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
-"""``QuantTensor`` — a wrapper ``torch.Tensor`` subclass that stores weight
-quantization buffers (``qweight``, ``scales``, ``qzeros``) but presents the
+"""QuantTensor — wrapper ``torch.Tensor`` subclass for weight-quantized parameters.
+
+It stores quantization buffers (``qweight``, ``scales``, ``qzeros``) but presents the
shape / dtype / device of the dequantized full-precision weight.
Design notes:
@@ -51,7 +52,7 @@
def implements(*torch_fns: Callable) -> Callable[[Callable], Callable]:
- """Decorator to register a torch-function override for ``QuantTensor``."""
+ """Register a torch-function override for ``QuantTensor``."""
def decorator(fn: Callable) -> Callable:
for torch_fn in torch_fns:
@@ -304,11 +305,11 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
aten = torch.ops.aten
if func in (aten.detach.default, aten.clone.default, aten.alias.default, aten.contiguous.default):
- (self_,) = args
- return self_._apply_fn_to_data(lambda x: func(x))
+ self_ = args[0]
+ return self_._apply_fn_to_data(func)
if func is aten._to_copy.default:
- (self_,) = args
+ self_ = args[0]
dtype = kwargs.get("dtype")
device = kwargs.get("device")
@@ -324,7 +325,8 @@ def _move(x: torch.Tensor) -> torch.Tensor:
return self_._apply_fn_to_data(_move)
if func is aten.copy_.default:
- self_, src = args
+ self_ = args[0]
+ src = args[1]
if not isinstance(src, QuantTensor):
raise TypeError(f"Cannot copy_ a non-QuantTensor source into a QuantTensor (got {type(src)})")
self_.qweight.copy_(src.qweight)
@@ -358,7 +360,7 @@ def _maybe_dense(x: Any) -> Any:
@implements(F.linear)
-def _linear(input: torch.Tensor, weight: QuantTensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+def _linear(input: torch.Tensor, weight: QuantTensor, bias: torch.Tensor | None = None) -> torch.Tensor: # noqa: A002
if torch.onnx.is_in_onnx_export():
raise RuntimeError(
"Olive QuantTensor cannot be traced by torch.onnx.export directly. "
@@ -378,7 +380,7 @@ def _linear(input: torch.Tensor, weight: QuantTensor, bias: torch.Tensor | None
@implements(F.embedding)
def _embedding(
- input: torch.Tensor,
+ input: torch.Tensor, # noqa: A002
weight: QuantTensor,
padding_idx: int | None = None,
max_norm: float | None = None,
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index d026144fc6..8869fcb987 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -14,10 +14,9 @@
from olive.common.quant.hf_utils import (
OliveHfQuantizationConfig,
OliveHfQuantizationMethod,
- replace_matching_submodules,
tie_quant_word_embeddings,
)
-from olive.common.quant.patterns import match_skip
+from olive.common.quant.selection import iter_quant_targets
from olive.common.quant.state_dict import install_quant_tensor_param
from olive.common.quant.tensor import QuantTensor
from olive.common.quant.utils import WeightQuantizer
@@ -152,73 +151,35 @@ def prepare_model(
else:
excluded_attn_inputs.update(attn_inputs[:2])
- # Collect every ``nn.Module`` under any experts subtree, so we can
- # honour the ``moe`` category flag the same way ``lm_head`` /
- # ``embeds`` are honoured today.
- expert_module_ids: set[int] = set()
- expert_owners: list[tuple[torch.nn.Module, str]] = [] # (experts_module, dotted_name)
- for layer_wrapper in wrapper.get_layer_wrappers():
- experts_module = layer_wrapper.get_experts(return_name=False)
- if experts_module is None:
- continue
- # Find the dotted name of this experts module relative to the model root.
- experts_name = None
- for name, mod in wrapper.model.named_modules():
- if mod is experts_module:
- experts_name = name
- break
- expert_owners.append((experts_module, experts_name or ""))
- for sub in experts_module.modules():
- expert_module_ids.add(id(sub))
-
skip_patterns = list(getattr(qcfg, "modules_to_not_convert", None) or [])
- def should_quantize(module: torch.nn.Module, name: str) -> bool:
- if module in excluded_attn_inputs:
- return False
- # already-quantized (QuantTensor weight) — leave it alone so the user
- # can compose passes on top of a partially-quantized model.
- weight = getattr(module, "weight", None)
- if weight is not None and isinstance(weight, QuantTensor):
- return False
- if match_skip(name, skip_patterns):
- return False
- # category-flag skips (lm_head / embeds / moe) — first rule wins
- if id(module) in expert_module_ids and not getattr(qcfg, "moe", False):
- return False
- if isinstance(module, torch.nn.Linear):
- return name != lm_head_name or qcfg.lm_head
- if qcfg.embeds and isinstance(module, torch.nn.Embedding):
- return name == embeds_name
- return False
-
- def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module:
- # TODO(jambayk): validate that the module and config are compatible
- qargs = qcfg.get_qlinear_init_args(name)
- module.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs))
- new_qargs[name] = qargs
- return module
-
- replace_matching_submodules(wrapper.model, should_quantize, add_quant_info, description="Preparing model")
-
- # Fused-3D MoE: experts modules expose 3D nn.Parameters directly (e.g.
- # ``gate_up_proj`` of shape ``(num_experts, *, *)``). Annotate the
- # experts module with a per-parameter quant_info_3d dict so
- # ``finalize`` can replace each parameter with a QuantTensor.
- if getattr(qcfg, "moe", False):
- for experts_module, experts_name in expert_owners:
- param_qinfos: dict[str, QuantInfo] = {}
- for pname, param in experts_module.named_parameters(recurse=False):
- if param.dim() != 3:
- continue
- full_name = f"{experts_name}.{pname}" if experts_name else pname
- if match_skip(full_name, skip_patterns):
- continue
- qargs = qcfg.get_qlinear_init_args(full_name)
- param_qinfos[pname] = QuantInfo(quantizer=WeightQuantizer(**qargs))
- new_qargs[full_name] = qargs
- if param_qinfos:
- experts_module.quant_info_3d = param_qinfos
+ fused_targets_by_module: dict[int, dict[str, QuantInfo]] = {}
+ for target in iter_quant_targets(
+ wrapper.model,
+ quantize_lm_head=qcfg.lm_head,
+ quantize_embeds=qcfg.embeds,
+ quantize_moe=getattr(qcfg, "moe", False),
+ skip_patterns=skip_patterns,
+ extra_skip_modules=excluded_attn_inputs,
+ ):
+ qargs = qcfg.get_qlinear_init_args(target.full_name)
+ new_qargs[target.full_name] = qargs
+ if target.kind == "fused_experts":
+ fused_targets_by_module.setdefault(id(target.module), {})[target.param_name] = QuantInfo(
+ quantizer=WeightQuantizer(**qargs)
+ )
+ # Stash the experts module on the dict so we can install
+ # ``quant_info_3d`` below without re-walking.
+ fused_targets_by_module[id(target.module)]["__module__"] = target.module # type: ignore[assignment]
+ else:
+ target.module.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs))
+
+ # Fused-3D MoE: stamp the per-parameter QuantInfo dict onto each
+ # experts module so ``finalize`` can replace each parameter with a
+ # QuantTensor.
+ for entry in fused_targets_by_module.values():
+ module = entry.pop("__module__") # type: ignore[arg-type]
+ module.quant_info_3d = entry # type: ignore[assignment]
# remove overrides for modules not being quantized
for name in list(qcfg.overrides or {}):
diff --git a/test/common/quant/test_selection.py b/test/common/quant/test_selection.py
new file mode 100644
index 0000000000..d36dd197e3
--- /dev/null
+++ b/test/common/quant/test_selection.py
@@ -0,0 +1,167 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for ``olive.common.quant.selection.iter_quant_targets``."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from olive.common.quant.selection import iter_quant_targets
+
+
+class _Toy(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.embed_tokens = nn.Embedding(16, 8)
+ self.linear = nn.Linear(8, 8, bias=False)
+ self.lm_head = nn.Linear(8, 16, bias=False)
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ def get_output_embeddings(self):
+ return self.lm_head
+
+
+def _names(targets):
+ return sorted((t.full_name, t.kind) for t in targets)
+
+
+def test_default_skips_lm_head_and_embeds():
+ m = _Toy()
+ targets = list(iter_quant_targets(m, quantize_lm_head=False, quantize_embeds=False, quantize_moe=False))
+ assert _names(targets) == [("linear", "linear")]
+
+
+def test_include_lm_head_and_embeds():
+ m = _Toy()
+ targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False))
+ assert _names(targets) == [
+ ("embed_tokens", "embedding"),
+ ("linear", "linear"),
+ ("lm_head", "linear"),
+ ]
+
+
+def test_skip_patterns_filter_by_name():
+ m = _Toy()
+ targets = list(
+ iter_quant_targets(
+ m,
+ quantize_lm_head=True,
+ quantize_embeds=True,
+ quantize_moe=False,
+ skip_patterns=["re:.*_head"],
+ )
+ )
+ assert _names(targets) == [("embed_tokens", "embedding"), ("linear", "linear")]
+
+
+def test_extra_skip_modules_skip_by_identity():
+ m = _Toy()
+ targets = list(
+ iter_quant_targets(
+ m,
+ quantize_lm_head=True,
+ quantize_embeds=False,
+ quantize_moe=False,
+ extra_skip_modules={m.linear},
+ )
+ )
+ assert _names(targets) == [("lm_head", "linear")]
+
+
+def test_already_quantized_param_is_skipped():
+ from olive.common.quant.tensor import QuantTensor
+
+ m = _Toy()
+ # Build a tiny QuantTensor placeholder and stamp it onto m.linear.weight.
+ qt = QuantTensor.from_packed(
+ qweight=torch.zeros((8, 4), dtype=torch.uint8),
+ scales=torch.zeros((8, 1), dtype=torch.float32),
+ qzeros=None,
+ bits=4,
+ group_size=8,
+ symmetric=True,
+ shape=(8, 8),
+ dtype=torch.float32,
+ )
+ m.linear.weight = nn.Parameter(qt, requires_grad=False)
+
+ targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=False, quantize_moe=False))
+ assert _names(targets) == [("lm_head", "linear")]
+
+
+class _ExpertList(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.experts = nn.ModuleList([nn.Linear(8, 8, bias=False) for _ in range(2)])
+
+
+def test_moe_disabled_skips_submodules_under_experts(monkeypatch):
+ """When ``quantize_moe=False``, every nn.Module under an experts subtree is skipped."""
+ from olive.common.hf import wrapper as wrapper_mod
+
+ class FakeLayerWrapper:
+ def __init__(self, experts, name):
+ self._experts = experts
+ self._name = name
+
+ def get_experts(self, return_name=True):
+ return (self._experts, self._name) if return_name else self._experts
+
+ class FakeWrapper:
+ def __init__(self, model):
+ self.model = model
+
+ def get_layer_wrappers(self):
+ return [FakeLayerWrapper(self.model.experts, "experts")]
+
+ monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m)))
+
+ m = _ExpertList()
+ targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False))
+ assert _names(targets) == []
+
+
+def test_moe_enabled_yields_3d_fused_params(monkeypatch):
+ from olive.common.hf import wrapper as wrapper_mod
+
+ class FusedExperts(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False)
+ self.down_proj = nn.Parameter(torch.zeros(4, 16, 8), requires_grad=False)
+
+ class _Model(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.experts = FusedExperts()
+
+ class FakeLayerWrapper:
+ def __init__(self, experts, name):
+ self._experts = experts
+ self._name = name
+
+ def get_experts(self, return_name=True):
+ return (self._experts, self._name) if return_name else self._experts
+
+ class FakeWrapper:
+ def __init__(self, model):
+ self.model = model
+
+ def get_layer_wrappers(self):
+ return [FakeLayerWrapper(self.model.experts, "experts")]
+
+ monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m)))
+
+ m = _Model()
+ targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True))
+ fused = [(t.full_name, t.kind, t.shape) for t in targets if t.kind == "fused_experts"]
+ assert sorted(fused) == [
+ ("experts.down_proj", "fused_experts", (4, 16, 8)),
+ ("experts.gate_up_proj", "fused_experts", (4, 8, 16)),
+ ]
diff --git a/test/common/quant/test_utils.py b/test/common/quant/test_utils.py
index a69b1d71f9..d5004db897 100644
--- a/test/common/quant/test_utils.py
+++ b/test/common/quant/test_utils.py
@@ -402,7 +402,9 @@ def test_pack_unpack_device_consistency(self, device):
class TestNDimensional:
- """Verify that the quantizer / pack helpers operate identically on N-D tensors,
+ """N-D quantization tests.
+
+ Verify that the quantizer / pack helpers operate identically on N-D tensors,
always quantizing along the last dim, without an explicit leading-dim loop.
"""
diff --git a/test/passes/onnx/test_model_builder.py b/test/passes/onnx/test_model_builder.py
index 4f574063ec..cf5f394c8c 100644
--- a/test/passes/onnx/test_model_builder.py
+++ b/test/passes/onnx/test_model_builder.py
@@ -199,8 +199,10 @@ def fake_create_model(
def test_olive_quantized_model_raises_for_moe():
- """ModelBuilder cannot consume Olive-quantized MoE checkpoints; it must error out
- cleanly so the user reaches for Mobius (or re-runs RTN without moe=True).
+ """ModelBuilder must reject Olive-quantized MoE checkpoints.
+
+ Errors out cleanly so the user reaches for an alternative builder
+ or re-runs RTN without ``moe=True``.
"""
from olive.passes.onnx.model_builder import OliveQuantizedModel
From 4b66da5a397f049d8b5d1a00bc99bf657ae790fc Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 00:57:57 +0000
Subject: [PATCH 09/18] Drop dead writes to self.modules_to_not_convert /
self._skip_patterns
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Neither attribute is read anywhere — the post-walk loop that produced
the literal skip-name list was a leftover from before the refactor.
The configured patterns already live on quantization_config.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/quant/hf_utils.py | 14 --------------
1 file changed, 14 deletions(-)
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index 58f9866c78..fdc9d01959 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -190,9 +190,7 @@ def _process_model_before_weight_loading(
skip_patterns.extend(self.quantization_config.modules_to_not_convert)
if keep_in_fp32_modules:
skip_patterns.extend(keep_in_fp32_modules)
- self._skip_patterns = skip_patterns
- skip_literal_names: set[str] = set()
for target in iter_quant_targets(
model,
quantize_lm_head=self.quantization_config.lm_head,
@@ -211,18 +209,6 @@ def _process_model_before_weight_loading(
)
install_quant_tensor_param(target.module, target.param_name, qt)
- # Record literal names of every nn.Linear/Embedding that was
- # *not* converted so ``modules_to_not_convert`` reflects reality
- # (mirrors HF quantizer conventions).
- for name, module in model.named_modules():
- if not isinstance(module, (nn.Linear, nn.Embedding)):
- continue
- if not isinstance(module._parameters.get("weight"), nn.Parameter) or not isinstance(
- module._parameters["weight"].data, QuantTensor
- ):
- skip_literal_names.add(name)
- self.modules_to_not_convert = sorted(skip_literal_names)
-
if self.quantization_config.tie_word_embeddings:
# doing first time so that the weight load doesn't complain about missing weights
tie_quant_word_embeddings(model)
From f8bc45315b454afd2edd133d34f2ab4c7ad4f420 Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 01:10:39 +0000
Subject: [PATCH 10/18] Make quant_info a parameter-level attribute; unify rank
handling
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Stop tagging modules with quant_info / quant_info_3d and stop branching
on 2D vs 3D in the quantization passes. The quantizer already operates
along the last dim regardless of rank, so a single iteration over
parameters that carry a quant_info attribute is enough.
- QuantTarget slims to (module, module_name, pname, full_name) with a
.param property; the caller reads shape/dtype/device from the
parameter directly. No more 'kind' field.
- prepare_model writes target.param.quant_info in one pass — both 2D
linear/embedding weights and fused experts parameters use the same
code path. The quant_info_3d dict-stash on experts modules is gone.
- finalize iterates every parameter that has quant_info, calls
QuantTensor.from_float (already rank-generic), and installs in place.
- GPTQ and AutoClip read module.weight.quant_info; module discovery
uses hasattr(module.weight, 'quant_info') instead of a module-level
attribute.
- HF placeholder install pulls shape/dtype/device off target.param and
the placeholder builder is now rank-generic.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/quant/hf_utils.py | 22 ++---
olive/common/quant/selection.py | 140 ++++++++++++----------------
olive/passes/pytorch/autoclip.py | 14 +--
olive/passes/pytorch/gptq.py | 40 ++++----
olive/passes/pytorch/quant_utils.py | 116 +++++++++++------------
test/common/quant/test_selection.py | 25 ++---
6 files changed, 155 insertions(+), 202 deletions(-)
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index fdc9d01959..bb790851d2 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -199,15 +199,16 @@ def _process_model_before_weight_loading(
skip_patterns=skip_patterns,
):
qargs = self.quantization_config.get_qlinear_init_args(target.full_name)
+ param = target.param
qt = _build_placeholder_quant_tensor(
- shape=target.shape,
+ shape=tuple(param.shape),
bits=qargs["bits"],
symmetric=qargs["symmetric"],
group_size=qargs["group_size"],
- dtype=target.dtype,
- device=target.device,
+ dtype=param.dtype,
+ device=param.device,
)
- install_quant_tensor_param(target.module, target.param_name, qt)
+ install_quant_tensor_param(target.module, target.pname, qt)
if self.quantization_config.tie_word_embeddings:
# doing first time so that the weight load doesn't complain about missing weights
@@ -248,13 +249,7 @@ def _build_placeholder_quant_tensor(
quantizer = WeightQuantizer(bits=bits, symmetric=symmetric, group_size=group_size, signed=False)
packing_factor = 8 // bits
qparam_shape = quantizer.get_qparam_shape(shape)
-
- if len(shape) == 2:
- qweight_shape = (shape[0], math.ceil(shape[1] / packing_factor))
- elif len(shape) == 3:
- qweight_shape = (shape[0], shape[1], math.ceil(shape[2] / packing_factor))
- else:
- raise ValueError(f"QuantTensor placeholder only supports 2D/3D shapes, got {shape}")
+ qweight_shape = (*shape[:-1], math.ceil(shape[-1] / packing_factor))
qweight = torch.zeros(qweight_shape, dtype=torch.uint8, device=device)
scales = torch.zeros(qparam_shape, dtype=dtype, device=device)
@@ -262,10 +257,7 @@ def _build_placeholder_quant_tensor(
if symmetric:
qzeros = None
else:
- if len(qparam_shape) == 2:
- qz_shape = (qparam_shape[0], math.ceil(qparam_shape[1] / packing_factor))
- else:
- qz_shape = (qparam_shape[0], qparam_shape[1], math.ceil(qparam_shape[2] / packing_factor))
+ qz_shape = (*qparam_shape[:-1], math.ceil(qparam_shape[-1] / packing_factor))
qzeros = torch.zeros(qz_shape, dtype=torch.uint8, device=device)
return QuantTensor.from_packed(
diff --git a/olive/common/quant/selection.py b/olive/common/quant/selection.py
index f376d33054..69aed8681a 100644
--- a/olive/common/quant/selection.py
+++ b/olive/common/quant/selection.py
@@ -9,6 +9,12 @@
:class:`QuantTensor` placeholders before weight loading) and the
PyTorch RTN/GPTQ passes (which attach calibration metadata) consume
the same set of targets — only the per-target action differs.
+
+Every target is a single ``nn.Parameter``. The selector makes no
+distinction between 2D linear/embedding weights and 3D fused-MoE
+parameters: downstream code reads ``target.param.shape`` and lets
+:class:`~olive.common.quant.utils.WeightQuantizer` handle any rank
+along the last dim.
"""
from __future__ import annotations
@@ -16,7 +22,6 @@
from dataclasses import dataclass
from typing import TYPE_CHECKING
-import torch
import torch.nn as nn
from olive.common.quant.patterns import match_skip
@@ -25,6 +30,8 @@
from collections.abc import Iterable, Iterator
from typing import Callable
+ import torch
+
from olive.common.hf.wrapper import ModelWrapper
@@ -35,30 +42,23 @@ class QuantTarget:
Attributes:
module: The owning ``nn.Module``.
module_name: Dotted name of ``module`` relative to the model root.
- param_name: Name of the parameter on ``module`` (e.g. ``"weight"``
- for ``nn.Linear``/``nn.Embedding``; an arbitrary 3D
- parameter name for a fused-experts module).
- full_name: ``f"{module_name}.{param_name}"`` (with an exception
- for ``param_name == "weight"`` on ``nn.Linear``/``nn.Embedding``,
- in which case ``full_name == module_name`` — matching how
- overrides and skip patterns have always been keyed).
- kind: ``"linear"`` for ``nn.Linear``, ``"embedding"`` for
- ``nn.Embedding``, ``"fused_experts"`` for a 3D parameter on
- an experts module.
- shape: Logical (dequantized) shape of the parameter.
- dtype: Original parameter dtype.
- device: Original parameter device.
+ pname: Name of the parameter on ``module``.
+ full_name: Key used for overrides / skip-pattern lookups.
+ For ``"weight"`` on ``nn.Linear``/``nn.Embedding`` this is
+ just ``module_name`` (matching the long-standing override
+ convention); otherwise it is ``f"{module_name}.{pname}"``.
"""
module: nn.Module
module_name: str
- param_name: str
+ pname: str
full_name: str
- kind: str
- shape: tuple[int, ...]
- dtype: torch.dtype
- device: torch.device
+
+ @property
+ def param(self) -> torch.nn.Parameter:
+ """The selected parameter."""
+ return self.module._parameters[self.pname]
def _collect_experts(
@@ -85,31 +85,30 @@ def iter_quant_targets(
skip_patterns: Iterable[str] = (),
extra_skip_modules: Iterable[nn.Module] = (),
skip_already_quantized: bool = True,
- consider_linears: bool = True,
- consider_embeddings: bool = True,
) -> Iterator[QuantTarget]:
"""Walk ``model`` once and yield every parameter selected for quantization.
+ Yielded parameters are:
+
+ * ``nn.Linear.weight`` and ``nn.Embedding.weight`` (2D), and
+ * direct ``nn.Parameter`` attributes on each experts module
+ (typically 3D fused-MoE weights), when ``quantize_moe=True``.
+
Selection rules (first matching skip wins):
- * ``extra_skip_modules`` (caller-supplied set, e.g. attention inputs
- excluded by GPTQ) skips the module by identity.
+ * ``extra_skip_modules`` (caller-supplied set, e.g. attention
+ inputs excluded by GPTQ) skips the module by identity.
* ``quantize_lm_head=False`` skips the output embedding module.
* ``quantize_embeds=False`` skips the input embedding module.
- * ``quantize_moe=False`` skips every ``nn.Module`` under any experts
- subtree (this both leaves fused-3D experts alone and prevents
- silently quantizing per-expert ``nn.Linear``s inside
- ``ModuleList(Expert)`` blocks).
+ * ``quantize_moe=False`` skips every ``nn.Module`` under any
+ experts subtree — this both leaves fused parameters alone *and*
+ prevents silently quantizing per-expert ``nn.Linear``s inside
+ ``ModuleList(Expert)`` blocks.
* ``skip_patterns`` matches the parameter's ``full_name`` via the
shared HF-style substring / ``re:``-prefixed regex matcher.
- * When ``skip_already_quantized=True``, weights that are already a
- :class:`QuantTensor` are skipped (idempotent re-runs).
-
- When ``quantize_moe=True`` and an experts module exposes a 3D
- ``nn.Parameter`` (fused experts), that parameter is yielded as a
- ``"fused_experts"`` target. Per-expert 2D ``nn.Linear``s inside an
- ``nn.ModuleList`` experts wrapper continue to come through the
- regular linear walk.
+ * When ``skip_already_quantized=True`` (default), parameters whose
+ underlying tensor is already a :class:`QuantTensor` are skipped
+ (idempotent re-runs).
"""
from olive.common.hf.wrapper import ModelWrapper
from olive.common.quant.tensor import QuantTensor
@@ -117,10 +116,9 @@ def iter_quant_targets(
try:
wrapper = ModelWrapper.from_model(model)
except Exception: # pylint: disable=broad-except
- # Not every model is wrappable (e.g., random sklearn-like
- # test fixtures). Without the wrapper we cannot honour MoE /
- # lm_head / embeds category flags; fall back to the unfiltered
- # 2D walk.
+ # Not every model is wrappable (e.g., random test fixtures).
+ # Without the wrapper we cannot honour MoE / lm_head / embeds
+ # category flags; fall back to the unfiltered 2D walk.
wrapper = None
lm_head_module: nn.Module | None = None
@@ -131,6 +129,7 @@ def iter_quant_targets(
embed_module = model.get_input_embeddings()
expert_modules = _collect_experts(model, wrapper)
+ expert_module_ids = {id(m) for m, _ in expert_modules}
# ID-based skip set for fast identity checks during the named_modules walk.
skip_ids: set[int] = {id(m) for m in extra_skip_modules}
@@ -150,54 +149,31 @@ def _is_skipped(module: nn.Module, full_name: str) -> bool:
return True
return bool(patterns) and match_skip(full_name, patterns)
- # 2D pass: every nn.Linear / nn.Embedding under the model.
+ def _is_already_quantized(param) -> bool:
+ return skip_already_quantized and (isinstance(param, QuantTensor) or isinstance(param.data, QuantTensor))
+
for name, module in model.named_modules():
- if isinstance(module, nn.Linear) and consider_linears:
- kind = "linear"
- elif isinstance(module, nn.Embedding) and consider_embeddings:
- kind = "embedding"
- else:
- continue
- if _is_skipped(module, name):
- continue
- weight = module.weight
- if skip_already_quantized and isinstance(weight, QuantTensor):
+ # 2D pass: nn.Linear / nn.Embedding ``weight`` (legacy
+ # override-key convention: full_name == module_name).
+ if isinstance(module, (nn.Linear, nn.Embedding)):
+ if _is_skipped(module, name):
+ continue
+ weight = module.weight
+ if weight is None or _is_already_quantized(weight):
+ continue
+ yield QuantTarget(module=module, module_name=name, pname="weight", full_name=name)
continue
- if weight is None:
+
+ # Fused-MoE pass: direct parameters on experts modules.
+ if not quantize_moe or id(module) not in expert_module_ids:
continue
- yield QuantTarget(
- module=module,
- module_name=name,
- param_name="weight",
- full_name=name,
- kind=kind,
- shape=tuple(weight.shape),
- dtype=weight.dtype,
- device=weight.device,
- )
-
- # 3D pass: fused-experts modules only when MoE quantization is requested.
- if not quantize_moe:
- return
- for experts_module, experts_name in expert_modules:
- for pname, param in experts_module.named_parameters(recurse=False):
- if param is None or param.dim() != 3:
- continue
- if skip_already_quantized and isinstance(param.data, QuantTensor):
+ for pname, param in module.named_parameters(recurse=False):
+ if param is None or _is_already_quantized(param):
continue
- full_name = f"{experts_name}.{pname}" if experts_name else pname
- if _is_skipped(experts_module, full_name):
+ full_name = f"{name}.{pname}" if name else pname
+ if _is_skipped(module, full_name):
continue
- yield QuantTarget(
- module=experts_module,
- module_name=experts_name,
- param_name=pname,
- full_name=full_name,
- kind="fused_experts",
- shape=tuple(param.shape),
- dtype=param.dtype,
- device=param.device,
- )
+ yield QuantTarget(module=module, module_name=name, pname=pname, full_name=full_name)
def for_each_target(
diff --git a/olive/passes/pytorch/autoclip.py b/olive/passes/pytorch/autoclip.py
index 5b8d7b2960..52e5af3038 100644
--- a/olive/passes/pytorch/autoclip.py
+++ b/olive/passes/pytorch/autoclip.py
@@ -131,9 +131,9 @@ def _get_oc_batch_size(out_features: int) -> int:
@staticmethod
def accumulate_inputs(module: torch.nn.Module, inputs: tuple, _: torch.Tensor) -> None:
- if module.quant_info.data is None:
- module.quant_info.data = {"inputs": []}
- module.quant_info.data["inputs"].append(inputs[0].detach().cpu())
+ if module.weight.quant_info.data is None:
+ module.weight.quant_info.data = {"inputs": []}
+ module.weight.quant_info.data["inputs"].append(inputs[0].detach().cpu())
@classmethod
def process_module(
@@ -144,11 +144,11 @@ def process_module(
max_shrink: float,
n_sample_token: int,
) -> None:
- if module.quant_info.data is None or not module.quant_info.data.get("inputs"):
+ if module.weight.quant_info.data is None or not module.weight.quant_info.data.get("inputs"):
raise ValueError(f"Module {module} does not have cached inputs initialized!")
- input_feat = torch.cat(module.quant_info.data["inputs"], dim=0)
- module.quant_info.data = None
+ input_feat = torch.cat(module.weight.quant_info.data["inputs"], dim=0)
+ module.weight.quant_info.data = None
module.to(device)
cls._auto_clip_layer(
@@ -173,7 +173,7 @@ def _auto_clip_layer(
if weight.dim() != 2:
raise ValueError("AutoClip expects a 2D linear weight tensor.")
- quantizer = module.quant_info.quantizer
+ quantizer = module.weight.quant_info.quantizer
effective_group_size = weight.shape[1] if quantizer.group_size <= 0 else quantizer.group_size
if weight.shape[1] % effective_group_size != 0:
raise ValueError("Weight in_features must be divisible by group_size.")
diff --git a/olive/passes/pytorch/gptq.py b/olive/passes/pytorch/gptq.py
index 03dff6e28e..389ff7686e 100644
--- a/olive/passes/pytorch/gptq.py
+++ b/olive/passes/pytorch/gptq.py
@@ -121,8 +121,8 @@ def accumulate_hessian(module: torch.nn.Module, inp: tuple, _: Any) -> None:
_: Unused output parameter.
"""
- if module.quant_info.data is None:
- module.quant_info.data = {
+ if module.weight.quant_info.data is None:
+ module.weight.quant_info.data = {
"H": torch.zeros((module.in_features, module.in_features), device=inp[0].device),
"N": 0,
}
@@ -130,10 +130,10 @@ def accumulate_hessian(module: torch.nn.Module, inp: tuple, _: Any) -> None:
batch_size = inp[0].shape[0]
inp = inp[0].reshape(-1, module.in_features).t()
- module.quant_info.data["H"] *= module.quant_info.data["N"] / (module.quant_info.data["N"] + batch_size)
- module.quant_info.data["N"] += batch_size
- inp = math.sqrt(2 / module.quant_info.data["N"]) * inp.float()
- module.quant_info.data["H"] += inp.matmul(inp.t())
+ module.weight.quant_info.data["H"] *= module.weight.quant_info.data["N"] / (module.weight.quant_info.data["N"] + batch_size)
+ module.weight.quant_info.data["N"] += batch_size
+ inp = math.sqrt(2 / module.weight.quant_info.data["N"]) * inp.float()
+ module.weight.quant_info.data["H"] += inp.matmul(inp.t())
@staticmethod
def process_module(
@@ -148,18 +148,18 @@ def process_module(
actorder: Whether to use act-order quantization scheme.
"""
- if module.quant_info.data is None:
+ if module.weight.quant_info.data is None:
raise ValueError(f"Module {module} does not have quant_info.data initialized!")
if actorder is None:
- actorder = module.quant_info.quantizer.group_size == -1
+ actorder = module.weight.quant_info.quantizer.group_size == -1
elif actorder is True:
- assert module.quant_info.quantizer.group_size == -1, (
+ assert module.weight.quant_info.quantizer.group_size == -1, (
"actorder can only be True when group_size is -1, but got group_size="
- f"{module.quant_info.quantizer.group_size}"
+ f"{module.weight.quant_info.quantizer.group_size}"
)
- H = module.quant_info.data["H"]
+ H = module.weight.quant_info.data["H"]
W = module.weight.data.clone().float().to(H.device)
num_cols = H.shape[0]
@@ -189,9 +189,9 @@ def process_module(
now_idx = 1
# create a per-channel quantizer
quantizer = WeightQuantizer(
- bits=module.quant_info.quantizer.bits, symmetric=module.quant_info.quantizer.symmetric, group_size=-1
+ bits=module.weight.quant_info.quantizer.bits, symmetric=module.weight.quant_info.quantizer.symmetric, group_size=-1
)
- if module.quant_info.quantizer.group_size == -1:
+ if module.weight.quant_info.quantizer.group_size == -1:
# this can be before or after actorder permutation since there's only one group
active_scale, active_zp = quantizer.find_qparams(W)
else:
@@ -211,13 +211,13 @@ def process_module(
w = W1[:, i]
d = Hinv1[i, i]
- if module.quant_info.quantizer.group_size != -1:
- if (i1 + i) % module.quant_info.quantizer.group_size == 0:
+ if module.weight.quant_info.quantizer.group_size != -1:
+ if (i1 + i) % module.weight.quant_info.quantizer.group_size == 0:
active_scale, active_zp = quantizer.find_qparams(
- W[:, (i1 + i) : (i1 + i + module.quant_info.quantizer.group_size)]
+ W[:, (i1 + i) : (i1 + i + module.weight.quant_info.quantizer.group_size)]
)
- if ((i1 + i) // module.quant_info.quantizer.group_size) - now_idx == -1:
+ if ((i1 + i) // module.weight.quant_info.quantizer.group_size) - now_idx == -1:
all_scales.append(active_scale)
all_zp.append(active_zp)
now_idx += 1
@@ -243,9 +243,9 @@ def process_module(
all_zp.append(active_zp)
module.weight.data = Q.to(module.weight.data.device).to(module.weight.data.dtype)
- module.quant_info.scales = torch.cat(all_scales, dim=1).to("cpu")
- module.quant_info.zero_points = torch.cat(all_zp, dim=1).to("cpu")
+ module.weight.quant_info.scales = torch.cat(all_scales, dim=1).to("cpu")
+ module.weight.quant_info.zero_points = torch.cat(all_zp, dim=1).to("cpu")
- module.quant_info.data = None
+ module.weight.quant_info.data = None
if torch.cuda.is_available():
torch.cuda.empty_cache()
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index 8869fcb987..5906fc905c 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -153,7 +153,6 @@ def prepare_model(
skip_patterns = list(getattr(qcfg, "modules_to_not_convert", None) or [])
- fused_targets_by_module: dict[int, dict[str, QuantInfo]] = {}
for target in iter_quant_targets(
wrapper.model,
quantize_lm_head=qcfg.lm_head,
@@ -164,22 +163,7 @@ def prepare_model(
):
qargs = qcfg.get_qlinear_init_args(target.full_name)
new_qargs[target.full_name] = qargs
- if target.kind == "fused_experts":
- fused_targets_by_module.setdefault(id(target.module), {})[target.param_name] = QuantInfo(
- quantizer=WeightQuantizer(**qargs)
- )
- # Stash the experts module on the dict so we can install
- # ``quant_info_3d`` below without re-walking.
- fused_targets_by_module[id(target.module)]["__module__"] = target.module # type: ignore[assignment]
- else:
- target.module.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs))
-
- # Fused-3D MoE: stamp the per-parameter QuantInfo dict onto each
- # experts module so ``finalize`` can replace each parameter with a
- # QuantTensor.
- for entry in fused_targets_by_module.values():
- module = entry.pop("__module__") # type: ignore[arg-type]
- module.quant_info_3d = entry # type: ignore[assignment]
+ target.param.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs))
# remove overrides for modules not being quantized
for name in list(qcfg.overrides or {}):
@@ -403,7 +387,7 @@ def run_layerwise_quantization(
for layer_idx, layer in enumerate(wrapper.get_layers(return_name=False)):
pbar.set_postfix(module=f"layers.{layer_idx}", refresh=False)
- quantizable_modules = [module for module in layer.modules() if hasattr(module, "quant_info")]
+ quantizable_modules = [module for module in layer.modules() if _module_weight_has_quant_info(module)]
handles = [module.register_forward_hook(input_hook) for module in quantizable_modules]
if update_before_process:
@@ -455,6 +439,30 @@ def run_layerwise_quantization(
return device
+def _module_weight_has_quant_info(module: torch.nn.Module) -> bool:
+ """Return True if ``module.weight`` carries a ``quant_info`` attribute.
+
+ Used by the layerwise discovery in :func:`run_layerwise_quantization`
+ to locate ``nn.Linear`` modules selected for calibrated quantization
+ without depending on a module-level attribute.
+ """
+ weight = getattr(module, "weight", None)
+ return weight is not None and hasattr(weight, "quant_info")
+
+
+def _iter_quant_info_params(model: torch.nn.Module):
+ """Yield ``(module, pname, param, quant_info)`` for every selected parameter."""
+ for sub_module in model.modules():
+ for pname in list(sub_module._parameters):
+ param = sub_module._parameters.get(pname)
+ if param is None:
+ continue
+ info = getattr(param, "quant_info", None)
+ if info is None:
+ continue
+ yield sub_module, pname, param, info
+
+
def finalize(
model: HfModelHandler,
output_model_path: str,
@@ -463,56 +471,38 @@ def finalize(
device: str,
retie_word_embeddings: bool = False,
) -> HfModelHandler:
- """Finalize quantization by installing ``QuantTensor`` parameters on the host modules.
+ """Finalize quantization by installing ``QuantTensor`` parameters in place.
- For every module marked with ``quant_info`` (2D linear / embedding) and
- every fused-3D MoE experts module marked with ``quant_info_3d``, build a
- ``QuantTensor`` from the float weight + computed qparams and install it
- via :func:`install_quant_tensor_param` so that:
+ Walks every ``nn.Parameter`` whose tensor has a ``quant_info``
+ attribute (set by :func:`prepare_model`), builds a ``QuantTensor``
+ from the float tensor plus computed qparams, and installs it via
+ :func:`install_quant_tensor_param` so that:
- * ``module.`` is an ``nn.Parameter(QuantTensor)`` whose dispatch
- still drives the original eager forward;
+ * ``module.`` is an ``nn.Parameter(QuantTensor)`` whose
+ dispatch still drives the original eager forward;
* ``module._qweight`` / ``_scales`` / ``_qzeros`` are plain
buffers (aliasing the QuantTensor's inner tensors), so the model
- saves cleanly via ``save_pretrained`` / safetensors with **no
- tensor subclass on disk**.
+ saves cleanly via ``save_pretrained`` / safetensors with no
+ tensor subclass on disk.
+
+ The same code path handles 2D linear/embedding weights and any
+ higher-rank fused parameter (e.g. 3D MoE experts) — quantization
+ is always along the last dim.
"""
- for sub_module in wrapper.model.modules():
- # ---- 2D path: nn.Linear / nn.Embedding ----------------------
- qinfo = getattr(sub_module, "quant_info", None)
- if qinfo is not None and isinstance(sub_module, (torch.nn.Linear, torch.nn.Embedding)):
- sub_module.to(device)
- quantizer = qinfo.quantizer
- with torch.no_grad():
- qt = QuantTensor.from_float(
- sub_module.weight.detach(),
- bits=quantizer.bits,
- symmetric=quantizer.symmetric,
- group_size=quantizer.group_size,
- scales=qinfo.scales,
- zero_points=qinfo.zero_points,
- ).to("cpu")
- sub_module.to("cpu")
- install_quant_tensor_param(sub_module, "weight", qt)
- del sub_module.quant_info
-
- # ---- fused-3D MoE path -------------------------------------
- info_3d: dict | None = getattr(sub_module, "quant_info_3d", None)
- if info_3d:
- for pname, qinfo_3d in info_3d.items():
- param = sub_module._parameters.get(pname)
- if param is None:
- continue
- quantizer = qinfo_3d.quantizer
- with torch.no_grad():
- qt = QuantTensor.from_float(
- param.detach().to(device),
- bits=quantizer.bits,
- symmetric=quantizer.symmetric,
- group_size=quantizer.group_size,
- ).to("cpu")
- install_quant_tensor_param(sub_module, pname, qt)
- del sub_module.quant_info_3d
+ for sub_module, pname, param, info in list(_iter_quant_info_params(wrapper.model)):
+ quantizer = info.quantizer
+ sub_module.to(device)
+ with torch.no_grad():
+ qt = QuantTensor.from_float(
+ param.data.detach(),
+ bits=quantizer.bits,
+ symmetric=quantizer.symmetric,
+ group_size=quantizer.group_size,
+ scales=info.scales,
+ zero_points=info.zero_points,
+ ).to("cpu")
+ sub_module.to("cpu")
+ install_quant_tensor_param(sub_module, pname, qt)
if retie_word_embeddings:
tie_quant_word_embeddings(wrapper.model)
diff --git a/test/common/quant/test_selection.py b/test/common/quant/test_selection.py
index d36dd197e3..6bf6e014e6 100644
--- a/test/common/quant/test_selection.py
+++ b/test/common/quant/test_selection.py
@@ -27,23 +27,19 @@ def get_output_embeddings(self):
def _names(targets):
- return sorted((t.full_name, t.kind) for t in targets)
+ return sorted(t.full_name for t in targets)
def test_default_skips_lm_head_and_embeds():
m = _Toy()
targets = list(iter_quant_targets(m, quantize_lm_head=False, quantize_embeds=False, quantize_moe=False))
- assert _names(targets) == [("linear", "linear")]
+ assert _names(targets) == ["linear"]
def test_include_lm_head_and_embeds():
m = _Toy()
targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False))
- assert _names(targets) == [
- ("embed_tokens", "embedding"),
- ("linear", "linear"),
- ("lm_head", "linear"),
- ]
+ assert _names(targets) == ["embed_tokens", "linear", "lm_head"]
def test_skip_patterns_filter_by_name():
@@ -57,7 +53,7 @@ def test_skip_patterns_filter_by_name():
skip_patterns=["re:.*_head"],
)
)
- assert _names(targets) == [("embed_tokens", "embedding"), ("linear", "linear")]
+ assert _names(targets) == ["embed_tokens", "linear"]
def test_extra_skip_modules_skip_by_identity():
@@ -71,14 +67,13 @@ def test_extra_skip_modules_skip_by_identity():
extra_skip_modules={m.linear},
)
)
- assert _names(targets) == [("lm_head", "linear")]
+ assert _names(targets) == ["lm_head"]
def test_already_quantized_param_is_skipped():
from olive.common.quant.tensor import QuantTensor
m = _Toy()
- # Build a tiny QuantTensor placeholder and stamp it onto m.linear.weight.
qt = QuantTensor.from_packed(
qweight=torch.zeros((8, 4), dtype=torch.uint8),
scales=torch.zeros((8, 1), dtype=torch.float32),
@@ -92,7 +87,7 @@ def test_already_quantized_param_is_skipped():
m.linear.weight = nn.Parameter(qt, requires_grad=False)
targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=False, quantize_moe=False))
- assert _names(targets) == [("lm_head", "linear")]
+ assert _names(targets) == ["lm_head"]
class _ExpertList(nn.Module):
@@ -160,8 +155,8 @@ def get_layer_wrappers(self):
m = _Model()
targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True))
- fused = [(t.full_name, t.kind, t.shape) for t in targets if t.kind == "fused_experts"]
- assert sorted(fused) == [
- ("experts.down_proj", "fused_experts", (4, 16, 8)),
- ("experts.gate_up_proj", "fused_experts", (4, 8, 16)),
+ fused = sorted((t.full_name, tuple(t.param.shape)) for t in targets)
+ assert fused == [
+ ("experts.down_proj", (4, 16, 8)),
+ ("experts.gate_up_proj", (4, 8, 16)),
]
From a488f233e6b327bad55a6fd448ef10c84b1ea1da Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 01:22:02 +0000
Subject: [PATCH 11/18] Drop QuantTarget dataclass; yield (module, pname,
full_name) tuples
The dataclass had four fields and a one-line .param property, used by
two callers. A plain tuple is shorter, matches how the layerwise
quantization loop already iterates over (module, pname, param, info)
tuples, and removes the unused module_name field and dead
for_each_target helper.
QuantTarget remains as a type alias for the public signature.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/quant/hf_utils.py | 8 ++--
olive/common/quant/selection.py | 60 +++++++----------------------
olive/passes/pytorch/quant_utils.py | 8 ++--
test/common/quant/test_selection.py | 4 +-
4 files changed, 23 insertions(+), 57 deletions(-)
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index bb790851d2..aa544bd14e 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -191,15 +191,15 @@ def _process_model_before_weight_loading(
if keep_in_fp32_modules:
skip_patterns.extend(keep_in_fp32_modules)
- for target in iter_quant_targets(
+ for module, pname, full_name in iter_quant_targets(
model,
quantize_lm_head=self.quantization_config.lm_head,
quantize_embeds=self.quantization_config.embeds,
quantize_moe=self.quantization_config.moe,
skip_patterns=skip_patterns,
):
- qargs = self.quantization_config.get_qlinear_init_args(target.full_name)
- param = target.param
+ qargs = self.quantization_config.get_qlinear_init_args(full_name)
+ param = module._parameters[pname]
qt = _build_placeholder_quant_tensor(
shape=tuple(param.shape),
bits=qargs["bits"],
@@ -208,7 +208,7 @@ def _process_model_before_weight_loading(
dtype=param.dtype,
device=param.device,
)
- install_quant_tensor_param(target.module, target.pname, qt)
+ install_quant_tensor_param(module, pname, qt)
if self.quantization_config.tie_word_embeddings:
# doing first time so that the weight load doesn't complain about missing weights
diff --git a/olive/common/quant/selection.py b/olive/common/quant/selection.py
index 69aed8681a..87afc3e540 100644
--- a/olive/common/quant/selection.py
+++ b/olive/common/quant/selection.py
@@ -10,16 +10,19 @@
PyTorch RTN/GPTQ passes (which attach calibration metadata) consume
the same set of targets — only the per-target action differs.
-Every target is a single ``nn.Parameter``. The selector makes no
-distinction between 2D linear/embedding weights and 3D fused-MoE
-parameters: downstream code reads ``target.param.shape`` and lets
+Every target is a single ``nn.Parameter``, yielded as a
+``(module, pname, full_name)`` tuple. ``full_name`` is the key used
+for overrides / skip-pattern lookups (``module_name`` for ``"weight"``
+on ``nn.Linear`` / ``nn.Embedding``; ``f"{module_name}.{pname}"``
+otherwise). The selector makes no distinction between 2D linear /
+embedding weights and 3D fused-MoE parameters — downstream code reads
+the parameter's own shape and lets
:class:`~olive.common.quant.utils.WeightQuantizer` handle any rank
along the last dim.
"""
from __future__ import annotations
-from dataclasses import dataclass
from typing import TYPE_CHECKING
import torch.nn as nn
@@ -28,37 +31,12 @@
if TYPE_CHECKING:
from collections.abc import Iterable, Iterator
- from typing import Callable
-
- import torch
from olive.common.hf.wrapper import ModelWrapper
-@dataclass
-class QuantTarget:
- """A single parameter selected for quantization.
-
- Attributes:
- module: The owning ``nn.Module``.
- module_name: Dotted name of ``module`` relative to the model root.
- pname: Name of the parameter on ``module``.
- full_name: Key used for overrides / skip-pattern lookups.
- For ``"weight"`` on ``nn.Linear``/``nn.Embedding`` this is
- just ``module_name`` (matching the long-standing override
- convention); otherwise it is ``f"{module_name}.{pname}"``.
-
- """
-
- module: nn.Module
- module_name: str
- pname: str
- full_name: str
-
- @property
- def param(self) -> torch.nn.Parameter:
- """The selected parameter."""
- return self.module._parameters[self.pname]
+QuantTarget = tuple[nn.Module, str, str]
+"""``(module, pname, full_name)`` for a single parameter selected for quantization."""
def _collect_experts(
@@ -153,15 +131,15 @@ def _is_already_quantized(param) -> bool:
return skip_already_quantized and (isinstance(param, QuantTensor) or isinstance(param.data, QuantTensor))
for name, module in model.named_modules():
- # 2D pass: nn.Linear / nn.Embedding ``weight`` (legacy
- # override-key convention: full_name == module_name).
+ # nn.Linear / nn.Embedding ``weight`` — legacy override-key
+ # convention: full_name == module_name.
if isinstance(module, (nn.Linear, nn.Embedding)):
if _is_skipped(module, name):
continue
weight = module.weight
if weight is None or _is_already_quantized(weight):
continue
- yield QuantTarget(module=module, module_name=name, pname="weight", full_name=name)
+ yield module, "weight", name
continue
# Fused-MoE pass: direct parameters on experts modules.
@@ -173,16 +151,4 @@ def _is_already_quantized(param) -> bool:
full_name = f"{name}.{pname}" if name else pname
if _is_skipped(module, full_name):
continue
- yield QuantTarget(module=module, module_name=name, pname=pname, full_name=full_name)
-
-
-def for_each_target(
- model: nn.Module,
- handler: Callable[[QuantTarget], None],
- **selection_kwargs,
-) -> list[QuantTarget]:
- """Run ``handler`` once per target and return the materialised list."""
- targets = list(iter_quant_targets(model, **selection_kwargs))
- for t in targets:
- handler(t)
- return targets
+ yield module, pname, full_name
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index 5906fc905c..401cafe176 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -153,7 +153,7 @@ def prepare_model(
skip_patterns = list(getattr(qcfg, "modules_to_not_convert", None) or [])
- for target in iter_quant_targets(
+ for module, pname, full_name in iter_quant_targets(
wrapper.model,
quantize_lm_head=qcfg.lm_head,
quantize_embeds=qcfg.embeds,
@@ -161,9 +161,9 @@ def prepare_model(
skip_patterns=skip_patterns,
extra_skip_modules=excluded_attn_inputs,
):
- qargs = qcfg.get_qlinear_init_args(target.full_name)
- new_qargs[target.full_name] = qargs
- target.param.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs))
+ qargs = qcfg.get_qlinear_init_args(full_name)
+ new_qargs[full_name] = qargs
+ module._parameters[pname].quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs))
# remove overrides for modules not being quantized
for name in list(qcfg.overrides or {}):
diff --git a/test/common/quant/test_selection.py b/test/common/quant/test_selection.py
index 6bf6e014e6..325bd4f9c8 100644
--- a/test/common/quant/test_selection.py
+++ b/test/common/quant/test_selection.py
@@ -27,7 +27,7 @@ def get_output_embeddings(self):
def _names(targets):
- return sorted(t.full_name for t in targets)
+ return sorted(full_name for _, _, full_name in targets)
def test_default_skips_lm_head_and_embeds():
@@ -155,7 +155,7 @@ def get_layer_wrappers(self):
m = _Model()
targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True))
- fused = sorted((t.full_name, tuple(t.param.shape)) for t in targets)
+ fused = sorted((full_name, tuple(module._parameters[pname].shape)) for module, pname, full_name in targets)
assert fused == [
("experts.down_proj", (4, 16, 8)),
("experts.gate_up_proj", (4, 8, 16)),
From 9ecdd501726202097641b76e509b21fee0096d53 Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 01:34:08 +0000
Subject: [PATCH 12/18] Address rubber-duck review findings
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* Filter fused-MoE params to 2D/3D ranks in iter_quant_targets so a
1D bias-like parameter fails at selection time instead of much later
in finalize.
* refresh_quant_tensor_refs: also check isinstance(param.data,
QuantTensor) for forward-compat with future torch versions that may
not return the underlying subclass from nn.Parameter().
* OliveHfQuantizationConfig: replace bare '# pylint: disable' with the
specific super-init-not-called rule; use output.get(k) in to_dict.
* finalize: log a warning when moe=True that the resulting checkpoint
isn't directly ONNX-exportable via the Olive conversion pass — it
must be consumed by an MoE-aware model builder.
* Add regression test that _module_weight_has_quant_info ignores
nn.LayerNorm / nn.Conv2d / unmarked nn.Linear (defends GPTQ/AutoClip
discovery against future drift).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/quant/hf_utils.py | 6 ++----
olive/common/quant/selection.py | 2 +-
olive/common/quant/state_dict.py | 2 +-
olive/passes/pytorch/quant_utils.py | 7 +++++++
test/passes/pytorch/test_quant_utils.py | 23 +++++++++++++++++++++++
5 files changed, 34 insertions(+), 6 deletions(-)
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index aa544bd14e..f37de79fdf 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -80,8 +80,7 @@ class OliveHfQuantizationConfig(QuantizationConfigMixin):
"""
- # pylint: disable
- def __init__(
+ def __init__( # pylint: disable=super-init-not-called
self,
bits: int,
symmetric: bool,
@@ -94,7 +93,6 @@ def __init__(
tie_word_embeddings: bool = False,
**kwargs,
):
- # pylint: disable=W0231
self.quant_method = OliveHfQuantizationMethod.OLIVE
self.bits = bits
@@ -123,7 +121,7 @@ def to_dict(self) -> dict:
overrides = {}
for module_name, override in self.overrides.items():
# remove None or default values from the override
- cleaned_override = {k: v for k, v in override.__dict__.items() if v is not None and v != output[k]}
+ cleaned_override = {k: v for k, v in override.__dict__.items() if v is not None and v != output.get(k)}
if cleaned_override:
overrides[module_name] = cleaned_override
output["overrides"] = sort_layers_by_name(overrides) or None
diff --git a/olive/common/quant/selection.py b/olive/common/quant/selection.py
index 87afc3e540..d45e7256a7 100644
--- a/olive/common/quant/selection.py
+++ b/olive/common/quant/selection.py
@@ -146,7 +146,7 @@ def _is_already_quantized(param) -> bool:
if not quantize_moe or id(module) not in expert_module_ids:
continue
for pname, param in module.named_parameters(recurse=False):
- if param is None or _is_already_quantized(param):
+ if param is None or param.dim() not in (2, 3) or _is_already_quantized(param):
continue
full_name = f"{name}.{pname}" if name else pname
if _is_skipped(module, full_name):
diff --git a/olive/common/quant/state_dict.py b/olive/common/quant/state_dict.py
index c3a8aff187..f663b19351 100644
--- a/olive/common/quant/state_dict.py
+++ b/olive/common/quant/state_dict.py
@@ -139,7 +139,7 @@ def refresh_quant_tensor_refs(module: torch.nn.Module) -> None:
# ``param`` itself is the QuantTensor instance stored on the
# module (``nn.Parameter(qt)`` for a tensor subclass returns
# the underlying QuantTensor — see ``torch.nn.Parameter.__new__``).
- if param is None or not isinstance(param, QuantTensor):
+ if param is None or not (isinstance(param, QuantTensor) or isinstance(param.data, QuantTensor)):
continue
qname, sname, zname = buffer_names(pname)
qweight = sub_module._buffers.get(qname)
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index 401cafe176..617ae6b967 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -508,6 +508,13 @@ def finalize(
tie_quant_word_embeddings(wrapper.model)
quant_config.tie_word_embeddings = True
+ if getattr(quant_config, "moe", False):
+ logger.warning(
+ "MoE weights have been quantized as 3D tensor parameters. The resulting checkpoint is "
+ "save/load compatible via transformers but is not directly exportable to ONNX with the "
+ "Olive ONNX conversion pass — consume it via the ORT GenAI model_builder or Mobius."
+ )
+
wrapper.model.quantization_method = quant_config.quant_method
wrapper.model.config.quantization_config = quant_config
diff --git a/test/passes/pytorch/test_quant_utils.py b/test/passes/pytorch/test_quant_utils.py
index 402377dacb..6686e85ba0 100644
--- a/test/passes/pytorch/test_quant_utils.py
+++ b/test/passes/pytorch/test_quant_utils.py
@@ -88,3 +88,26 @@ def test_install_on_linear_module(self):
y = linear(x)
assert y.shape == (2, 32)
assert not isinstance(y, QuantTensor)
+
+
+def test_module_weight_has_quant_info_only_for_marked_params():
+ """Regression: discovery must not pick up LayerNorm / Conv2d weights.
+
+ GPTQ / AutoClip discover quantizable layers via
+ ``_module_weight_has_quant_info``. Modules that happen to expose a
+ ``weight`` attribute but never had ``quant_info`` stamped on it must
+ be left alone.
+ """
+ from olive.common.quant.utils import WeightQuantizer
+ from olive.passes.pytorch.quant_utils import QuantInfo, _module_weight_has_quant_info
+
+ ln = nn.LayerNorm(16)
+ conv = nn.Conv2d(3, 8, kernel_size=3)
+ linear_unmarked = nn.Linear(16, 32, bias=False)
+ linear_marked = nn.Linear(16, 32, bias=False)
+ linear_marked.weight.quant_info = QuantInfo(quantizer=WeightQuantizer(bits=4, symmetric=True, group_size=16))
+
+ assert not _module_weight_has_quant_info(ln)
+ assert not _module_weight_has_quant_info(conv)
+ assert not _module_weight_has_quant_info(linear_unmarked)
+ assert _module_weight_has_quant_info(linear_marked)
From de4c04efc467e36165dea027581762c3ae8bbbef Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Sat, 16 May 2026 01:53:01 +0000
Subject: [PATCH 13/18] Address Copilot PR feedback + add forward parity tests
- __torch_dispatch__ clone/contiguous now forwards extra args/kwargs.
- iter_quant_targets skips ALL nn.Embedding when embeds=False (positional
/ token-type embeddings like GPT-2 wpe are no longer silently quantized).
- WeightQuantizer assertion message: 2/4/8-bit (was 4/8-bit).
- tie_quant_word_embeddings: mark dst aliased buffers non-persistent so
safetensors save emits one copy of qweight/scales/qzeros.
- finalize: group selected params by host module so each module's
to(device)/to(cpu) cycle runs once for MoE experts modules carrying
multiple 3D weight params.
- state_dict: add ensure_state_dict_hooks(model) defensive walk that
installs the save hook on every host module that owns a QuantTensor
parameter (idempotent).
- Add test_forward_parity.py: bit-exact eager parity for full models
(embedding + linears) and fused 3D MoE forwards, plus end-to-end
ONNX export -> onnxruntime numerical parity for Olive-quantized
nn.Linear via make_export_compatible_quant.
- Enable pylint by adding file-level protected-access disables on the
files that intentionally touch nn.Module._parameters.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
olive/common/quant/hf_utils.py | 7 +-
olive/common/quant/selection.py | 11 +-
olive/common/quant/state_dict.py | 19 +++
olive/common/quant/tensor.py | 4 +-
olive/common/quant/utils.py | 2 +-
olive/passes/pytorch/quant_utils.py | 35 +++--
test/common/quant/test_forward_parity.py | 165 +++++++++++++++++++++++
test/common/quant/test_utils.py | 2 +-
8 files changed, 224 insertions(+), 21 deletions(-)
create mode 100644 test/common/quant/test_forward_parity.py
diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py
index f37de79fdf..a42fd71518 100644
--- a/olive/common/quant/hf_utils.py
+++ b/olive/common/quant/hf_utils.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
from __future__ import annotations
import math
@@ -351,11 +352,15 @@ def tie_quant_word_embeddings(model: PreTrainedModel) -> None:
return
qname, sname, zname = buffer_names("weight")
- # tie buffers
+ # Tie buffers, marking the destination copy non-persistent so
+ # ``state_dict()`` / safetensors only emit one set of keys for the
+ # shared tensors (avoids duplicate qweight/scales on disk).
for n in (qname, sname, zname):
src_buf = src._buffers.get(n)
if src_buf is None:
continue
+ if n in dst._buffers:
+ dst._non_persistent_buffers_set.add(n)
dst._buffers[n] = src_buf
# tie the QuantTensor parameter itself (same Python Parameter object,
diff --git a/olive/common/quant/selection.py b/olive/common/quant/selection.py
index d45e7256a7..2330c2fd51 100644
--- a/olive/common/quant/selection.py
+++ b/olive/common/quant/selection.py
@@ -100,11 +100,8 @@ def iter_quant_targets(
wrapper = None
lm_head_module: nn.Module | None = None
- embed_module: nn.Module | None = None
if hasattr(model, "get_output_embeddings"):
lm_head_module = model.get_output_embeddings()
- if hasattr(model, "get_input_embeddings"):
- embed_module = model.get_input_embeddings()
expert_modules = _collect_experts(model, wrapper)
expert_module_ids = {id(m) for m, _ in expert_modules}
@@ -113,8 +110,6 @@ def iter_quant_targets(
skip_ids: set[int] = {id(m) for m in extra_skip_modules}
if not quantize_lm_head and lm_head_module is not None:
skip_ids.add(id(lm_head_module))
- if not quantize_embeds and embed_module is not None:
- skip_ids.add(id(embed_module))
if not quantize_moe:
for experts, _ in expert_modules:
for sub in experts.modules():
@@ -132,8 +127,12 @@ def _is_already_quantized(param) -> bool:
for name, module in model.named_modules():
# nn.Linear / nn.Embedding ``weight`` — legacy override-key
- # convention: full_name == module_name.
+ # convention: full_name == module_name. When ``quantize_embeds``
+ # is False every ``nn.Embedding`` is skipped (positional /
+ # token-type / etc.), not just ``model.get_input_embeddings()``.
if isinstance(module, (nn.Linear, nn.Embedding)):
+ if isinstance(module, nn.Embedding) and not quantize_embeds:
+ continue
if _is_skipped(module, name):
continue
weight = module.weight
diff --git a/olive/common/quant/state_dict.py b/olive/common/quant/state_dict.py
index f663b19351..0f29b437a6 100644
--- a/olive/common/quant/state_dict.py
+++ b/olive/common/quant/state_dict.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
"""State-dict helpers for Olive's quantized weight representation.
Olive's quantization layout:
@@ -76,6 +77,24 @@ def install_state_dict_hooks(module: torch.nn.Module) -> None:
setattr(module, _INSTALLED_FLAG, True)
+def ensure_state_dict_hooks(model: torch.nn.Module) -> None:
+ """Install the save hook on every submodule that hosts a QuantTensor parameter.
+
+ Belt-and-suspenders for paths that may install a ``QuantTensor``
+ parameter without going through :func:`install_quant_tensor_param`
+ (e.g. retie helpers, future loaders). The hook is idempotent.
+ """
+ from olive.common.quant.tensor import QuantTensor
+
+ for sub_module in model.modules():
+ for param in sub_module._parameters.values():
+ if param is None:
+ continue
+ if isinstance(param, QuantTensor) or isinstance(getattr(param, "data", None), QuantTensor):
+ install_state_dict_hooks(sub_module)
+ break
+
+
def install_quant_tensor_param(
module: torch.nn.Module,
pname: str,
diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py
index b81e32d7f3..7ad1d2fced 100644
--- a/olive/common/quant/tensor.py
+++ b/olive/common/quant/tensor.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access,super-init-not-called,redefined-builtin,not-callable
"""QuantTensor — wrapper ``torch.Tensor`` subclass for weight-quantized parameters.
It stores quantization buffers (``qweight``, ``scales``, ``qzeros``) but presents the
@@ -306,7 +307,8 @@ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
if func in (aten.detach.default, aten.clone.default, aten.alias.default, aten.contiguous.default):
self_ = args[0]
- return self_._apply_fn_to_data(func)
+ extra_args = args[1:]
+ return self_._apply_fn_to_data(lambda x: func(x, *extra_args, **kwargs))
if func is aten._to_copy.default:
self_ = args[0]
diff --git a/olive/common/quant/utils.py b/olive/common/quant/utils.py
index 7b478ad906..289d4e4e90 100644
--- a/olive/common/quant/utils.py
+++ b/olive/common/quant/utils.py
@@ -36,7 +36,7 @@ def __init__(self, bits: int = 4, symmetric: bool = True, group_size: int = 0, s
signed: Whether to use signed quantization (default is False, meaning unsigned)
"""
- assert bits in [2, 4, 8], "Only 4-bit and 8-bit quantization supported"
+ assert bits in [2, 4, 8], "Only 2-bit, 4-bit and 8-bit quantization supported"
self.bits = bits
self.symmetric = symmetric
self.group_size = group_size
diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py
index 617ae6b967..23ecab1127 100644
--- a/olive/passes/pytorch/quant_utils.py
+++ b/olive/passes/pytorch/quant_utils.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
from __future__ import annotations
import logging
@@ -489,20 +490,32 @@ def finalize(
higher-rank fused parameter (e.g. 3D MoE experts) — quantization
is always along the last dim.
"""
- for sub_module, pname, param, info in list(_iter_quant_info_params(wrapper.model)):
- quantizer = info.quantizer
+ # Group selected params by their owning module so the module is
+ # moved to ``device`` once even when it owns multiple quantized
+ # parameters (fused-MoE experts modules carry two 3D tensors).
+ by_module: dict[int, tuple[torch.nn.Module, list[tuple[str, torch.nn.Parameter, QuantInfo]]]] = {}
+ for sub_module, pname, param, info in _iter_quant_info_params(wrapper.model):
+ entry = by_module.setdefault(id(sub_module), (sub_module, []))
+ entry[1].append((pname, param, info))
+
+ for sub_module, params in by_module.values():
sub_module.to(device)
with torch.no_grad():
- qt = QuantTensor.from_float(
- param.data.detach(),
- bits=quantizer.bits,
- symmetric=quantizer.symmetric,
- group_size=quantizer.group_size,
- scales=info.scales,
- zero_points=info.zero_points,
- ).to("cpu")
+ built = []
+ for pname, param, info in params:
+ quantizer = info.quantizer
+ qt = QuantTensor.from_float(
+ param.data.detach(),
+ bits=quantizer.bits,
+ symmetric=quantizer.symmetric,
+ group_size=quantizer.group_size,
+ scales=info.scales,
+ zero_points=info.zero_points,
+ ).to("cpu")
+ built.append((pname, qt))
sub_module.to("cpu")
- install_quant_tensor_param(sub_module, pname, qt)
+ for pname, qt in built:
+ install_quant_tensor_param(sub_module, pname, qt)
if retie_word_embeddings:
tie_quant_word_embeddings(wrapper.model)
diff --git a/test/common/quant/test_forward_parity.py b/test/common/quant/test_forward_parity.py
new file mode 100644
index 0000000000..69aadcecbd
--- /dev/null
+++ b/test/common/quant/test_forward_parity.py
@@ -0,0 +1,165 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Numerical-parity tests for QuantTensor-backed modules.
+
+Confirms that ``nn.Linear`` / ``nn.Embedding`` / fused-3D MoE host modules
+whose weight has been swapped to a ``QuantTensor`` produce **bit-exact**
+outputs vs the same forward run against ``QuantTensor.to_dense()`` (the
+canonical eager dequant), and that the ONNX-export wrappers built by
+``make_export_compatible_quant`` round-trip through onnxruntime with
+matching outputs.
+"""
+
+from __future__ import annotations
+
+import copy
+
+import onnx
+import onnxruntime as ort
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from olive.common.hf.quant import make_export_compatible_quant
+from olive.common.quant.state_dict import install_quant_tensor_param
+from olive.common.quant.tensor import QuantTensor
+
+
+def _quantize_inplace(module: nn.Module, pname: str, *, bits: int, group_size: int, symmetric: bool) -> None:
+ param = module._parameters[pname]
+ qt = QuantTensor.from_float(
+ param.data.detach().clone(),
+ bits=bits,
+ group_size=group_size,
+ symmetric=symmetric,
+ )
+ install_quant_tensor_param(module, pname, qt)
+
+
+def _dense_reference(model: nn.Module) -> nn.Module:
+ """Return a deep copy of ``model`` with every QuantTensor weight materialized."""
+ ref = copy.deepcopy(model)
+ for module in ref.modules():
+ weight = module._parameters.get("weight")
+ if weight is None:
+ continue
+ if isinstance(weight.data, QuantTensor):
+ module._parameters["weight"] = nn.Parameter(weight.data.to_dense(), requires_grad=False)
+ return ref
+
+
+@pytest.mark.parametrize("bits,group_size,symmetric", [(4, 32, False), (4, -1, True), (8, 32, False)])
+def test_full_model_forward_parity(bits, group_size, symmetric):
+ torch.manual_seed(0)
+
+ class Toy(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.embed = nn.Embedding(32, 64)
+ self.fc1 = nn.Linear(64, 128, bias=False)
+ self.fc2 = nn.Linear(128, 64, bias=True)
+
+ def forward(self, ids: torch.Tensor) -> torch.Tensor:
+ return self.fc2(F.silu(self.fc1(self.embed(ids))))
+
+ model = Toy().eval()
+ ids = torch.randint(0, 32, (2, 16))
+ ref_out = model(ids)
+
+ for sub_module in (model.embed, model.fc1, model.fc2):
+ _quantize_inplace(sub_module, "weight", bits=bits, group_size=group_size, symmetric=symmetric)
+
+ quant_out = model(ids)
+ dense_out = _dense_reference(model)(ids)
+
+ # vs canonical dequant: must be bit-exact (same kernels, same data)
+ torch.testing.assert_close(quant_out, dense_out, rtol=0, atol=0)
+ # vs original fp: just a sanity bound on quant error
+ assert (quant_out - ref_out).abs().mean().item() < 1.0
+
+
+@pytest.mark.parametrize("bits,group_size,symmetric", [(4, 16, False), (8, -1, True)])
+def test_fused_moe_forward_parity(bits, group_size, symmetric):
+ """Fused 3D-expert forward: index into a QuantTensor 3D weight per expert."""
+ torch.manual_seed(0)
+ num_experts, in_dim, hidden = 4, 32, 64
+
+ class FusedMoE(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.gate_up = nn.Parameter(torch.randn(num_experts, hidden, in_dim))
+ self.down = nn.Parameter(torch.randn(num_experts, in_dim, hidden))
+
+ def forward(self, x: torch.Tensor, expert_ids: torch.Tensor) -> torch.Tensor:
+ # x: (tokens, in_dim); expert_ids: (tokens,) routing
+ outs = []
+ for t in range(x.shape[0]):
+ e = int(expert_ids[t].item())
+ h = F.linear(x[t : t + 1], self.gate_up[e]) # (1, hidden)
+ outs.append(F.linear(F.silu(h), self.down[e])) # (1, in_dim)
+ return torch.cat(outs, dim=0)
+
+ model = FusedMoE().eval()
+ x = torch.randn(8, in_dim)
+ expert_ids = torch.tensor([0, 1, 2, 3, 0, 2, 1, 3])
+ ref_out = model(x, expert_ids)
+
+ # Quantize the 3D expert tensors directly via QuantTensor.from_float.
+ for pname in ("gate_up", "down"):
+ _quantize_inplace(model, pname, bits=bits, group_size=group_size, symmetric=symmetric)
+
+ quant_out = model(x, expert_ids)
+
+ # Reference: same forward against dense-materialized 3D weights.
+ dense_model = FusedMoE()
+ with torch.no_grad():
+ dense_model.gate_up = nn.Parameter(model.gate_up.data.to_dense(), requires_grad=False)
+ dense_model.down = nn.Parameter(model.down.data.to_dense(), requires_grad=False)
+ dense_out = dense_model(x, expert_ids)
+
+ torch.testing.assert_close(quant_out, dense_out, rtol=0, atol=0)
+ assert (quant_out - ref_out).abs().mean().item() < 5.0
+
+
+@pytest.mark.parametrize("bits,group_size,symmetric", [(4, 32, False), (4, 32, True), (8, 32, False)])
+def test_onnx_export_parity(tmp_path, bits, group_size, symmetric):
+ """Olive-quantized nn.Linear -> ONNX MatMulNBits -> onnxruntime matches eager."""
+ torch.manual_seed(0)
+ in_dim, out_dim = 64, 32
+
+ class LinearOnly(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.fc = nn.Linear(in_dim, out_dim, bias=False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.fc(x)
+
+ model = LinearOnly().eval()
+ _quantize_inplace(model.fc, "weight", bits=bits, group_size=group_size, symmetric=symmetric)
+
+ x = torch.randn(1, 4, in_dim)
+ eager_out = model(x).detach()
+
+ # Build an export-compatible variant (QuantLinearNbit wrapper) and export to ONNX.
+ export_model = make_export_compatible_quant(copy.deepcopy(model), dynamo=False)
+
+ onnx_path = tmp_path / "model.onnx"
+ torch.onnx.export(
+ export_model,
+ (x,),
+ onnx_path,
+ input_names=["input"],
+ output_names=["output"],
+ dynamic_axes={"input": {0: "batch", 1: "seq"}, "output": {0: "batch", 1: "seq"}},
+ opset_version=21,
+ dynamo=False,
+ )
+
+ onnx.checker.check_model(str(onnx_path))
+ sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+ ort_out = sess.run(["output"], {"input": x.numpy()})[0]
+
+ torch.testing.assert_close(torch.from_numpy(ort_out), eager_out, rtol=1e-2, atol=1e-2)
diff --git a/test/common/quant/test_utils.py b/test/common/quant/test_utils.py
index d5004db897..8e9ea928b1 100644
--- a/test/common/quant/test_utils.py
+++ b/test/common/quant/test_utils.py
@@ -63,7 +63,7 @@ def test_initialization(self, bits, symmetric, group_size, signed):
def test_invalid_bits(self):
"""Test that invalid bits raise AssertionError."""
- with pytest.raises(AssertionError, match="Only 4-bit and 8-bit quantization supported"):
+ with pytest.raises(AssertionError, match="Only 2-bit, 4-bit and 8-bit quantization supported"):
WeightQuantizer(bits=16, symmetric=True, group_size=0)
def test_midq_calculation(self):
From fbd60bd6e625a5ac73757805b9724941a12ca063 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:13:57 +0000
Subject: [PATCH 14/18] Initial plan
From 4c50d8b30979d9fe14df69995dd7dba902ca484a Mon Sep 17 00:00:00 2001
From: copilot
Date: Tue, 21 Jul 2026 23:34:12 +0000
Subject: [PATCH 15/18] Override precedence: insertion-order first-match; add
ReDoS-safe regex validator + tests
---
olive/common/quant/patterns.py | 134 +++++++++++++++++++++++++----
test/common/quant/test_hf_utils.py | 16 +++-
test/common/quant/test_patterns.py | 72 ++++++++++++++--
3 files changed, 196 insertions(+), 26 deletions(-)
diff --git a/olive/common/quant/patterns.py b/olive/common/quant/patterns.py
index d2fed4af44..120f7f5f77 100644
--- a/olive/common/quant/patterns.py
+++ b/olive/common/quant/patterns.py
@@ -13,6 +13,15 @@
- Keys prefixed with ``re:`` opt into regular-expression matching via
``re.fullmatch``.
+When multiple ``overrides`` keys match the same target, the **first**
+matching key in config (insertion) order wins (``match_override``); this
+is the finalized precedence rule — "longest / most specific pattern" is
+deliberately not used.
+
+``re:`` patterns are validated for catastrophic-backtracking safety
+before compilation (``_assert_regex_safe``): over-long patterns and
+nested unbounded quantifiers (e.g. ``(a+)+``) are rejected.
+
These helpers are the single source of truth for the matching logic and
are used by both ``OliveHfQuantizationConfig`` and the Olive walker.
"""
@@ -24,42 +33,135 @@
REGEX_PREFIX = "re:"
+# Safety bound for user-supplied ``re:`` patterns. ``re.fullmatch`` has no timeout and a
+# length cap alone does not prevent catastrophic backtracking (short adversarial patterns
+# such as ``(a+)+$`` can already hang matching). We reject patterns that combine a group
+# quantified with an unbounded quantifier (``*`` / ``+`` / ``{n,}``) with a body that itself
+# contains an unbounded quantifier or a top-level alternation — the classic ReDoS shapes
+# ``(a+)+``, ``(a*)*``, ``(a|a)*`` — in addition to enforcing a conservative length cap.
+_MAX_REGEX_LEN = 200
+
def is_regex_pattern(pattern: str) -> bool:
"""Return True if ``pattern`` opts into regex matching."""
return isinstance(pattern, str) and pattern.startswith(REGEX_PREFIX)
+def _skip_char_class(raw: str, i: int) -> int:
+ """Return the index just past a ``[...]`` character class starting at ``raw[i] == '['``."""
+ n = len(raw)
+ j = i + 1
+ if j < n and raw[j] == "^":
+ j += 1
+ if j < n and raw[j] == "]": # a literal ``]`` as the first class member
+ j += 1
+ while j < n and raw[j] != "]":
+ if raw[j] == "\\":
+ j += 1
+ j += 1
+ return j + 1
+
+
+def _brace_is_unbounded(raw: str, i: int) -> bool:
+ """Whether the ``{...}`` quantifier starting at ``raw[i] == '{'`` has no finite upper bound."""
+ close = raw.find("}", i)
+ if close == -1:
+ return False
+ inner = raw[i + 1 : close]
+ # ``{n,}`` (no upper bound) is unbounded; ``{n}`` / ``{n,m}`` are bounded.
+ return inner.endswith(",") or ("," in inner and inner.split(",", 1)[1] == "")
+
+
+def _body_has_unbounded(body: str) -> bool:
+ """Whether a group body contains an unbounded quantifier or a top-level alternation."""
+ depth = 0
+ i = 0
+ n = len(body)
+ while i < n:
+ c = body[i]
+ if c == "\\":
+ i += 2
+ continue
+ if c == "[":
+ i = _skip_char_class(body, i)
+ continue
+ if c == "(":
+ depth += 1
+ elif c == ")":
+ depth = max(0, depth - 1)
+ elif depth == 0 and c == "|":
+ return True
+ elif c in "*+":
+ return True
+ elif c == "{" and _brace_is_unbounded(body, i):
+ return True
+ i += 1
+ return False
+
+
+def _assert_regex_safe(raw: str) -> None:
+ """Reject regex patterns prone to catastrophic backtracking (ReDoS).
+
+ Raises ``ValueError`` for over-long patterns and for nested unbounded quantifiers.
+ """
+ if len(raw) > _MAX_REGEX_LEN:
+ raise ValueError(
+ f"Regex override/skip pattern is too long ({len(raw)} > {_MAX_REGEX_LEN} chars); "
+ "keep patterns short to bound matching cost."
+ )
+ stack: list[int] = []
+ i = 0
+ n = len(raw)
+ while i < n:
+ c = raw[i]
+ if c == "\\":
+ i += 2
+ continue
+ if c == "[":
+ i = _skip_char_class(raw, i)
+ continue
+ if c == "(":
+ stack.append(i)
+ elif c == ")":
+ start = stack.pop() if stack else -1
+ nxt = raw[i + 1] if i + 1 < n else ""
+ quantified_unbounded = nxt in ("*", "+") or (nxt == "{" and _brace_is_unbounded(raw, i + 1))
+ if quantified_unbounded and start != -1 and _body_has_unbounded(raw[start + 1 : i]):
+ raise ValueError(
+ f"Regex override/skip pattern {raw!r} contains a nested unbounded quantifier "
+ "(e.g. `(a+)+`), which is vulnerable to catastrophic backtracking and is "
+ "rejected. Rewrite the pattern without nested `*`/`+`/`{n,}` quantifiers."
+ )
+ i += 1
+
+
@cache
def _compiled(pattern: str) -> re.Pattern:
- return re.compile(pattern[len(REGEX_PREFIX) :])
+ raw = pattern[len(REGEX_PREFIX) :]
+ _assert_regex_safe(raw)
+ return re.compile(raw)
def match_override(name: str, patterns) -> str | None:
- """Find the best override pattern for ``name``.
+ """Find the override pattern for ``name`` using first-match-wins semantics.
- Plain string patterns match by **literal equality**. ``re:`` patterns
- match by ``re.fullmatch``. The most specific match (longest pattern
- string) wins; ties are broken lexicographically so the choice is
- deterministic.
+ Patterns are evaluated in insertion order (config order); the **first** matching pattern
+ wins. Plain string patterns match by **literal equality**; ``re:`` patterns match by
+ ``re.fullmatch``. Insertion order (not "longest / most specific pattern") is the finalized,
+ documented precedence rule so that overlapping overrides resolve deterministically.
- Returns the original pattern (with prefix preserved) so the caller
- can look it up in the underlying overrides dict, or ``None`` if no
- pattern matched.
+ Returns the original pattern (with prefix preserved) so the caller can look it up in the
+ underlying overrides dict, or ``None`` if no pattern matched.
"""
if not patterns:
return None
- matches: list[str] = []
for pattern in patterns:
if is_regex_pattern(pattern):
if _compiled(pattern).fullmatch(name):
- matches.append(pattern)
+ return pattern
elif name == pattern:
- matches.append(pattern)
- if not matches:
- return None
- # longest pattern first, then lexical
- return sorted(matches, key=lambda p: (-len(p), p))[0]
+ return pattern
+ return None
def match_skip(name: str, patterns) -> bool:
diff --git a/test/common/quant/test_hf_utils.py b/test/common/quant/test_hf_utils.py
index f54b3faa2b..23c5e282fa 100644
--- a/test/common/quant/test_hf_utils.py
+++ b/test/common/quant/test_hf_utils.py
@@ -534,12 +534,24 @@ def test_get_qlinear_init_args_with_regex_override(self):
init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1")
assert init_args == {"bits": 8, "symmetric": True, "group_size": 32}
- def test_literal_beats_unrelated_regex(self):
+ def test_first_matching_override_wins_in_insertion_order(self):
+ # Two overlapping overrides match the same target. Per the finalized precedence rule
+ # (design item 6), the FIRST key in insertion order wins — not the longer/literal one.
overrides = {
"re:.*\\.w1": {"bits": 8},
"model.layers.0.mlp.experts.0.w1": {"bits": 6},
}
config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128, overrides=overrides)
init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1")
- # literal (longer) wins
+ # The regex is first in insertion order, so it wins even though the literal is "more specific".
+ assert init_args["bits"] == 8
+
+ def test_reordering_overlapping_overrides_flips_the_winner(self):
+ # Same two overlapping overrides, literal placed first -> literal wins.
+ overrides = {
+ "model.layers.0.mlp.experts.0.w1": {"bits": 6},
+ "re:.*\\.w1": {"bits": 8},
+ }
+ config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128, overrides=overrides)
+ init_args = config.get_qlinear_init_args("model.layers.0.mlp.experts.0.w1")
assert init_args["bits"] == 6
diff --git a/test/common/quant/test_patterns.py b/test/common/quant/test_patterns.py
index 1a66287983..be3ab9dd9d 100644
--- a/test/common/quant/test_patterns.py
+++ b/test/common/quant/test_patterns.py
@@ -4,7 +4,7 @@
# --------------------------------------------------------------------------
import pytest
-from olive.common.quant.patterns import is_regex_pattern, match_override, match_skip
+from olive.common.quant.patterns import _assert_regex_safe, is_regex_pattern, match_override, match_skip
class TestIsRegexPattern:
@@ -34,22 +34,28 @@ def test_regex_fullmatch(self):
# not a fullmatch
assert match_override("prefix_layers.0.experts", ["re:layers\\.\\d+\\.experts"]) is None
- def test_longest_pattern_wins(self):
- # both match — the more specific (longer) wins
+ def test_first_matching_pattern_wins_in_insertion_order(self):
+ # both match — the first in config (insertion) order wins, regardless of length
patterns = ["re:.*\\.experts\\..*", "re:layers\\.\\d+\\.experts\\.gate_proj"]
match = match_override("layers.0.experts.gate_proj", patterns)
- assert match == "re:layers\\.\\d+\\.experts\\.gate_proj"
+ assert match == "re:.*\\.experts\\..*"
+ # reversing the order flips the winner
+ assert match_override("layers.0.experts.gate_proj", list(reversed(patterns))) == (
+ "re:layers\\.\\d+\\.experts\\.gate_proj"
+ )
def test_tie_break_is_deterministic(self):
- # same-length patterns: lexically smallest wins
+ # same-length patterns: first (and here only) match wins
patterns = ["re:foo.bar", "re:foo.baz"]
# only the first matches
assert match_override("foo.bar", patterns) == "re:foo.bar"
- def test_literal_beats_shorter_regex_when_longer(self):
+ def test_first_match_wins_even_when_later_pattern_is_more_specific(self):
+ # A broad ``re:.*`` placed first wins over a later literal, per insertion order.
patterns = ["re:.*", "model.embed_tokens"]
- # literal pattern (len 18) is longer than "re:.*" (len 5)
- assert match_override("model.embed_tokens", patterns) == "model.embed_tokens"
+ assert match_override("model.embed_tokens", patterns) == "re:.*"
+ # But a literal placed first wins over a later broad regex.
+ assert match_override("model.embed_tokens", ["model.embed_tokens", "re:.*"]) == "model.embed_tokens"
class TestMatchSkip:
@@ -81,3 +87,53 @@ def test_none_or_empty_patterns(self):
)
def test_parametrized(self, name, patterns, expected):
assert match_skip(name, patterns) is expected
+
+
+class TestRegexSafetyBound:
+ """Design item 8: ``re:`` patterns must have a real adversarial-input safety bound."""
+
+ @pytest.mark.parametrize(
+ "adversarial",
+ [
+ "(a+)+$",
+ "(a*)*$",
+ "(a+)*$",
+ "(a*)+$",
+ "(a|a)+$",
+ "(a|aa)+$",
+ "([a-z]+)+$",
+ "(x+x+)+y",
+ "(.*)*$",
+ "(a+){2,}",
+ ],
+ )
+ def test_rejects_nested_unbounded_quantifiers(self, adversarial):
+ with pytest.raises(ValueError, match="nested unbounded quantifier"):
+ _assert_regex_safe(adversarial)
+
+ def test_rejects_overlong_pattern(self):
+ with pytest.raises(ValueError, match="too long"):
+ _assert_regex_safe("a" * 500)
+
+ @pytest.mark.parametrize(
+ "safe",
+ [
+ "layers\\.\\d+\\.experts\\.gate_proj",
+ ".*\\.experts\\..*",
+ "model\\.layers\\.\\d+\\.mlp\\.(gate|up|down)_proj",
+ "(abc)+def", # quantified group but body has no unbounded quantifier
+ "a{2,4}b+",
+ "[a-z]+\\.[0-9]+",
+ ],
+ )
+ def test_accepts_safe_patterns(self, safe):
+ # Should not raise, and should compile / match through the public API.
+ _assert_regex_safe(safe)
+
+ def test_adversarial_pattern_rejected_through_match_apis(self):
+ # The safety bound is enforced when the pattern is actually used, not only when
+ # validated directly.
+ with pytest.raises(ValueError, match="nested unbounded quantifier"):
+ match_override("aaaaaaaaaaaaaaaaaaaa!", ["re:(a+)+$"])
+ with pytest.raises(ValueError, match="nested unbounded quantifier"):
+ match_skip("aaaaaaaaaaaaaaaaaaaa!", ["re:(a+)+$"])
From a0aa4cc9f4aa9d634ef2ab169c682115f7400a5c Mon Sep 17 00:00:00 2001
From: copilot
Date: Tue, 21 Jul 2026 23:37:28 +0000
Subject: [PATCH 16/18] Add MoE selection/tensor regression tests + PyTorch
Native RTN docs
---
docs/source/features/quantization.md | 99 ++++++++++++++++++++++++++++
test/common/quant/test_selection.py | 81 +++++++++++++++++++++++
test/common/quant/test_tensor.py | 36 ++++++++++
3 files changed, 216 insertions(+)
diff --git a/docs/source/features/quantization.md b/docs/source/features/quantization.md
index a5b5ec6fbe..471d483ecd 100644
--- a/docs/source/features/quantization.md
+++ b/docs/source/features/quantization.md
@@ -71,6 +71,105 @@ This pass supports ONNX models and can quantize `MatMul` and `Gather` nodes to 4
}
```
+## PyTorch Native RTN
+
+The `Rtn` pass applies RTN weight quantization directly to a PyTorch (Hugging Face) model, before any ONNX
+export. Unlike `OnnxBlockWiseRtnQuantization` (which operates on an already-exported ONNX graph), `Rtn` runs on
+the `HfModelHandler` and replaces the weight *storage* of quantizable parameters in place with a quantized
+tensor representation, while keeping the surrounding modules (`nn.Linear` / `nn.Embedding`, and MoE experts)
+otherwise unchanged. This lets you compose it with other PyTorch quantization passes (see the `Gptq` example
+below) and share settings via per-module `overrides`.
+
+By default `Rtn` quantizes `nn.Linear` weights and leaves the embeddings, the language-model head, and any
+Mixture-of-Experts (MoE) experts at full precision. Three independent category flags opt those groups in:
+
+| Flag | Default | Effect when `true` |
+| --- | --- | --- |
+| `lm_head` | `false` | Also quantize the language-model head. |
+| `embeds` | `false` | Also quantize the input embeddings. |
+| `moe` | `false` | Also quantize MoE expert weights (both 3D fused-parameter experts such as gpt-oss / newer Qwen3-MoE, and `nn.ModuleList`-style experts such as Mixtral). |
+
+The `moe` flag is **fail-closed**: when `moe` is `false`, every module under an experts subtree is skipped even
+if it looks like a plain `nn.Linear`. If the model config indicates an MoE architecture but Olive cannot resolve
+the experts subtree for that (unrecognized) architecture, the pass raises a clear error *before* modifying any
+parameter, rather than silently quantizing the experts.
+
+Only weight parameters are quantized. On fused-expert modules that also expose 2D bias parameters (e.g.
+gpt-oss's `gate_up_proj_bias` / `down_proj_bias`), the biases are left in full precision.
+
+### `moe` and ONNX export
+
+MoE quantization in Olive is **storage-only**: Olive does not export 3D fused-expert `QuantTensor`s to ONNX.
+Attempting to `torch.onnx.export` a model with 3D-quantized experts raises a clear error directing you to
+Mobius / ORT GenAI `ModelBuilder` for the experts. Non-MoE parts (attention projections, router/gate,
+embeddings, lm_head) still export through the existing `MatMulNBits` / `GatherBlockQuantized` path.
+
+### `modules_to_not_convert` and `overrides`
+
+`modules_to_not_convert` lists module-name patterns to exclude entirely, and `overrides` maps module-name
+patterns to per-module `{"bits", "symmetric", "group_size"}` settings. Both accept two key styles:
+
+- **Plain strings** keep the existing Hugging Face semantics (substring match for `modules_to_not_convert`,
+ literal match for `overrides`).
+- **`re:`-prefixed keys** are treated as regular expressions matched with `re.fullmatch` (e.g.
+ `"re:model\\.layers\\.\\d+\\.mlp\\..*"`).
+
+For safety, `re:` patterns are validated before use: overly long patterns and patterns with nested unbounded
+quantifiers (catastrophic-backtracking / ReDoS shapes such as `(a+)+`) are rejected with a clear error.
+
+### Resolution order and override precedence
+
+When multiple rules could apply to the same target, the **first** rule that matches wins, in this order:
+
+1. `modules_to_not_convert` (hard exclude)
+2. category flags (`lm_head` / `embeds` / `moe`) — hard excludes; `overrides` can never re-include what a
+ category flag skipped
+3. `overrides`
+4. pass-level defaults (`bits` / `group_size` / `sym`)
+
+When several `overrides` entries match the same target, precedence is **insertion order in the config, first
+match wins** — not "longest / most specific pattern". Order your `overrides` from most specific to least
+specific accordingly.
+
+### Example Configuration
+```json
+{
+ "type": "Rtn",
+ "bits": 4,
+ "group_size": 128,
+ "sym": false,
+ "embeds": true,
+ "moe": true,
+ "overrides": {
+ "re:.*\\.experts\\..*": { "bits": 4, "group_size": 32 }
+ }
+}
+```
+
+### Composing with `Gptq`
+
+`Rtn` can run on an already-quantized model, so you can quantize the transformer `nn.Linear` layers with a
+calibration-based pass such as `Gptq` first, then cover the parts `Gptq` doesn't handle (embeddings, lm_head,
+MoE experts) with `Rtn`:
+
+```json
+[
+ { "type": "Gptq" },
+ { "type": "Rtn", "moe": true, "embeds": true }
+]
+```
+
+The reverse order is not supported: calibration-based passes assume a clean full-precision starting point and
+will reject an already-quantized model.
+
+### Migration note: removal of `QuantLinear` / `QuantEmbedding`
+
+The previous `nn.Module` wrappers `olive.common.quant.nn.QuantLinear` and `QuantEmbedding` have been **removed**
+as a sanctioned breaking change. Quantized weights are now stored as a `QuantTensor` on the parameter itself
+rather than by swapping the parent module. State dicts saved by the previous pipeline continue to load
+correctly. Only models persisted with `torch.save(model)` (pickling live `QuantLinear` / `QuantEmbedding`
+instances) are affected — re-run the `Rtn` pass to regenerate such checkpoints, or save/load via `state_dict`.
+
## HQQ
`HQQ (Half-Quadratic Quantization)` is a fast, calibration-free weight quantization method that enables low-bit quantization of large models without relying on gradient-based optimization. Unlike data-dependent approaches like GPTQ, [HQQ](https://dropbox.github.io/hqq_blog/) uses half-quadratic splitting to minimize weight quantization error efficiently.
diff --git a/test/common/quant/test_selection.py b/test/common/quant/test_selection.py
index 325bd4f9c8..ece5b0d4ae 100644
--- a/test/common/quant/test_selection.py
+++ b/test/common/quant/test_selection.py
@@ -160,3 +160,84 @@ def get_layer_wrappers(self):
("experts.down_proj", (4, 16, 8)),
("experts.gate_up_proj", (4, 8, 16)),
]
+
+
+def _install_fake_wrapper(monkeypatch, experts_by_layer):
+ """Patch ModelWrapper.from_model to expose ``experts_by_layer`` (list of (module, name))."""
+ from olive.common.hf import wrapper as wrapper_mod
+
+ class FakeLayerWrapper:
+ def __init__(self, experts, name):
+ self._experts = experts
+ self._name = name
+
+ def get_experts(self, return_name=True):
+ return (self._experts, self._name) if return_name else self._experts
+
+ class FakeWrapper:
+ def __init__(self, model):
+ self.model = model
+
+ def get_layer_wrappers(self):
+ return [FakeLayerWrapper(e, n) for e, n in experts_by_layer]
+
+ monkeypatch.setattr(wrapper_mod.ModelWrapper, "from_model", classmethod(lambda cls, m: FakeWrapper(m)))
+
+
+def test_moe_enabled_yields_only_3d_weights_and_skips_2d_bias(monkeypatch):
+ """Regression (gpt-oss gap): 2D bias params on a fused experts module must NOT be quantized."""
+
+ class FusedExpertsWithBias(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False)
+ self.gate_up_proj_bias = nn.Parameter(torch.zeros(4, 8), requires_grad=False) # 2D bias
+ self.down_proj = nn.Parameter(torch.zeros(4, 16, 8), requires_grad=False)
+ self.down_proj_bias = nn.Parameter(torch.zeros(4, 16), requires_grad=False) # 2D bias
+
+ class _Model(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.experts = FusedExpertsWithBias()
+
+ m = _Model()
+ _install_fake_wrapper(monkeypatch, [(m.experts, "experts")])
+
+ targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True))
+ assert _names(targets) == ["experts.down_proj", "experts.gate_up_proj"]
+ # every yielded param is 3D (no 2D bias slipped in)
+ assert all(module._parameters[pname].dim() == 3 for module, pname, _ in targets)
+
+
+def test_modulelist_experts_moe_flag_controls_selection(monkeypatch):
+ """Regression (ModuleList bug): per-expert Linears are quantized iff ``moe=True``."""
+ m = _ExpertList()
+ _install_fake_wrapper(monkeypatch, [(m.experts, "experts")])
+
+ off = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False))
+ assert _names(off) == []
+
+ on = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True))
+ assert _names(on) == ["experts.0", "experts.1"]
+
+
+def test_fail_closed_when_moe_arch_but_experts_not_discovered(monkeypatch):
+ """``moe=False`` must fail closed for an MoE arch whose experts can't be resolved."""
+
+ class _MoEConfig:
+ num_local_experts = 8
+
+ class _Model(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.config = _MoEConfig()
+ self.some_linear = nn.Linear(8, 8, bias=False)
+
+ m = _Model()
+ # Wrapper resolves no experts (unrecognized architecture).
+ _install_fake_wrapper(monkeypatch, [])
+
+ import pytest
+
+ with pytest.raises(ValueError, match="Mixture-of-Experts"):
+ list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False))
diff --git a/test/common/quant/test_tensor.py b/test/common/quant/test_tensor.py
index 3f36cf4354..f3ab58243e 100644
--- a/test/common/quant/test_tensor.py
+++ b/test/common/quant/test_tensor.py
@@ -136,3 +136,39 @@ def test_linear_raises_when_in_onnx_export(self, w2d, monkeypatch):
monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True)
with pytest.raises(RuntimeError, match="QuantTensor cannot be traced"):
F.linear(x, qt)
+
+
+class TestQuantTensor3DExpertRouting:
+ def test_tensor_index_routing_preserves_quantized_storage(self, w3d):
+ """Advanced/tensor-index expert selection (how real MoE routes) stays quantized."""
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ expert_ids = torch.tensor([0, 2, 2, 1])
+ selected = qt[expert_ids]
+ assert isinstance(selected, QuantTensor)
+ assert selected.shape == (4, *w3d.shape[1:])
+ # Values match a dense gather.
+ ref = qt.to_dense()[expert_ids]
+ assert torch.allclose(selected.to_dense(), ref, atol=1e-6)
+
+ def test_list_index_routing_preserves_quantized_storage(self, w3d):
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ selected = qt[[0, 3]]
+ assert isinstance(selected, QuantTensor)
+ assert selected.shape == (2, *w3d.shape[1:])
+
+ def test_unsupported_3d_indexing_raises_instead_of_dequantizing(self, w3d):
+ """Multi-axis / advanced indexing that isn't leading-dim-only must raise, not OOM-dequant."""
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ with pytest.raises(RuntimeError, match="Unsupported indexing pattern"):
+ _ = qt[:, 0, :]
+
+ def test_unsupported_3d_op_raises_during_onnx_export(self, w3d, monkeypatch):
+ """Central guard: any 3D QuantTensor op reaching the dense fallback under export must raise.
+
+ ``torch.index_select`` is an op that is not individually special-cased, so it exercises the
+ central ``_maybe_dense`` rejection rather than an op-specific check.
+ """
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True)
+ with pytest.raises(RuntimeError, match="ModelBuilder|Mobius|MoE"):
+ _ = torch.index_select(qt, 1, torch.tensor([0, 1]))
From c41a61878c2e649eeb7c41f2a4e2742eb49018ef Mon Sep 17 00:00:00 2001
From: copilot
Date: Tue, 21 Jul 2026 23:51:13 +0000
Subject: [PATCH 17/18] Central 3D movement-op guard + real-HF Mixtral
round-trip test + docs
---
olive/common/quant/tensor.py | 47 +++++++++++++
test/common/quant/test_forward_parity.py | 85 ++++++++++++++++++++++++
test/common/quant/test_tensor.py | 27 +++++++-
3 files changed, 157 insertions(+), 2 deletions(-)
diff --git a/olive/common/quant/tensor.py b/olive/common/quant/tensor.py
index adf4c16104..d219941439 100644
--- a/olive/common/quant/tensor.py
+++ b/olive/common/quant/tensor.py
@@ -512,3 +512,50 @@ def _move(x: torch.Tensor) -> torch.Tensor:
return x.to(**move_kwargs) if move_kwargs else x
return self._apply_fn_to_data(_move)
+
+
+# ----------------------------------------------------------------------
+# Movement / view ops
+# ----------------------------------------------------------------------
+# Storage-only quantization intercepts ``F.linear`` / ``F.embedding`` (and integer
+# expert selection). Shape-movement ops (transpose / reshape / view / permute / …) are
+# *not* implemented against the packed layout: letting them fall through to the default
+# tensor-subclass machinery silently produces a malformed ``QuantTensor`` (constructed via
+# ``__new__`` without the packed metadata), which then fails deep inside an unrelated op.
+# We reject them centrally with a clear, actionable error instead — and under ONNX export
+# surface the same MoE-export guidance as every other unsupported 3D op.
+_UNSUPPORTED_MOVEMENT_FNS = (
+ torch.transpose,
+ torch.Tensor.transpose,
+ torch.t,
+ torch.Tensor.t,
+ torch.permute,
+ torch.Tensor.permute,
+ torch.swapaxes,
+ torch.Tensor.swapaxes,
+ torch.movedim,
+ torch.Tensor.movedim,
+ torch.reshape,
+ torch.Tensor.reshape,
+ torch.Tensor.view,
+ torch.flatten,
+ torch.Tensor.flatten,
+ torch.Tensor.expand,
+ torch.squeeze,
+ torch.Tensor.squeeze,
+ torch.unsqueeze,
+ torch.Tensor.unsqueeze,
+)
+
+
+@implements(*_UNSUPPORTED_MOVEMENT_FNS)
+def _unsupported_movement(*args, **kwargs):
+ self = next((a for a in args if isinstance(a, QuantTensor)), None)
+ if self is not None and self.dim() >= 3 and torch.onnx.is_in_onnx_export():
+ raise RuntimeError(_MOE_ONNX_EXPORT_MSG)
+ raise RuntimeError(
+ "Shape-movement / view ops (transpose, reshape, view, permute, flatten, expand, …) are "
+ "not supported on an Olive QuantTensor: quantization is storage-only and these ops would "
+ "require fully dequantizing the packed weight. Access the dense weight explicitly via "
+ "``.to_dense()`` if you really need to reshape it (e.g. outside a memory-sensitive path)."
+ )
diff --git a/test/common/quant/test_forward_parity.py b/test/common/quant/test_forward_parity.py
index 69aadcecbd..6c704fe438 100644
--- a/test/common/quant/test_forward_parity.py
+++ b/test/common/quant/test_forward_parity.py
@@ -163,3 +163,88 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
ort_out = sess.run(["output"], {"input": x.numpy()})[0]
torch.testing.assert_close(torch.from_numpy(ort_out), eager_out, rtol=1e-2, atol=1e-2)
+
+
+def _quantize_model_targets(model, *, moe, embeds, lm_head, bits=4, group_size=32, symmetric=False):
+ """Quantize every ``iter_quant_targets`` selection in place with a QuantTensor."""
+ from olive.common.quant.selection import iter_quant_targets
+
+ for module, pname, _ in list(
+ iter_quant_targets(
+ model,
+ quantize_lm_head=lm_head,
+ quantize_embeds=embeds,
+ quantize_moe=moe,
+ )
+ ):
+ _quantize_inplace(module, pname, bits=bits, group_size=group_size, symmetric=symmetric)
+
+
+def test_real_hf_mixtral_moe_round_trip():
+ """Round-trip a real (config-only, no weight download) HF MoE architecture.
+
+ Builds a tiny ``MixtralForCausalLM`` from config, quantizes it with ``moe=True`` (the
+ fused 3D expert weights plus embeddings / lm_head), then saves and reloads the state
+ dict and verifies every quantized weight dequantizes bit-identically after reload.
+
+ Note: the model's own ``grouped_mm`` expert forward is not exercised here — Olive's
+ storage-only MoE quantization dispatches ``F.linear`` / ``F.embedding`` and does not
+ implement fused ``grouped_mm``; ONNX export / execution of the experts is delegated to
+ ORT GenAI ModelBuilder / Mobius. This test therefore validates the buffer-backed
+ save/reload path (the actual round-trip contract) on a real architecture.
+ """
+ transformers = pytest.importorskip("transformers")
+
+ from olive.common.quant.state_dict import refresh_quant_tensor_refs
+
+ def build():
+ torch.manual_seed(0)
+ cfg = transformers.MixtralConfig(
+ vocab_size=64,
+ hidden_size=32,
+ intermediate_size=32,
+ num_hidden_layers=1,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ num_local_experts=4,
+ num_experts_per_tok=2,
+ max_position_embeddings=64,
+ tie_word_embeddings=False,
+ )
+ return transformers.MixtralForCausalLM(cfg).eval()
+
+ model = build()
+ _quantize_model_targets(model, moe=True, embeds=True, lm_head=True)
+
+ # every fused expert weight must now be a QuantTensor; biases (if any) stay dense.
+ experts = model.model.layers[0].mlp.experts
+ assert isinstance(experts._parameters["gate_up_proj"].data, QuantTensor)
+ assert isinstance(experts._parameters["down_proj"].data, QuantTensor)
+
+ # Snapshot the dequantized value of every quantized parameter.
+ def dequant_snapshot(m):
+ snap = {}
+ for name, sub in m.named_modules():
+ for pname, param in sub._parameters.items():
+ if param is not None and isinstance(param.data, QuantTensor):
+ snap[f"{name}.{pname}"] = param.data.to_dense().clone()
+ return snap
+
+ before = dequant_snapshot(model)
+ assert any("experts" in k for k in before), "expected at least one quantized expert weight"
+
+ # Save state dict (QuantTensor params dropped; only plain buffers persisted) then reload
+ # onto a freshly quantized model and re-verify each dequantized weight is bit-identical.
+ state = model.state_dict()
+ assert not any(isinstance(v, QuantTensor) for v in state.values())
+
+ reloaded = build()
+ _quantize_model_targets(reloaded, moe=True, embeds=True, lm_head=True)
+ _missing, unexpected = reloaded.load_state_dict(state, strict=False)
+ assert not unexpected, f"unexpected keys on reload: {unexpected}"
+ refresh_quant_tensor_refs(reloaded)
+
+ after = dequant_snapshot(reloaded)
+ assert after.keys() == before.keys()
+ for key, ref in before.items():
+ torch.testing.assert_close(after[key], ref, rtol=0, atol=0)
diff --git a/test/common/quant/test_tensor.py b/test/common/quant/test_tensor.py
index f3ab58243e..3d624b6df7 100644
--- a/test/common/quant/test_tensor.py
+++ b/test/common/quant/test_tensor.py
@@ -170,5 +170,28 @@ def test_unsupported_3d_op_raises_during_onnx_export(self, w3d, monkeypatch):
"""
qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True)
- with pytest.raises(RuntimeError, match="ModelBuilder|Mobius|MoE"):
- _ = torch.index_select(qt, 1, torch.tensor([0, 1]))
+ with pytest.raises(RuntimeError, match=r"ModelBuilder|Mobius|MoE"):
+ torch.index_select(qt, 1, torch.tensor([0, 1]))
+
+ @pytest.mark.parametrize(
+ "op",
+ [
+ lambda qt: qt.transpose(-2, -1),
+ lambda qt: qt.reshape(4, -1),
+ lambda qt: qt.view(4, -1),
+ lambda qt: qt.permute(0, 2, 1),
+ lambda qt: qt.flatten(),
+ lambda qt: torch.transpose(qt, -2, -1),
+ ],
+ )
+ def test_movement_ops_raise_instead_of_silently_misbehaving(self, w3d, op):
+ """Shape-movement / view ops must raise a clear error rather than produce a malformed tensor."""
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ with pytest.raises(RuntimeError, match="storage-only"):
+ op(qt)
+
+ def test_movement_op_under_onnx_export_raises_moe_message(self, w3d, monkeypatch):
+ qt = QuantTensor.from_float(w3d, bits=4, symmetric=False, group_size=32)
+ monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True)
+ with pytest.raises(RuntimeError, match=r"ModelBuilder|Mobius|MoE"):
+ qt.transpose(-2, -1)
From 8035bb9265e6402861ef377cab7f35c903d344d4 Mon Sep 17 00:00:00 2001
From: copilot
Date: Wed, 22 Jul 2026 00:05:18 +0000
Subject: [PATCH 18/18] Add Gptq+Rtn composition test; make lintrunner clean
across changed files
---
olive/common/hf/quant.py | 1 +
olive/common/quant/patterns.py | 6 +--
olive/passes/pytorch/gptq.py | 8 +++-
olive/passes/pytorch/kquant.py | 7 +++-
test/common/quant/test_forward_parity.py | 7 ++--
test/common/quant/test_selection.py | 48 ++++++++++++++++++++++++
test/common/quant/test_tensor.py | 1 +
test/passes/pytorch/test_gptq.py | 1 +
test/passes/pytorch/test_kquant.py | 1 +
test/passes/pytorch/test_quant_utils.py | 6 +--
test/passes/pytorch/test_rtn.py | 1 +
11 files changed, 73 insertions(+), 14 deletions(-)
diff --git a/olive/common/hf/quant.py b/olive/common/hf/quant.py
index 8d9ceb0048..0c112e3435 100644
--- a/olive/common/hf/quant.py
+++ b/olive/common/hf/quant.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=attribute-defined-outside-init,protected-access
from __future__ import annotations
import logging
diff --git a/olive/common/quant/patterns.py b/olive/common/quant/patterns.py
index 120f7f5f77..d9cd372b71 100644
--- a/olive/common/quant/patterns.py
+++ b/olive/common/quant/patterns.py
@@ -89,11 +89,7 @@ def _body_has_unbounded(body: str) -> bool:
depth += 1
elif c == ")":
depth = max(0, depth - 1)
- elif depth == 0 and c == "|":
- return True
- elif c in "*+":
- return True
- elif c == "{" and _brace_is_unbounded(body, i):
+ elif (depth == 0 and c == "|") or c in "*+" or (c == "{" and _brace_is_unbounded(body, i)):
return True
i += 1
return False
diff --git a/olive/passes/pytorch/gptq.py b/olive/passes/pytorch/gptq.py
index 389ff7686e..e4c00d0507 100644
--- a/olive/passes/pytorch/gptq.py
+++ b/olive/passes/pytorch/gptq.py
@@ -130,7 +130,9 @@ def accumulate_hessian(module: torch.nn.Module, inp: tuple, _: Any) -> None:
batch_size = inp[0].shape[0]
inp = inp[0].reshape(-1, module.in_features).t()
- module.weight.quant_info.data["H"] *= module.weight.quant_info.data["N"] / (module.weight.quant_info.data["N"] + batch_size)
+ module.weight.quant_info.data["H"] *= module.weight.quant_info.data["N"] / (
+ module.weight.quant_info.data["N"] + batch_size
+ )
module.weight.quant_info.data["N"] += batch_size
inp = math.sqrt(2 / module.weight.quant_info.data["N"]) * inp.float()
module.weight.quant_info.data["H"] += inp.matmul(inp.t())
@@ -189,7 +191,9 @@ def process_module(
now_idx = 1
# create a per-channel quantizer
quantizer = WeightQuantizer(
- bits=module.weight.quant_info.quantizer.bits, symmetric=module.weight.quant_info.quantizer.symmetric, group_size=-1
+ bits=module.weight.quant_info.quantizer.bits,
+ symmetric=module.weight.quant_info.quantizer.symmetric,
+ group_size=-1,
)
if module.weight.quant_info.quantizer.group_size == -1:
# this can be before or after actorder permutation since there's only one group
diff --git a/olive/passes/pytorch/kquant.py b/olive/passes/pytorch/kquant.py
index 2e3b849da1..3413f70323 100644
--- a/olive/passes/pytorch/kquant.py
+++ b/olive/passes/pytorch/kquant.py
@@ -18,7 +18,12 @@
from olive.passes import Pass
from olive.passes.pass_config import PassConfigParam
-from olive.passes.pytorch.quant_utils import _module_weight_has_quant_info, finalize, get_quantizer_config, prepare_model
+from olive.passes.pytorch.quant_utils import (
+ _module_weight_has_quant_info,
+ finalize,
+ get_quantizer_config,
+ prepare_model,
+)
if TYPE_CHECKING:
from olive.hardware.accelerator import AcceleratorSpec
diff --git a/test/common/quant/test_forward_parity.py b/test/common/quant/test_forward_parity.py
index 6c704fe438..65784ab65b 100644
--- a/test/common/quant/test_forward_parity.py
+++ b/test/common/quant/test_forward_parity.py
@@ -1,3 +1,4 @@
+# pylint: disable=protected-access,not-callable
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
@@ -50,7 +51,7 @@ def _dense_reference(model: nn.Module) -> nn.Module:
return ref
-@pytest.mark.parametrize("bits,group_size,symmetric", [(4, 32, False), (4, -1, True), (8, 32, False)])
+@pytest.mark.parametrize(("bits", "group_size", "symmetric"), [(4, 32, False), (4, -1, True), (8, 32, False)])
def test_full_model_forward_parity(bits, group_size, symmetric):
torch.manual_seed(0)
@@ -80,7 +81,7 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
assert (quant_out - ref_out).abs().mean().item() < 1.0
-@pytest.mark.parametrize("bits,group_size,symmetric", [(4, 16, False), (8, -1, True)])
+@pytest.mark.parametrize(("bits", "group_size", "symmetric"), [(4, 16, False), (8, -1, True)])
def test_fused_moe_forward_parity(bits, group_size, symmetric):
"""Fused 3D-expert forward: index into a QuantTensor 3D weight per expert."""
torch.manual_seed(0)
@@ -123,7 +124,7 @@ def forward(self, x: torch.Tensor, expert_ids: torch.Tensor) -> torch.Tensor:
assert (quant_out - ref_out).abs().mean().item() < 5.0
-@pytest.mark.parametrize("bits,group_size,symmetric", [(4, 32, False), (4, 32, True), (8, 32, False)])
+@pytest.mark.parametrize(("bits", "group_size", "symmetric"), [(4, 32, False), (4, 32, True), (8, 32, False)])
def test_onnx_export_parity(tmp_path, bits, group_size, symmetric):
"""Olive-quantized nn.Linear -> ONNX MatMulNBits -> onnxruntime matches eager."""
torch.manual_seed(0)
diff --git a/test/common/quant/test_selection.py b/test/common/quant/test_selection.py
index ece5b0d4ae..85ce9e3e09 100644
--- a/test/common/quant/test_selection.py
+++ b/test/common/quant/test_selection.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
"""Tests for ``olive.common.quant.selection.iter_quant_targets``."""
from __future__ import annotations
@@ -241,3 +242,50 @@ def __init__(self):
with pytest.raises(ValueError, match="Mixture-of-Experts"):
list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=False))
+
+
+def test_gptq_then_rtn_moe_composition_skips_already_quantized(monkeypatch):
+ """Regression for `Gptq`-then-`Rtn(moe=True, embeds=True)` composition.
+
+ Emulates GPTQ having quantized only the `nn.Linear` layers first (they become
+ ``QuantTensor``-backed), then runs the RTN selection with ``moe=True`` / ``embeds=True``:
+ the already-quantized Linears must be skipped (kept as their GPTQ tensors) while the MoE
+ experts and embeddings are newly selected — no conflict, no double quantization.
+ """
+ from olive.common.quant.tensor import QuantTensor
+
+ class FusedExperts(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.gate_up_proj = nn.Parameter(torch.zeros(4, 8, 16), requires_grad=False)
+ self.down_proj = nn.Parameter(torch.zeros(4, 16, 8), requires_grad=False)
+
+ class _MoEModel(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.embed_tokens = nn.Embedding(16, 8)
+ self.q_proj = nn.Linear(8, 8, bias=False)
+ self.o_proj = nn.Linear(8, 8, bias=False)
+ self.experts = FusedExperts()
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ m = _MoEModel()
+ _install_fake_wrapper(monkeypatch, [(m.experts, "experts")])
+
+ # Emulate GPTQ: quantize only the nn.Linear weights (8-bit here so we can tell them apart).
+ for linear in (m.q_proj, m.o_proj):
+ qt = QuantTensor.from_float(linear.weight.data.clone(), bits=8, group_size=8, symmetric=True)
+ linear.weight = nn.Parameter(qt, requires_grad=False)
+
+ targets = list(iter_quant_targets(m, quantize_lm_head=True, quantize_embeds=True, quantize_moe=True))
+
+ # RTN must only pick up what GPTQ left un-quantized: the embeddings and the MoE experts.
+ assert _names(targets) == ["embed_tokens", "experts.down_proj", "experts.gate_up_proj"]
+
+ # The GPTQ-quantized Linears are untouched (still 8-bit QuantTensor).
+ assert isinstance(m.q_proj.weight.data, QuantTensor)
+ assert m.q_proj.weight.data.bits == 8
+ assert isinstance(m.o_proj.weight.data, QuantTensor)
+ assert m.o_proj.weight.data.bits == 8
diff --git a/test/common/quant/test_tensor.py b/test/common/quant/test_tensor.py
index 3d624b6df7..87e92f167a 100644
--- a/test/common/quant/test_tensor.py
+++ b/test/common/quant/test_tensor.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=redefined-outer-name,not-callable
import pytest
import torch
import torch.nn as nn
diff --git a/test/passes/pytorch/test_gptq.py b/test/passes/pytorch/test_gptq.py
index 9ec4d24bc5..4652428502 100644
--- a/test/passes/pytorch/test_gptq.py
+++ b/test/passes/pytorch/test_gptq.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
from pathlib import Path
import pytest
diff --git a/test/passes/pytorch/test_kquant.py b/test/passes/pytorch/test_kquant.py
index b824a9fe9f..087b834917 100644
--- a/test/passes/pytorch/test_kquant.py
+++ b/test/passes/pytorch/test_kquant.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
from pathlib import Path
import pytest
diff --git a/test/passes/pytorch/test_quant_utils.py b/test/passes/pytorch/test_quant_utils.py
index e7edaaf4e9..00f08ba9cf 100644
--- a/test/passes/pytorch/test_quant_utils.py
+++ b/test/passes/pytorch/test_quant_utils.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access,redefined-outer-name,not-callable
import logging
from copy import deepcopy
from types import SimpleNamespace
@@ -497,9 +498,6 @@ def test_prepare_model_locks_default_quantized_qkv_member_without_override(input
for V -- that would disagree with V's on-disk weights. Instead, Q/K should be demoted
to V's existing default config.
"""
- from olive.common.quant.state_dict import install_quant_tensor_param
- from olive.common.quant.tensor import QuantTensor
-
qu = quant_utils_module
existing = {
@@ -544,6 +542,8 @@ def fake_load(model_handler, **kwargs):
assert qcfg.get_qlinear_init_args(f"model.layers.0.self_attn.{proj}") == default
# No new override added for V (it stays at defaults on disk).
assert "model.layers.0.self_attn.v_proj" not in (qcfg.overrides or {})
+
+
# ---------------------------------------------------------------------------
# State-dict helper tests (install_quant_tensor_param, 2D + 3D fused MoE)
# ---------------------------------------------------------------------------
diff --git a/test/passes/pytorch/test_rtn.py b/test/passes/pytorch/test_rtn.py
index f05da40cc8..07447c0f72 100644
--- a/test/passes/pytorch/test_rtn.py
+++ b/test/passes/pytorch/test_rtn.py
@@ -2,6 +2,7 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
+# pylint: disable=protected-access
from pathlib import Path
import pytest