Extend PyTorch RTN weight quantization to MoE experts

[Copilot]

Describe your changes

Olive's native PyTorch quantization only handled nn.Linear / nn.Embedding, so MoE expert weights (the bulk of a MoE model's parameters) either stayed at full precision (fused-3D params were silently skipped) or were quantized whenever the pass ran with no opt-in. This resumes and completes the jambayk/moe-quant branch: storage-only MoE quantization via a QuantTensor tensor-subclass and a param-level quantization walk, merged onto current main.

Merge & invariants

• Merged jambayk/moe-quant onto main, reconciling the param-level walk with main's QKV-aware override renormalization and replacing the old isinstance(module, QuantLinear/QuantEmbedding) already-quantized check with a QuantTensor-based _collect_already_quantized_names.

Selection (common/quant/selection.py)

• New moe category flag (parallel to lm_head/embeds); when moe=False, every module under an experts subtree is skipped.
• Fused-expert targets now require dim() == 3, fixing gpt-oss 2D *_bias params being silently quantized.
• Fail-closed: raises before mutating any parameter when the config indicates MoE but experts discovery fails (no unfiltered fallback).

QuantTensor (common/quant/tensor.py)

• Central rejection of unsupported 3D ops under ONNX export (not just matmul/bmm).
• Leading-dim (expert) selection — including tensor/list index routing — stays on the quantized fast path; other 3D indexing raises instead of fully dequantizing (OOM guard).
• Shape-movement/view ops (transpose/reshape/view/permute/…) raise a clear error instead of producing a malformed subclass.

Override matching (common/quant/patterns.py)

• Precedence is now insertion-order, first-match-wins (removed longest-pattern rule).
• re: regex keys validated against catastrophic backtracking (length cap + nested-unbounded-quantifier rejection).

model_builder.py

• OliveQuantizedModel rejects moe=True checkpoints with a clear error; non-MoE checkpoints migrate <pname>_qweight/_scales/_qzeros keys so existing Linear/Embedding checkpoints keep loading.

Compatibility

• olive.common.quant.nn.QuantLinear/QuantEmbedding removed as a documented breaking change (state dicts still load; only torch.save(model) pickles are affected — re-run RTN).

Docs & tests

• New "PyTorch Native RTN" section in quantization.md (moe flag, re: syntax, precedence, migration note).
• Regression coverage: bias/ModuleList/fail-closed selection, tensor-index routing, export & movement-op guards, adversarial regex, Gptq→Rtn(moe=True, embeds=True) composition, and a real (config-only) Mixtral round-trip.

[
    { "type": "Gptq" },
    { "type": "Rtn", "moe": true, "embeds": true }
]

Note: transformers 5.14 Mixtral experts use fused grouped_mm, incompatible with storage-only F.linear dispatch by design; the real-HF test validates the buffer-backed save/reload (dequant parity) rather than the grouped_mm forward.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

Release note: The PyTorch Rtn pass gains a moe flag to quantize MoE expert weights (fused-3D and ModuleList experts), plus re:-prefixed regex support in overrides/modules_to_not_convert. QuantLinear/QuantEmbedding modules are removed (breaking change — see migration note).

(Optional) Issue link

Follow-up (out of scope): a tracked issue must be filed against the Mobius repo for preprocess_olive_weights silently mis-reshaping 3D expert weights as 2D.

[azure-pipelines]

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

[Copilot]

Pull request overview

This PR extends Olive’s native PyTorch RTN-style weight quantization to support Mixture-of-Experts (MoE) expert weights by moving from module-swapping (QuantLinear/QuantEmbedding) to a parameter-level, storage-only representation (QuantTensor) that can also handle fused 3D expert parameters. It also centralizes quantization target selection and pattern matching, updates ONNX export/model-builder integration boundaries, and adds extensive regression and parity tests.

Changes:

• Introduces QuantTensor + state-dict helpers to quantize weights (including fused 3D MoE expert tensors) without swapping parent modules, and removes olive.common.quant.nn.
• Adds centralized quantization target selection (iter_quant_targets) and pattern matching helpers (re: regex support with safety validation; insertion-order override precedence).
• Updates PyTorch passes (RTN/GPTQ/KQuant/AutoClip), ONNX ModelBuilder behavior, documentation, and adds broad test coverage (selection, tensor behavior, regex safety, parity, and model-builder rejection for MoE).

Reviewed changes

Copilot reviewed 28 out of 28 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
test/passes/pytorch/test_rtn.py	Updates RTN tests to assert quantization via QuantTensor-backed weights instead of QuantLinear/QuantEmbedding.
test/passes/pytorch/test_quant_utils.py	Adjusts quant-utils tests for parameter-level quant_info and adds new state-dict helper coverage.
test/passes/pytorch/test_kquant.py	Updates KQuant tests to use QuantTensor checks and bit assertions.
test/passes/pytorch/test_gptq.py	Updates GPTQ tests to validate QuantTensor-backed weights and composition behavior.
test/passes/onnx/test_model_builder.py	Adds coverage ensuring ModelBuilder rejects Olive-quantized MoE checkpoints.
test/common/quant/test_utils.py	Extends quant utils tests to validate N-D quantization and pack/unpack helpers.
test/common/quant/test_tensor.py	New tests for QuantTensor 2D/3D behavior, indexing, movement guards, and ONNX-export guards.
test/common/quant/test_selection.py	New tests for iter_quant_targets selection rules including MoE gating, fail-closed behavior, and already-quantized skipping.
test/common/quant/test_patterns.py	New tests for override/skip matching semantics and regex safety validation.
test/common/quant/test_nn.py	Removes tests for deprecated QuantLinear/QuantEmbedding module wrappers.
test/common/quant/test_hf_utils.py	Updates HF quantizer tests for QuantTensor-based layout, adds MoE and regex-override coverage.
test/common/quant/test_forward_parity.py	New numerical parity tests (eager vs dense reference) plus ONNX export parity for the export-compatible wrappers.
olive/passes/pytorch/rtn.py	Enables MoE support in RTN pass config via allow_moe=True.
olive/passes/pytorch/quant_utils.py	Refactors quantization to parameter-level selection, adds MoE config/options, and installs QuantTensor params during finalize.
olive/passes/pytorch/kquant.py	Adjusts KQuant to use weight-level quant_info and shared _module_weight_has_quant_info.
olive/passes/pytorch/gptq.py	Migrates GPTQ calibration/processing to store metadata on module.weight.quant_info.
olive/passes/pytorch/autoclip.py	Migrates AutoClip input caching and processing to use module.weight.quant_info.
olive/passes/onnx/model_builder.py	Rejects Olive-quantized MoE checkpoints; adds suffix mapping so Olive’s *_qweight/*_scales/*_qzeros buffers load into expected ModelBuilder attributes.
olive/common/quant/utils.py	Generalizes quantization and packing utilities to N-D tensors (quantize/pack along last dim).
olive/common/quant/tensor.py	New QuantTensor tensor-subclass implementing storage-only quantization with dispatch, 3D MoE indexing behavior, and ONNX/movement guards.
olive/common/quant/state_dict.py	New state-dict utilities for installing QuantTensor parameters + buffer aliases and refreshing references after load.
olive/common/quant/selection.py	New shared quantization target selection logic with MoE-aware rules and fail-closed behavior.
olive/common/quant/patterns.py	New pattern matching helpers for overrides/skip patterns with re: support and regex safety validation.
olive/common/quant/nn.py	Removes deprecated QuantLinear/QuantEmbedding module implementations.
olive/common/quant/hf_utils.py	Updates HF quantization config and quantizer implementation to use QuantTensor placeholders + state-dict refresh; adds regex overrides.
olive/common/hf/wrapper.py	Adds MoE experts/router conventions and accessors on LayerWrapper.
olive/common/hf/quant.py	Adds export-compatible wrappers creation from QuantTensor and improves symmetric handling (optional qzeros) for contrib ops.
docs/source/features/quantization.md	Documents PyTorch Native RTN, MoE flag behavior, regex semantics/safety, precedence rules, and migration away from QuantLinear/QuantEmbedding.

[Copilot]

Pull request overview

Copilot reviewed 28 out of 28 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

olive/common/quant/hf_utils.py:80

• The overrides docstring says "longest pattern wins", but get_qlinear_init_args uses match_override(...) which is insertion-order / first-match-wins. This is misleading and contradicts the behavior documented elsewhere in this PR (and tested in test_patterns.py).

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

olive/common/quant/selection.py:103

• The docstring states quantize_embeds=False only skips the input embedding module, but the implementation currently skips all nn.Embedding modules when quantize_embeds is false. The docstring should match the actual selection semantics (and the pass config description: "input embeddings").

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

olive/common/quant/selection.py:176

• When quantize_embeds=True, the current logic will quantize every nn.Embedding in the model (positional/token-type/etc.), not just the input embeddings. This contradicts the config/CLI description ("quantize the input embeddings") and can unintentionally change model behavior. Consider restricting embedding quantization to model.get_input_embeddings() when available.

        if isinstance(module, (nn.Linear, nn.Embedding)):
            if isinstance(module, nn.Embedding) and not quantize_embeds:
                continue
            if _is_skipped(module, name):
                continue

[titaiwangms]

Review team synthesis (5 parallel reviewers: readability / code / critical / deep / integration)

This PR was reviewed against the acceptance criteria in #2583. Several claims were independently verified/reproduced against the actual diff before being reported here. Overall: not yet mergeable — two Critical findings have concrete repros, plus several Major gaps in the OOM guard, legacy-checkpoint compatibility, and required tracked-issue/test coverage.

🔴 Critical

1. Override precedence flips across to_dict() serialization (spec deviation + load-time correctness bug)
hf_utils.py::to_dict() calls sort_layers_by_name(overrides), which reorders overrides into name-sorted order. The finalized rule (issue #2583 §6, and patterns.py::match_override's own docstring) is insertion-order, first-match-wins — explicitly not longest-pattern or sorted. The reload path (_process_model_before_weight_loading → get_qlinear_init_args → match_override) does depend on this order to pick bits/group_size for the placeholder QuantTensor.

Reproduced independently:

insertion order: ['re:.*\.w1' -> bits 8, 'model.layers.0.mlp.experts.0.w1' -> bits 6]
pre-serialize winner bits  = 8   (regex first, per insertion-order rule) ✔ correct
serialized overrides order = ['model.layers.0.mlp.experts.0.w1', 're:.*\.w1']  (sorted!)
post-roundtrip winner bits = 6   ← precedence FLIPPED

This means quantizing and saving a checkpoint, then reloading it, can silently resolve overlapping overrides differently — producing a shape/metadata mismatch or silently wrong dequantization on reload. Existing AC #4 regression tests only exercise the in-memory path, never a to_dict → reconstruct round-trip.
Fix: stop sorting overrides in to_dict() (preserve insertion order), and add a serialize→reload precedence regression test.

2. Real ReDoS bypass in the "safe" regex validator
_assert_regex_safe in patterns.py rejects nested unbounded quantifiers like (a+)+, but a pattern where the inner group is bounded and only the outer quantifier is unbounded — e.g. (a{1,2})+$ — passes validation while still causing catastrophic backtracking. Reproduced directly:

pattern = re.compile(r'(a{1,2})+$')
20 chars: 1.3ms
25 chars: 14.6ms
28 chars: 61ms
30 chars: 160ms   (clearly exponential)

A user-supplied re: override/skip pattern of this shape can hang model loading. The validator needs to also reject "quantified group whose body itself contains any repetition (bounded or unbounded) or alternation," not only nested-unbounded bodies.

🟠 Major

3. Eager-mode OOM guard does not hold for unregistered ops
The PR description claims "other 3D indexing raises instead of fully dequantizing" as an OOM guard, but this only holds under torch.onnx.is_in_onnx_export(). In plain eager mode, an unregistered op like torch.index_select(qt, 0, ids) falls through __torch_dispatch__ to to_dense(), fully materializing the 3D expert tensor — and the returned object is a malformed QuantTensor missing bits/qweight, causing a subsequent AttributeError on any further use. This should either raise consistently (in and out of export), or the OOM-guard claim in the description should be scoped down to export-time only.

4. Legitimate MoE routing index patterns are rejected
_indexes_leading_dim_only in tensor.py is overly restrictive: it rejects tuple-form indices like w[expert_ids, :, :] and rank->1 tensor indices (common for top-k routing), forcing them onto the slow/error path even though they only index the leading (expert) dimension. This will break real top-k MoE routing code that uses these forms.

5. Old QuantLinear/QuantEmbedding-era checkpoints cannot be reloaded
Verified directly in state_dict.py/hf_utils.py: there is no migration for legacy bare qweight/scales/qzeros buffer names (from the deleted classes) to the new <pname>_qweight convention on Olive's own HF-checkpoint reload path. This contradicts the PR's stated compatibility claim ("state dicts still load"). Note: this is a different code path from model_builder.py's OliveQuantizedModel, which does correctly migrate key names for ModelBuilder consumption — that path is fine; the gap is specifically in Olive's own HF model reload.

6. Fail-closed MoE detection has a partial-discovery gap
The fail-closed check in selection.py only fires when expert_modules is entirely empty. A hybrid model where some layers resolve their experts subtree and others don't will bypass the guard: with moe=False, nn.Linears in the unresolved subtree still get quantized; with moe=True, unresolved fused weights silently stay at full precision either way.

7. Gptq → Rtn(moe=True, embeds=True) composability is not tested end-to-end
The current regression test emulates GPTQ output by hand-installing QuantTensor placeholders and asserting iter_quant_targets skips them — it never runs the actual Gptq pass followed by the actual Rtn pass. Per #2583's "Composability with Gptq" section, this needs a true two-pass test verifying Linear layers keep GPTQ quantization while MoE experts/embeds get RTN-quantized without conflict (including the save/reload round-trip).

8. Required Mobius tracked issue not filed
#2583's acceptance criteria explicitly require a tracked issue (not a prose mention) documenting the preprocess_olive_weights 3D-mis-reshape risk. The PR description only states "a tracked issue must be filed" with no link — this AC item appears unmet.

9. Non-MoE key-migration path in model_builder.py has no regression test
The set_tensor suffix-migration logic (weight_qweight → qweight, etc.) that keeps existing Linear/Embedding ModelBuilder checkpoints loading is untested — only the moe=True-rejection path has a test. Per #2583 this test bullet was explicitly required.

🟡 Minor

• Stale/contradictory docstrings: hf_utils.py's OliveHfQuantizationConfig/get_qlinear_init_args docstrings still say "the longest pattern wins (ties broken lexically)," directly contradicting the shipped and documented first-match/insertion-order rule. Please sync with patterns.py's docstring and quantization.md.
• 2-bit RTN quantization cannot currently be exported to ONNX (QuantLinearNbit only supports 4/8-bit) — should be documented or rejected earlier with a clear error.
• The compiled-regex cache (@cache in patterns.py) is unbounded for the process lifetime; consider a bounded cache for long-lived services.
• olive/passes/pytorch/kquant.py (+ its test) is modified but not listed in Resume jambayk/moe-quant: extend RTN weight quantization to MoE experts #2583's enumerated in-scope file list — the change looks like a necessary mechanical ripple, but should be acknowledged explicitly rather than silently included.
• _assert_regex_safe validates lazily (only when a pattern is first matched against some module name); an invalid re: pattern that never matches anything is never validated. Consider validating all re: keys eagerly at config construction.

✅ What's solid

• N-D quantization math generalization is clean and verified: test_quantizer_3d_matches_2d_per_slice proves the 3D path is bit-identical to independently quantizing each expert slice across bits∈{2,4,8} and group_size∈{-1,16,32}.
• The gpt-oss 2D-bias mis-quantization bug is correctly fixed (dim() == 3 requirement) with a dedicated regression test.
• Fail-closed ordering is correct where it does fire (raises before any parameter is mutated).
• No remaining production references to the deleted QuantLinear/QuantEmbedding classes anywhere in the repo (verified via full-repo search).
• The central ONNX-export rejection path (_maybe_dense routing everything unregistered through a guard) correctly closes the "silent partial export" gap for movement ops and non-fast-path indexing under export.

Recommendation: address the two Critical items (override-serialization precedence, ReDoS bypass) and the eager-mode OOM guard / legacy-checkpoint-load gaps before merge; file the Mobius issue and add the two missing regression tests (model_builder non-MoE migration, true two-pass Gptq→Rtn composition) to close out the remaining acceptance criteria.