From bf9ccadd4a50cdcf5c2948d8163a8ede611986c5 Mon Sep 17 00:00:00 2001
From: Vishal Jain <vishal.jain@amd.com>
Date: Fri, 28 Aug 2026 03:04:09 -0700
Subject: [PATCH 1/4] Add support for quant_auto format in model builder
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds QuantAutoModel to handle checkpoints using quant_method: quant_auto —
an asymmetric int4 format with no bitwise packing, where weights are stored
as float16 containing integer values (0-15) and scales/zeros as
(out_features * num_groups, 1) tensors in output-first flat order.

Key changes:

- loaders/quant_auto.py: new QuantAutoModel class with:
  - _dequantize_embedding(): fp16 embedding for Gather, native int4 for
    lm_head (preserves trained asymmetric zeros so EOS logit is not
    suppressed by symmetric re-quantization)
  - set_properties(): infers bits/group_size from tensor shapes
  - repack(): casts F16 int4 weights to int32 and calls pack_ort_format

- loaders/base.py: normalize_vlm_weight_name() maps .zeros to .qzeros
  so existing load paths handle the quant_auto naming convention

- loaders/quant_model.py: dispatch quant_auto to QuantAutoModel

- builders/base.py: fix tied-embedding export for pre-quantized lm_head:
  - make_tied_quantized_embedding_input_names() returns MatMulNBits
    initializer names (not to_nbits names) for pre-quantized models
  - make_embedding() inserts Reshape nodes to promote 1D scales/zeros
    (flattened by pack_ort_format) to 2D as required by GatherBlockQuantized

- docs/quant-auto-support.md: design and implementation notes

Co-Authored-By: Claude <noreply@anthropic.com>
---
 docs/quant-auto-support.md                  | 169 ++++++++++++++++++++
 src/python/py/models/builders/base.py       |  33 ++++
 src/python/py/models/loaders/base.py        |  68 +++++++-
 src/python/py/models/loaders/quant_auto.py  | 135 ++++++++++++++++
 src/python/py/models/loaders/quant_model.py |   3 +
 5 files changed, 402 insertions(+), 6 deletions(-)
 create mode 100644 docs/quant-auto-support.md
 create mode 100644 src/python/py/models/loaders/quant_auto.py

diff --git a/docs/quant-auto-support.md b/docs/quant-auto-support.md
new file mode 100644
index 0000000000..8ce9b8a050
--- /dev/null
+++ b/docs/quant-auto-support.md
@@ -0,0 +1,169 @@
+# quant_auto Support in quantized_model.py
+
+## 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).
+
+**File changed:** `src/python/py/models/quantized_model.py`
+
+---
+
+## 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:** The base `_load_quant_config()` 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._load_quant_config()` uses `.get()` with safe defaults:
+```python
+self.global_bits       = quant_attrs["config"].get("bits", 4)
+self.global_group_size = quant_attrs["config"].get("group_size", -1)  # -1 = infer per-module
+```
+
+---
+
+### 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 **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..876dd8d737 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(self.weights, "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(self.weights, "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..331b3bfa4f 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
@@ -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..5c0e5c613d
--- /dev/null
+++ b/src/python/py/models/loaders/quant_auto.py
@@ -0,0 +1,135 @@
+# -------------------------------------------------------------------------
+# 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 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_vlm_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):
+        super().__init__(quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers)
+
+        # 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 _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))
+            if "model.embed_tokens.scales" not in weights:
+                continue
+            w  = weights["model.embed_tokens.weight"]   # (vocab, hidden) int4 values as F16
+            sc = weights["model.embed_tokens.scales"]   # (vocab*ng, 1)
+            zp = weights["model.embed_tokens.zeros"]    # (vocab*ng, 1)
+            gs = 32
+
+            # 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
+            break
+
+    def _load_quant_config(self, quant_attrs):
+        # quant_auto config has no bits/group_size fields; use defaults, infer per-module in set_properties
+        self.global_bits = quant_attrs["config"].get("bits", 4)
+        self.global_group_size = quant_attrs["config"].get("group_size", -1)
+
+    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.")
 

From 6122dfde636a2a557e04dbc048ea2132607343c5 Mon Sep 17 00:00:00 2001
From: Vishal Jain <vishal.jain@amd.com>
Date: Fri, 28 Aug 2026 05:54:06 -0700
Subject: [PATCH 2/4] Address PR review comments for quant_auto loader

- loaders/base.py: use self.normalize_weight_name() (polymorphic) instead
  of QuantizedModel.normalize_weight_name(self, ...) so subclass overrides
  are respected; fixes ValueError on .zeros tensors during loading

- loaders/quant_auto.py:
  - pass global_bits/global_group_size explicitly to super().__init__ to
    avoid KeyError when quant_auto config omits these fields
  - override normalize_weight_name to map .zeros -> .qzeros
  - infer embedding group size from tensor shapes instead of hard-coding 32
  - rename _dequantize_embedding -> dequantize_embedding and
    _configure -> configure to follow project naming conventions

- builders/base.py:
  - rename _wlm/_wlm_prequantized to wlm/wlm_prequantized
  - guard self.weights access with getattr(self, "weights", None) so
    make_tied_quantized_embedding_input_names works in unit test context
    (fixes AttributeError when tests call it without make_model)

- docs/quant-auto-support.md: update file paths and method names to
  reflect the refactored loader structure

Co-Authored-By: Claude <noreply@anthropic.com>
---
 docs/quant-auto-support.md                 | 19 +++++++----
 src/python/py/models/builders/base.py      | 20 ++++++------
 src/python/py/models/loaders/base.py       |  2 +-
 src/python/py/models/loaders/quant_auto.py | 38 ++++++++++++++--------
 4 files changed, 48 insertions(+), 31 deletions(-)

diff --git a/docs/quant-auto-support.md b/docs/quant-auto-support.md
index 8ce9b8a050..2e0659c06c 100644
--- a/docs/quant-auto-support.md
+++ b/docs/quant-auto-support.md
@@ -1,10 +1,14 @@
-# quant_auto Support in quantized_model.py
+# 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).
 
-**File changed:** `src/python/py/models/quantized_model.py`
+**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
 
 ---
 
@@ -87,12 +91,13 @@ proj.group_size = proj.in_features // n_groups
 
 ### 5. Config Loading — bits/group_size Absent
 
-**Problem:** The base `_load_quant_config()` 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.
+**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._load_quant_config()` uses `.get()` with safe defaults:
+**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
-self.global_bits       = quant_attrs["config"].get("bits", 4)
-self.global_group_size = quant_attrs["config"].get("group_size", -1)  # -1 = infer per-module
+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)
 ```
 
 ---
@@ -103,7 +108,7 @@ self.global_group_size = quant_attrs["config"].get("group_size", -1)  # -1 = inf
 
 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 **embedding Gather** table is dequantized to float16 so token lookups return proper activations:
+**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
diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py
index 876dd8d737..89cf8ee5df 100644
--- a/src/python/py/models/builders/base.py
+++ b/src/python/py/models/builders/base.py
@@ -1053,10 +1053,10 @@ def make_tied_quantized_embedding_input_names(self):
         # 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(self.weights, "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
+        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",
@@ -2503,12 +2503,12 @@ def make_embedding(self, embedding):
             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(self.weights, "lm_head", None)
-            _wlm_prequantized = _wlm is not None and getattr(_wlm, "qweight", None) is not None
+            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
+                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])
@@ -2516,8 +2516,8 @@ def make_embedding(self, embedding):
                 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
+                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])
diff --git a/src/python/py/models/loaders/base.py b/src/python/py/models/loaders/base.py
index 331b3bfa4f..54a1054125 100644
--- a/src/python/py/models/loaders/base.py
+++ b/src/python/py/models/loaders/base.py
@@ -281,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
 
diff --git a/src/python/py/models/loaders/quant_auto.py b/src/python/py/models/loaders/quant_auto.py
index 5c0e5c613d..c850032740 100644
--- a/src/python/py/models/loaders/quant_auto.py
+++ b/src/python/py/models/loaders/quant_auto.py
@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 
 import os
+import re
 
 import torch
 from safetensors.torch import load_file
@@ -17,15 +18,21 @@ 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_vlm_weight_name)
+    - 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):
-        super().__init__(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)
+        self.dequantize_embedding(input_path)
 
         for i, layer in enumerate(self.layers):
             if i >= self.num_layers:
@@ -45,7 +52,14 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
             self.repack(self.lm_head)
             self.lm_head.g_idx = None
 
-    def _dequantize_embedding(self, input_path):
+    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
@@ -68,7 +82,10 @@ def _dequantize_embedding(self, input_path):
             w  = weights["model.embed_tokens.weight"]   # (vocab, hidden) int4 values as F16
             sc = weights["model.embed_tokens.scales"]   # (vocab*ng, 1)
             zp = weights["model.embed_tokens.zeros"]    # (vocab*ng, 1)
-            gs = 32
+
+            # 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()
@@ -88,15 +105,10 @@ def _dequantize_embedding(self, input_path):
             self.lm_head = lm
             break
 
-    def _load_quant_config(self, quant_attrs):
-        # quant_auto config has no bits/group_size fields; use defaults, infer per-module in set_properties
-        self.global_bits = quant_attrs["config"].get("bits", 4)
-        self.global_group_size = quant_attrs["config"].get("group_size", -1)
-
     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):
+        def configure(proj):
             if proj.qweight is None:
                 return
             proj.out_features = proj.qweight.shape[0]
@@ -109,7 +121,7 @@ def _configure(proj):
             self.set_g_idx(proj)
 
         if isinstance(self.lm_head, QuantizedTensorModule):
-            _configure(self.lm_head)
+            configure(self.lm_head)
 
         for module in self.layers:
             for proj in [
@@ -121,7 +133,7 @@ def _configure(proj):
                 module.mlp.up_proj,
                 module.mlp.down_proj,
             ]:
-                _configure(proj)
+                configure(proj)
 
     def repack(self, module):
         """Weights are already integer-valued F16; cast to int32 and repack to ORT format directly."""

From f22c6c7d3a63109a1d0b48b513a811f48342c067 Mon Sep 17 00:00:00 2001
From: Vishal Jain <vishal.jain@amd.com>
Date: Fri, 28 Aug 2026 08:11:32 -0700
Subject: [PATCH 3/4] Address follow-up review comments for quant_auto loader

- loaders/quant_auto.py: dequantize_embedding() now accepts both .zeros
  and .qzeros naming for the zero-point tensor, and guards against sharded
  checkpoints where not all embedding tensors appear in the same shard

- test/python/builder/test_tied_embeddings.py: add three tests covering
  the pre-quantized lm_head path in make_tied_quantized_embedding_input_names:
  with zeros, without zeros, and when weights are not yet loaded (unit-test
  context without make_model)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/python/py/models/loaders/quant_auto.py  |  9 ++++-
 test/python/builder/test_tied_embeddings.py | 42 +++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/python/py/models/loaders/quant_auto.py b/src/python/py/models/loaders/quant_auto.py
index c850032740..b95c69ba7c 100644
--- a/src/python/py/models/loaders/quant_auto.py
+++ b/src/python/py/models/loaders/quant_auto.py
@@ -77,11 +77,16 @@ def dequantize_embedding(self, input_path):
             if not weight_file.endswith(".safetensors"):
                 continue
             weights = load_file(os.path.join(input_path, weight_file))
-            if "model.embed_tokens.scales" not in weights:
+            # 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["model.embed_tokens.zeros"]    # (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]
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

From 688de82847e0dfef0a209123dbe99bd3823e0e9d Mon Sep 17 00:00:00 2001
From: Vishal Jain <vishal.jain@amd.com>
Date: Fri, 28 Aug 2026 09:54:31 -0700
Subject: [PATCH 4/4] Raise RuntimeError in dequantize_embedding when no shard
 matches

Without this, a misconfigured checkpoint would silently export raw int4
values as the embedding table with no indication of what went wrong.

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/python/py/models/loaders/quant_auto.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/python/py/models/loaders/quant_auto.py b/src/python/py/models/loaders/quant_auto.py
index b95c69ba7c..96b22c8095 100644
--- a/src/python/py/models/loaders/quant_auto.py
+++ b/src/python/py/models/loaders/quant_auto.py
@@ -108,7 +108,12 @@ def dequantize_embedding(self, input_path):
             lm.bits = self.global_bits
             lm.group_size = gs
             self.lm_head = lm
-            break
+            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.