Add support for quant_auto format in model builder

[Vishal Jain (VishalX)]

Summary

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.

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

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Adds a new quantized checkpoint loader path for quant_method: quant_auto and updates the model builder so tied embeddings + pre-quantized lm_head can export correctly with ORT’s MatMulNBits / GatherBlockQuantized expectations.

Changes:

• Add a new QuantAutoModel loader and register it in the quantized-model dispatch.
• Extend the generic loader parsing logic to handle quant_auto tensor layout/splitting and .zeros-style naming.
• Update the ONNX builder to correctly reference/reshape pre-packed MatMulNBits scales/zero-points for tied quantized embeddings, plus add design notes.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
src/python/py/models/loaders/quant_model.py	Registers the new quant_auto loader in the factory dispatch.
src/python/py/models/loaders/quant_auto.py	Implements the quant_auto loader, embedding dequantization, and repacking to ORT format.
src/python/py/models/loaders/base.py	Adds quant_auto-specific name/shape handling in the shared safetensors parsing paths.
src/python/py/models/builders/base.py	Fixes tied-embedding export to reference the correct initializers and reshape packed scales/zeros when pre-quantized.
docs/quant-auto-support.md	Documents the quant_auto format and why loader/builder changes are needed.

Suppressed comments (1)

src/python/py/models/builders/base.py:2508

• Leading-underscore locals (_wlm, _wlm_prequantized) are disallowed by the model builder guidelines; rename them to non-underscored identifiers and update the subsequent references in this block.

            # 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:

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Suppressed comments (3)

Previously missed (2) — in code that hasn't changed since the last review.

src/python/py/models/loaders/base.py:296

• The embedding quantization-param skip is currently unconditional, but the comment says this is specific to quant_auto. Gating this on self.quant_type == "quant_auto" avoids silently dropping similarly-named tensors if a future format (or model variant) legitimately provides them.

                        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

docs/quant-auto-support.md:15

• The heading says "Six Layers of Incompatibility", but the document enumerates 7 sections (including "### 7. Tied lm_head..."). Please update the count to match the actual structure.

## Root Cause Analysis: Six Layers of Incompatibility

src/python/py/models/loaders/base.py:286

• normalize_vlm_weight_name() is added in this PR but never called. To keep the code/documentation consistent (and avoid dead normalization logic), apply it in the weight-loading loop before calling self.normalize_weight_name() (so per-loader overrides still run).

                # Map weights to modules
                for raw_name, tensor in weights.items():
                    name = self.normalize_weight_name(raw_name)
                    if name is None:
                        continue

[Vishal Jain (VishalX)]

Already fixed in commit 6122dfd (part of the previous review round). The call at base.py:284 was changed from QuantizedModel.normalize_weight_name(self, raw_name) to self.normalize_weight_name(raw_name), which enables the subclass dispatch and lets QuarkModel (and QuantAutoModel)'s overrides run correctly. CodeQL may have scanned an older commit.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

Previously missed (2) — in code that hasn't changed since the last review.

src/python/py/models/loaders/base.py:41

• normalize_vlm_weight_name() is defined but never referenced anywhere in the repository, so it adds dead code and can mislead readers (especially since the quant_auto .zeros handling is actually implemented via QuantAutoModel.normalize_weight_name() + QuantizedModel.__init__ calling self.normalize_weight_name). Consider removing this unused helper or wiring it into the actual normalization path.

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.

docs/quant-auto-support.md:25

• This section says the fix was “one line in normalize_vlm_weight_name()”, but that helper isn’t used by the loader. The working mechanism is (1) QuantizedModel.__init__ calling self.normalize_weight_name() so overrides apply, and (2) QuantAutoModel.normalize_weight_name() remapping .zeros → .qzeros. Updating the doc will prevent future confusion.

This issue also appears on line 159 of the same file.

**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.

**docs/quant-auto-support.md:159**
* This paragraph references `.zeros` normalization happening in `normalize_vlm_weight_name()`, but the loader path actually relies on `QuantAutoModel.normalize_weight_name()` (enabled by calling `self.normalize_weight_name` in the base loader). Please update the function name in the doc to match the implementation.

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.

</details>