diff --git a/docs/quant-auto-support.md b/docs/quant-auto-support.md new file mode 100644 index 0000000000..2e0659c06c --- /dev/null +++ b/docs/quant-auto-support.md @@ -0,0 +1,174 @@ +# quant_auto Support + +## Context + +A model using `quant_method: quant_auto` could not be exported through OGA's model builder. The format uses a custom quantization scheme developed for NPU deployment that differs from all four formats OGA previously supported (AWQ, GPTQ, Olive, Quark). + +**Files changed:** +- `src/python/py/models/loaders/quant_auto.py` — new `QuantAutoModel` loader +- `src/python/py/models/loaders/base.py` — `.zeros` → `.qzeros` name normalization +- `src/python/py/models/loaders/quant_model.py` — dispatch for `quant_auto` +- `src/python/py/models/builders/base.py` — tied-embedding export fixes + +--- + +## Root Cause Analysis: Six Layers of Incompatibility + +### 1. Tensor Naming — `.zeros` not recognized + +**Problem:** The quant_auto format uses `.zeros` as the zero-point suffix. OGA's regex patterns only match `(qzeros|weight_zero_point)`. Every `.zeros` tensor fell through to `NotImplementedError`. + +**Fix:** Added one line to `normalize_vlm_weight_name()`: +```python +name = re.sub(r'\.zeros$', '.qzeros', name) +``` +This remaps `.zeros` → `.qzeros` before any regex matching, requiring zero changes to the 400+ lines of existing pattern code. + +--- + +### 2. Fused-Layer Split Direction and Boundary + +**Problem — split direction:** `qkv_proj` and `gate_up_proj` are stored in `(out_features, in_features)` layout (same as Olive), requiring a dim=0 split. OGA's non-Olive path splits on dim=1, which produces the wrong result for GQA models where Q ≠ KV output size. + +**Fix:** Extended the existing Olive dim=0 split condition to include `quant_auto` for `qweight`: +```python +# Before: +if quant_type == "olive": +# After: +if quant_type in {"olive", "quant_auto"}: +``` + +**Problem — split boundary for scales/zeros:** Unlike Olive (which stores scales as a 2D `(out, n_groups)` matrix), quant_auto stores scales and zeros as a flat `(out*n_groups, 1)` column vector in output-first order. Splitting a flat `(491520, 1)` QKV scales tensor at row `q_size=3072` gives only 3072 rows instead of the required `q_size * n_groups = 294912` rows — leaving the other projections with wrong scale assignments. + +**Fix:** Separate branches for `quant_auto` scales/zeros that compute the split boundary as `q_size * ng` (not just `q_size`), preserving the flat output-first format that `pack_ort_format` expects: +```python +if quant_type == "quant_auto": + qkv_out = q_size + kv_size + kv_size + ng = tensor.shape[0] // qkv_out # infer n_groups from tensor shape + q_rows = q_size * ng + kv_rows = kv_size * ng + tensor_map["self_attn.q_proj.scales"] = tensor[:q_rows, :] + tensor_map["self_attn.k_proj.scales"] = tensor[q_rows : q_rows + kv_rows, :] + tensor_map["self_attn.v_proj.scales"] = tensor[q_rows + kv_rows :, :] +``` +Same pattern applied to qzeros for QKV, and scales/zeros for gate_up_proj using `intermediate_size * ng` as the midpoint. + +--- + +### 3. Container Format — No Bitwise Packing, uint8 Zero Points + +**Problem:** AWQ/GPTQ/Olive/Quark all store weights in packed int32 or uint8 containers (multiple int4 values per element). quant_auto stores weights as float16 with one integer value per element (range 0–15) — no packing at all. Calling the standard `unpack_on_row()` path on float16 data produces garbage via `torch.bitwise_right_shift` on floats. + +**Fix:** `QuantAutoModel.repack()` bypasses the unpack step entirely and calls `pack_ort_format()` directly after a dtype cast: +```python +def repack(self, module): + if module.qzeros is not None: + # Reshape flat (out*ng,1) → (out,ng) → transpose → (ng,out) so + # pack_zeros_ort_format's internal .T gives (out,ng) → ORT output-first layout + ng = module.qzeros.numel() // module.out_features + module.qzeros = module.qzeros.reshape(module.out_features, ng).T.to(torch.uint8).contiguous() + intweight = module.qweight.to(torch.int32) # F16 int values → int32 + self.pack_ort_format(module, intweight.T) # expects (in_features, out_features) +``` + +The zero point transpose is required because `pack_zeros_ort_format` applies its own `.T` internally before packing — so entering as `(ng, out)` produces the `(out, ng/2)` byte layout ORT's `MatMulNBits` kernel expects. + +--- + +### 4. Shape Inference — No Packing Factor + +**Problem:** `set_properties()` derives `in_features` from `qweight.shape[1] * 8 // bits` for Olive (which uses uint8 packing, so each column holds 2 values). For quant_auto, weight shape is `(out, in)` directly — no packing factor applies. + +**Fix:** `QuantAutoModel.set_properties()` reads shapes directly: +```python +proj.out_features = proj.qweight.shape[0] +proj.in_features = proj.qweight.shape[1] # no * 8 // bits +n_groups = proj.scales.reshape(proj.out_features, -1).shape[1] +proj.group_size = proj.in_features // n_groups +``` + +--- + +### 5. Config Loading — bits/group_size Absent + +**Problem:** `QuantizedModel.__init__` reads `config["bits"]` and `config["group_size"]` directly. The quant_auto config only contains `{"quant_method": "quant_auto"}` — accessing missing keys raises a `KeyError` immediately. + +**Fix:** `QuantAutoModel.__init__` extracts the values with `.get()` before calling `super().__init__()`, passing them as keyword arguments so the base class never touches the missing keys: +```python +global_bits = quant_attrs["config"].get("bits", 4) +global_group_size = quant_attrs["config"].get("group_size", -1) +super().__init__(..., global_bits=global_bits, global_group_size=global_group_size) +``` + +--- + +### 6. Embedding and lm_head Not Dequantized + +**Problem (discovered during inference testing):** The base class stores `model.embed_tokens.weight` as-is. For quant_auto, this is the raw int4 tensor (values 0.0–15.0 in float16). The model builder wrote these integer values directly into the ONNX embedding table, causing every token lookup to return a vector of values like `[3., 7., 8., 3., ...]` instead of proper float activations. This caused logit values of ~18,000 instead of the expected ±15 range. + +Additionally, `lm_head` (tied to the embedding via `tie_word_embeddings=True`) was being processed as a separate `QuantizedTensorModule` — packing the raw int4 embedding values through `pack_ort_format` and producing a `MatMulNBits` node with completely wrong scales. + +**Fix:** `QuantAutoModel.dequantize_embedding()` runs after `super().__init__()` and loads the embedding's scales and zeros from the safetensors files. The group size is inferred from tensor shapes (`ng = sc.numel() // vocab; gs = hidden // ng`). The **embedding Gather** table is dequantized to float16 so token lookups return proper activations: +```python +w_dq = ((w.float().reshape(-1, gs) - zp.float()) * sc.float()).reshape(w.shape).half() +self.embedding.weight = w_dq +``` + +The tied **lm_head** is handled separately — see next section for why it is kept quantized rather than dequantized. + +--- + +### 7. Tied lm_head: Native Asymmetric int4 (EOS Fix) + +**Problem (discovered during generation testing):** An earlier fix dequantized the tied lm_head to float16 and let the builder handle it. This produced runnable output but the model rarely emitted EOS — generations ran to the token cap with tail repetition / word-salad degeneration. + +**Root cause — the zero-point was silently dropped:** +1. A float16 lm_head has no `qweight`, so `make_matmul_int4` falls back to a plain float `MatMul` (base.py) carrying **no quantization metadata** — bits, group_size, and zero-point are all gone. +2. At save time `to_int4()` runs `MatMulNBitsQuantizer` with `is_symmetric=True` (the default), re-quantizing that float MatMul with a **symmetric** RTN grid — one scale per block, grid pinned to a fixed center, **no zero-point**. + +The QAT model was trained on an *asymmetric* grid. Collapsing to fp16 and re-fitting a symmetric grid introduced ~7% RMS logit noise, which suppressed the low-magnitude EOS logit → no-EOS / repetition. + +**Fix:** Keep the tied lm_head as a `QuantizedTensorModule` using the model's **own trained** scales/zeros, so it flows through the normal quantized path and exports as a 4-input **asymmetric** `MatMulNBits` (with a zero-point), exactly like every other linear layer. Properties are set here because `set_properties()` already ran during `super().__init__()`, when lm_head was still a plain `TensorModule`: +```python +lm = QuantizedTensorModule() +lm.qweight = w # (vocab, hidden) integer values stored as F16 +lm.scales = sc # (vocab*ng, 1) +lm.qzeros = zp # (vocab*ng, 1) +lm.out_features = w.shape[0] # vocab +lm.in_features = w.shape[1] # hidden +lm.bits = self.global_bits # 4 +lm.group_size = gs # 32 +self.lm_head = lm +``` +The shared `repack(self.lm_head)` in `__init__` then packs it. Because the node is already a `MatMulNBits`, the builder's `to_int4()` pass skips it — the trained zero-point is preserved end-to-end. + +--- + +## Factory Registration + +`QuantModel.from_pretrained()` dispatches `"quant_auto"` to the new class: +```python +elif quant_type == "quant_auto": + model = QuantAutoModel(quant_type, **kwargs) +``` + +--- + +## Backward Compatibility + +All four existing formats (AWQ, GPTQ, Olive, Quark) are **completely unchanged**. Every `quant_auto` code path is either a new `elif` branch or gated by `quant_type == "quant_auto"` / `quant_type in {"olive", "quant_auto"}`. The `.zeros` normalization in `normalize_vlm_weight_name()` is safe because no existing model tensor ends with the exact suffix `.zeros`. + +--- + +## quant_auto Tensor Format Reference + +| Tensor | Shape | dtype | Notes | +|---|---|---|---| +| `qkv_proj.weight` | `[out, in]` | F16 | `(out, in)` layout, values 0–15 | +| `qkv_proj.scales` | `[out*n_groups, 1]` | F16 | output-first flat order | +| `qkv_proj.zeros` | `[out*n_groups, 1]` | F16 | Same shape as scales | +| `gate_up_proj.weight` | `[2*intermediate_size, in]` | F16 | `2 × intermediate_size` rows | +| `down_proj.weight` | `[out, in]` | F16 | `(out, in)` | +| `embed_tokens.weight` | `[vocab, hidden]` | F16 | Quantized; dequantized at export | + +Key properties: weights stored as unpacked F16 integer values (one value per element), group_size inferred from scales shape, bits/group_size absent from config (use defaults: bits=4, group_size inferred per-module). diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index 6595827b9b..89cf8ee5df 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -1049,6 +1049,21 @@ def make_tied_quantized_embedding_input_names(self): # where rtn* = rtn, rtn_last # k_quant* = k_quant, k_quant_last, k_quant_linear, k_quant_mixed + # Pre-quantized lm_head (e.g. quant_auto): make_matmul_nbits registers weight under + # the MatMulNBits naming scheme rather than the to_nbits naming scheme. Return those + # names directly so make_embedding's GatherBlockQuantized references the right initializers. + # self.weights is the loaded model object (set in make_model before make_embedding runs). + wlm = getattr(getattr(self, "weights", None), "lm_head", None) + if wlm is not None and getattr(wlm, "qweight", None) is not None: + bits = wlm.bits + has_zeros = getattr(wlm, "qzeros", None) is not None + return ( + bits, + "lm_head.MatMulNBits.qweight", + "lm_head.MatMulNBits.scales", + "lm_head.MatMulNBits.qzeros" if has_zeros else "", + ) + base_method = self.quantization_algo placement = self.matmul_mixed_precision @@ -2486,9 +2501,27 @@ def make_embedding(self, embedding): # https://github.com/microsoft/onnxruntime/blob/0c9356cb986fd4cd2c5d510909d31186010ba226/onnxruntime/python/tools/quantization/neural_compressor/weight_only.py#L73 self.make_reshape(weight_reshape_name, weight_reshape_inputs, dtype=ir.DataType.UINT8, shape=[self.vocab_size, flat_dim]) input_names = [weight_reshape_output, self.input_names["input_ids"]] + # For pre-quantized lm_head, pack_ort_format flattens scales/zeros to 1D but + # GatherBlockQuantized requires the same rank as data (2D). Compute once here. + wlm = getattr(getattr(self, "weights", None), "lm_head", None) + wlm_prequantized = wlm is not None and getattr(wlm, "qweight", None) is not None if tied_weight_scale_name: + # Reshape scales from (vocab*ng,) to (vocab, ng). + if wlm_prequantized: + ng = wlm.scales.numel() // wlm.out_features + scale_reshape_name = f"{basename}/scales/Reshape" + self.make_reshape(scale_reshape_name, [tied_weight_scale_name, f"/model/constants/INT64/[{self.vocab_size}, {ng}]"], + dtype=self.io_dtype, shape=[self.vocab_size, ng]) + tied_weight_scale_name = f"{scale_reshape_name}/output_0" input_names.append(tied_weight_scale_name) if tied_weight_zp_name: + # Same rank requirement for zero points — reshape from 1D to 2D. + if wlm_prequantized: + ng_packed = wlm.qzeros.numel() // wlm.out_features + zp_reshape_name = f"{basename}/zeros/Reshape" + self.make_reshape(zp_reshape_name, [tied_weight_zp_name, f"/model/constants/INT64/[{self.vocab_size}, {ng_packed}]"], + dtype=ir.DataType.UINT8, shape=[self.vocab_size, ng_packed]) + tied_weight_zp_name = f"{zp_reshape_name}/output_0" input_names.append(tied_weight_zp_name) self.make_node( diff --git a/src/python/py/models/loaders/base.py b/src/python/py/models/loaders/base.py index 3cc4b6ce10..54a1054125 100644 --- a/src/python/py/models/loaders/base.py +++ b/src/python/py/models/loaders/base.py @@ -21,6 +21,26 @@ from safetensors.torch import load_file +def normalize_vlm_weight_name(name): + """Normalize a checkpoint tensor key for VLM/Quark conventions. + + Returns None if the tensor should be skipped (vision-tower weights), or + the normalized key string otherwise. + """ + # Skip vision tower weights in VLM checkpoints + if name.startswith(("model.visual.", "model.vision.", "visual.")): + return None + # Normalize common VLM prefix so existing LLM regex + parsing keeps working + if name.startswith("model.language_model."): + name = "model." + name[len("model.language_model."):] + # Normalize Quark weight_quantizer.* naming to flat weight_* naming + name = name.replace(".weight_quantizer.scale", ".weight_scale") + name = name.replace(".weight_quantizer.zero_point", ".weight_zero_point") + # Normalize quant_auto .zeros suffix to .qzeros + name = re.sub(r'\.zeros$', '.qzeros', name) + return name + + class QuantizedTensorModule: def __init__(self): self.qweight = None @@ -261,7 +281,7 @@ def __init__( # Map weights to modules for raw_name, tensor in weights.items(): - name = QuantizedModel.normalize_weight_name(self, raw_name) + name = self.normalize_weight_name(raw_name) if name is None: continue @@ -271,6 +291,9 @@ def __init__( if name == "model.embed_tokens.weight" or name == "transformer.embedding.word_embeddings.weight": self.embedding.weight = tensor + elif name in {"model.embed_tokens.scales", "model.embed_tokens.qzeros", "model.embed_tokens.g_idx"}: + # Embedding quantization params (quant_auto with tied weights); skip — embedding lookup uses float weights + continue elif name == "model.norm.weight" or name == "transformer.encoder.final_layernorm.weight": self.final_norm.weight = tensor elif name == "model.norm.bias" or name == "transformer.encoder.final_layernorm.bias": @@ -518,8 +541,8 @@ def __init__( # model.layers.layer_id.self_attention.query_key_value.qweight # model.layers.layer_id.self_attn.qkv_proj.weight # model.layers.layer_id.self_attention.query_key_value.weight - if quant_type == "olive": - # Olive: (out_features, in_features), split on dim=0 + if quant_type in {"olive", "quant_auto"}: + # Olive/QAT: (out_features, in_features), split on dim=0 tensor_map["self_attn.q_proj.qweight"] = tensor[:q_size, :] tensor_map["self_attn.k_proj.qweight"] = tensor[q_size : q_size + kv_size, :] tensor_map["self_attn.v_proj.qweight"] = tensor[q_size + kv_size :, :] @@ -540,7 +563,18 @@ def __init__( # model.layers.layer_id.self_attention.query_key_value.scales # model.layers.layer_id.self_attn.qkv_proj.weight_scale # model.layers.layer_id.self_attention.query_key_value.weight_scale - if quant_type == "olive": + if quant_type == "quant_auto": + # quant_auto: scales stored as (out*n_groups, 1) in output-first flat order. + # Split flat along dim=0 at q_size*ng and kv_size*ng boundaries to keep + # each slice in (out_i*n_groups, 1) flat form so pack_ort_format preserves order. + qkv_out = q_size + kv_size + kv_size + ng = tensor.shape[0] // qkv_out + q_rows = q_size * ng + kv_rows = kv_size * ng + tensor_map["self_attn.q_proj.scales"] = tensor[:q_rows, :] + tensor_map["self_attn.k_proj.scales"] = tensor[q_rows : q_rows + kv_rows, :] + tensor_map["self_attn.v_proj.scales"] = tensor[q_rows + kv_rows :, :] + elif quant_type == "olive": # Olive: (out_features, num_groups), split on dim=0 tensor_map["self_attn.q_proj.scales"] = tensor[:q_size, :] tensor_map["self_attn.k_proj.scales"] = tensor[q_size : q_size + kv_size, :] @@ -567,6 +601,15 @@ def __init__( tensor_map["self_attn.q_proj.qzeros"] = tensor[:q_dim, :] tensor_map["self_attn.k_proj.qzeros"] = tensor[q_dim : q_dim + kv_dim, :] tensor_map["self_attn.v_proj.qzeros"] = tensor[q_dim + kv_dim :, :] + elif quant_type == "quant_auto": + # quant_auto: zeros stored as (out*n_groups, 1) in output-first flat order. + qkv_out = q_size + kv_size + kv_size + ng = tensor.shape[0] // qkv_out + q_rows = q_size * ng + kv_rows = kv_size * ng + tensor_map["self_attn.q_proj.qzeros"] = tensor[:q_rows, :] + tensor_map["self_attn.k_proj.qzeros"] = tensor[q_rows : q_rows + kv_rows, :] + tensor_map["self_attn.v_proj.qzeros"] = tensor[q_rows + kv_rows :, :] else: # AWQ/GPTQ/Quark: int32 packing, split on dim=1 q_dim = ( @@ -605,8 +648,8 @@ def __init__( # model.layers.layer_id.mlp.dense_h_to_4h.qweight # model.layers.layer_id.mlp.gate_up_proj.weight # model.layers.layer_id.mlp.dense_h_to_4h.weight - if quant_type == "olive": - # Olive: (out_features, in_features), split on dim=0 + if quant_type in {"olive", "quant_auto"}: + # Olive/QAT: (out_features, in_features), split on dim=0 tensor_map["mlp.gate_proj.qweight"] = tensor[:intermediate_size, :] tensor_map["mlp.up_proj.qweight"] = tensor[intermediate_size:, :] else: @@ -628,7 +671,13 @@ def __init__( # model.layers.layer_id.mlp.dense_h_to_4h.scales # model.layers.layer_id.mlp.gate_up_proj.weight_scale # model.layers.layer_id.mlp.dense_h_to_4h.weight_scale - if quant_type == "olive": + if quant_type == "quant_auto": + # quant_auto: scales stored as (out*n_groups, 1) in output-first flat order. + ng = tensor.shape[0] // (2 * intermediate_size) + mid = intermediate_size * ng + tensor_map["mlp.gate_proj.scales"] = tensor[:mid, :] + tensor_map["mlp.up_proj.scales"] = tensor[mid:, :] + elif quant_type == "olive": # Olive: (out_features, num_groups), split on dim=0 tensor_map["mlp.gate_proj.scales"] = tensor[:intermediate_size, :] tensor_map["mlp.up_proj.scales"] = tensor[intermediate_size:, :] @@ -651,6 +700,12 @@ def __init__( intermediate_dim = intermediate_size // (8 // local_bits) tensor_map["mlp.gate_proj.qzeros"] = tensor[:intermediate_dim, :] tensor_map["mlp.up_proj.qzeros"] = tensor[intermediate_dim:, :] + elif quant_type == "quant_auto": + # quant_auto: zeros stored as (out*n_groups, 1) in output-first flat order. + ng = tensor.shape[0] // (2 * intermediate_size) + mid = intermediate_size * ng + tensor_map["mlp.gate_proj.qzeros"] = tensor[:mid, :] + tensor_map["mlp.up_proj.qzeros"] = tensor[mid:, :] else: # AWQ/GPTQ/Quark: int32 packing, split on dim=1 intermediate_dim = ( @@ -1061,3 +1116,4 @@ def pack_zeros_ort_format(self, module, reshape=False): module.qzeros = intzeros_pt.contiguous().byte() else: module.qzeros = intzeros_pt.contiguous() + diff --git a/src/python/py/models/loaders/quant_auto.py b/src/python/py/models/loaders/quant_auto.py new file mode 100644 index 0000000000..96b22c8095 --- /dev/null +++ b/src/python/py/models/loaders/quant_auto.py @@ -0,0 +1,157 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import os +import re + +import torch +from safetensors.torch import load_file + +from .base import QuantizedModel, QuantizedTensorModule + + +class QuantAutoModel(QuantizedModel): + """ + quant_auto format (quant_method: quant_auto): + - weight: (out_features, in_features) float16 containing integer values (no bitwise packing) + - scales: (out_features * num_groups, 1) float16 + - zeros: (out_features * num_groups, 1) float16 (renamed to qzeros by normalize_weight_name) + - bits and group_size are not in config; inferred from tensor shapes + """ + + def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers): + # quant_auto configs omit bits/group_size; supply defaults so the base class does not KeyError + global_bits = quant_attrs["config"].get("bits", 4) + global_group_size = quant_attrs["config"].get("group_size", -1) + super().__init__( + quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers, + global_bits=global_bits, global_group_size=global_group_size, + ) + + # Dequantize embedding: base class stored raw int4 F16 values; load scales/zeros and dequant + self.dequantize_embedding(input_path) + + for i, layer in enumerate(self.layers): + if i >= self.num_layers: + break + print(f"Repacking layer {i}") + self_attn = getattr(layer, "self_attn", None) or getattr(layer, "self_attention", None) + for _, q_tensors in self_attn.__dict__.items(): + if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None: + self.repack(q_tensors) + q_tensors.g_idx = None + for _, q_tensors in layer.mlp.__dict__.items(): + if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None: + self.repack(q_tensors) + q_tensors.g_idx = None + + if isinstance(self.lm_head, QuantizedTensorModule) and self.lm_head.qweight is not None: + self.repack(self.lm_head) + self.lm_head.g_idx = None + + def normalize_weight_name(self, name): + """Map .zeros suffix to .qzeros so existing loading patterns match.""" + name = super().normalize_weight_name(name) + if name is not None: + name = re.sub(r"\.zeros$", ".qzeros", name) + return name + + def dequantize_embedding(self, input_path): + """Set up the tied embedding and lm_head from the QAT model's int4 tensors. + + The embedding (and lm_head) is quantized in the QAT model: `weight` holds int4 + values (as F16), with trained `scales` and asymmetric `zeros`. The base class + skips the embedding's scales/zeros at load, so both must be read here directly. + + - Embedding Gather lookup: dequantized to fp16 (avoids raw int4 values in the + embedding table). + - lm_head (tied): kept as native asymmetric int4 using the model's OWN trained + scales/zeros, so it exports as a MatMulNBits with a zero-point like every other + linear layer. This prevents the builder's int4 pass from re-deriving a fresh + symmetric RTN scheme (which the QAT model was never trained for and which + suppresses the low-magnitude EOS token, causing no-EOS/repetition).""" + for weight_file in os.listdir(input_path): + if not weight_file.endswith(".safetensors"): + continue + weights = load_file(os.path.join(input_path, weight_file)) + # Accept either .zeros or .qzeros; require all three tensors in the same shard + zp_key = next( + (k for k in ("model.embed_tokens.zeros", "model.embed_tokens.qzeros") if k in weights), + None, + ) + if not all(k in weights for k in ("model.embed_tokens.weight", "model.embed_tokens.scales")) or zp_key is None: + continue + w = weights["model.embed_tokens.weight"] # (vocab, hidden) int4 values as F16 + sc = weights["model.embed_tokens.scales"] # (vocab*ng, 1) + zp = weights[zp_key] # (vocab*ng, 1) + + # Infer group size from tensor shapes rather than hard-coding + ng = sc.numel() // w.shape[0] + gs = w.shape[1] // ng + + # Embedding Gather uses dequantized fp16 weights + w_dq = ((w.float().reshape(-1, gs) - zp.float()) * sc.float()).reshape(w.shape).half() + self.embedding.weight = w_dq + + # lm_head (tied) uses the trained int4 quantization directly. Properties are + # set here (set_properties already ran in the base __init__ before this point, + # when lm_head was still a plain TensorModule) so the shared repack() applies. + lm = QuantizedTensorModule() + lm.qweight = w # (vocab, hidden) integer values stored as F16 + lm.scales = sc # (vocab*ng, 1) + lm.qzeros = zp # (vocab*ng, 1) + lm.out_features = w.shape[0] # vocab + lm.in_features = w.shape[1] # hidden + lm.bits = self.global_bits + lm.group_size = gs + self.lm_head = lm + return + + raise RuntimeError( + "dequantize_embedding: could not find model.embed_tokens.weight/scales/(q)zeros " + "in any safetensors shard. The embedding table cannot be dequantized." + ) + + def set_properties(self): + """Derive in_features, out_features, bits, and group_size from tensor shapes. + Weights are (out_features, in_features) F16 with 1 value per element — no packing factor.""" + def configure(proj): + if proj.qweight is None: + return + proj.out_features = proj.qweight.shape[0] + proj.in_features = proj.qweight.shape[1] + if proj.bits is None: + proj.bits = self.global_bits + # Infer group_size from scales shape: scales is (out * n_groups, 1) + n_groups = proj.scales.reshape(proj.out_features, -1).shape[1] + proj.group_size = proj.in_features // n_groups + self.set_g_idx(proj) + + if isinstance(self.lm_head, QuantizedTensorModule): + configure(self.lm_head) + + for module in self.layers: + for proj in [ + module.self_attn.q_proj, + module.self_attn.k_proj, + module.self_attn.v_proj, + module.self_attn.o_proj, + module.mlp.gate_proj, + module.mlp.up_proj, + module.mlp.down_proj, + ]: + configure(proj) + + def repack(self, module): + """Weights are already integer-valued F16; cast to int32 and repack to ORT format directly.""" + if module.qzeros is not None: + # Normalize qzeros to (ng, out) so pack_zeros_ort_format's .T gives (out, ng) + # and produces ORT's expected output-first packed layout. + # Fused layers arrive as (out, ng); unfused as (out*ng, 1) — flatten then reshape. + ng = module.qzeros.numel() // module.out_features + module.qzeros = module.qzeros.reshape(module.out_features, ng).T.to(torch.uint8).contiguous() + intweight = module.qweight.to(torch.int32) # (out, in) int values stored as F16 + self.pack_ort_format(module, intweight.T) # pack_ort_format expects (in_features, out_features) diff --git a/src/python/py/models/loaders/quant_model.py b/src/python/py/models/loaders/quant_model.py index 1e2dc004fd..b559f5d0dc 100644 --- a/src/python/py/models/loaders/quant_model.py +++ b/src/python/py/models/loaders/quant_model.py @@ -9,6 +9,7 @@ from .gptq import GPTQModel from .modelopt import ModeloptModel from .olive import OliveModel +from .quant_auto import QuantAutoModel from .quark import QuarkModel @@ -30,6 +31,8 @@ def from_pretrained(quant_type, **kwargs): model = QuarkModel(quant_type, **kwargs) elif quant_type in {"modelopt", "compressed-tensors"}: model = ModeloptModel(quant_type, **kwargs) + elif quant_type == "quant_auto": + model = QuantAutoModel(quant_type, **kwargs) else: raise NotImplementedError(f"The {quant_type} quantized model is not currently supported.") diff --git a/test/python/builder/test_tied_embeddings.py b/test/python/builder/test_tied_embeddings.py index c79111c197..e5b2a46ba4 100644 --- a/test/python/builder/test_tied_embeddings.py +++ b/test/python/builder/test_tied_embeddings.py @@ -297,6 +297,48 @@ def test_tied_quantized_embedding_weight_names_raise_for_unknown_algorithm(): model.make_tied_quantized_embedding_input_names() +def test_prequantized_lm_head_returns_matmul_nbits_names_with_zeros(): + """Pre-quantized lm_head (e.g. quant_auto) returns MatMulNBits initializer names.""" + lm_head = types.SimpleNamespace(qweight=object(), qzeros=object(), bits=4) + model = Model.__new__(Model) + model.weights = types.SimpleNamespace(lm_head=lm_head) + + bits, weight_name, scale_name, zp_name = model.make_tied_quantized_embedding_input_names() + + assert bits == 4 + assert weight_name == "lm_head.MatMulNBits.qweight" + assert scale_name == "lm_head.MatMulNBits.scales" + assert zp_name == "lm_head.MatMulNBits.qzeros" + + +def test_prequantized_lm_head_returns_matmul_nbits_names_without_zeros(): + """Pre-quantized lm_head without zero-points returns empty zp name.""" + lm_head = types.SimpleNamespace(qweight=object(), qzeros=None, bits=4) + model = Model.__new__(Model) + model.weights = types.SimpleNamespace(lm_head=lm_head) + + bits, weight_name, scale_name, zp_name = model.make_tied_quantized_embedding_input_names() + + assert bits == 4 + assert weight_name == "lm_head.MatMulNBits.qweight" + assert scale_name == "lm_head.MatMulNBits.scales" + assert zp_name == "" + + +def test_prequantized_lm_head_check_is_skipped_when_weights_not_loaded(): + """make_tied_quantized_embedding_input_names falls through to algo-based names when + self.weights hasn't been set yet (unit-test context without make_model).""" + model = Model.__new__(Model) + model.extra_options = {"algo_config": "rtn"} + model.quantization_algo, model.matmul_mixed_precision = desugar_algo_config(model.extra_options) + model.quant_attrs = {"is_symmetric": True, "matmul_block_size": 32} + # No model.weights set — simulates unit-test call without make_model + + bits, weight_name, _, _ = model.make_tied_quantized_embedding_input_names() + + assert weight_name == "lm_head.MatMul.weight_Q4G32" + + def _make_minimal_model_for_quantized_tied_embedding(*, algo_config, is_symmetric=True, quant_type=None): model = Model.__new__(Model) model.use_paged_attention = False