Extend PyTorch RTN weight quantization to MoE experts

[Copilot]

Summary

Extends Olive's native PyTorch RTN quantization (olive/common/quant/,
olive/passes/pytorch/rtn.py) to cover MoE fused-expert weights, in addition
to the existing nn.Linear / nn.Embedding support. Produces a standard HF
safetensors checkpoint with MoE experts already quantized, so downstream
consumers (Mobius / ORT GenAI ModelBuilder) don't need to run their own
inline quantization pass.

Design (carried over from jambayk/moe-quant, unchanged)

• Storage-only MoE quantization. No QuantExperts module re-implementing
  per-architecture forwards (that design was considered and rejected — see
  the original design notes). Instead, a torch.Tensor wrapper subclass,
  QuantTensor (modelled on Quark's quark.qtensor.QTensor), replaces the
  quantized parameter in-place. Host model forwards run unchanged.
• Uniform by parameter, not by layer type. Target selection
  (selection.py) treats every nn.Parameter the same way regardless of
  rank — 2D linear/embedding weights and 3D fused-expert weights both flow
  through the same WeightQuantizer, which quantizes along the last dim
  regardless of rank.
• Mobius owns ONNX export of MoE experts. Olive never attempts
  torch.onnx.export on a 3D QuantTensor; this is an explicit, permanent
  design decision, documented in-code and via onnxruntime/mobius#427.

What changed in this PR on top of jambayk/moe-quant

1. selection.py (new) — single generator (iter_quant_targets)
   unifying target selection for both the HF quantizer and RTN/GPTQ passes.
   Adds MoE-aware routing detection (_collect_experts,
   _layers_missing_experts, _config_indicates_moe) with fail-closed
   behavior: if the config looks like an MoE architecture but the experts
   subtree can't be resolved, quantization refuses to proceed rather than
   silently skipping the expert weights.
2. wrapper.py — LayerWrapper.get_experts() / get_router()
   accessors, generalizing the existing per-layer-type accessor pattern to
   MoE sub-modules.
3. Round-1 remediation (8 items from the first adversarial review pass)
   — assorted correctness/robustness fixes to selection, patterns, and
   QuantTensor construction.
4. Round-2 remediation (4 items):
   • patterns.py: reject nested-group alternation ((a|b) inside a
     repeated group) at any nesting depth — closes a ReDoS bypass of the
     skip-pattern regex safety check. Docstrings demoted from "prevents
     ReDoS" to "best-effort UX check, not a security boundary" (decision:
     no regex third-party dependency added).
   • tensor.py / hf_utils.py / state_dict.py: explicit
     is_placeholder flag threaded through QuantTensor's lifecycle so
     init-style ops (zero_, normal_, ...) only no-op on real
     placeholders and raise otherwise (previously any QuantTensor
     silently no-op'd on these ops, which could mask real bugs).
   • tensor.py: reject rank>1 boolean-mask indexing instead of
     misclassifying it as a safe leading-dim integer index.
   • selection.py / defaults.yaml: rewrite _config_indicates_moe to
     reuse the existing resolve_alias() nested-config mechanism (already
     used for HF I/O config resolution) plus a bounded sub-config sweep —
     fixes DBRX-style nested MoE config detection (ffn_config.moe_num_experts).

All changes verified: 314 tests passing (test/common/quant/,
test/passes/pytorch/test_rtn.py), lintrunner clean.

Known follow-ups (tracked separately, not blocking this PR)

• MoE quantization: 5 unresolved edge-case bugs in QuantTensor indexing / regex safety checks #2598 — 5 edge-case bugs found in round-3 adversarial review
  (QuantTensor indexing safety, patterns.py regex safety). None are
  reachable via the currently-supported RTN pipeline (round-to-nearest,
  no live forward pass), so they don't block merging this PR, but should
  be fixed before any pass that runs a live forward through a quantized
  MoE QuantTensor (see MoE quantization: extend native GPTQ pass (gptq.py) to support MoE experts #2599).
• MoE quantization: extend native GPTQ pass (gptq.py) to support MoE experts #2599 — extend the native GPTQ pass to support MoE experts (requires
  redesigning the calibration forward-hook mechanism for fused-3D expert
  tensors — not a simple flag flip).
• MoE quantization: investigate AWQ (autoawq) support for MoE experts #2600 — investigate whether the autoawq pass can support MoE
  experts (depends on upstream autoawq library capability, not just
  Olive-side plumbing).

Testing

• pytest test/common/quant/ test/passes/pytorch/test_rtn.py — 314 passed.
• lintrunner — clean.
• Full forward-parity tests (test_forward_parity.py) included, comparing
  quantized vs. unquantized model outputs for both 2D and 3D (MoE) targets.

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

[titaiwangms]

Full-PR review pass (readability, code, critical, deep, integration reviewers + QA)

Per request, ran the full 6-agent review team over the entire PR diff (not just the recent commits), to confirm mergeability. This surfaced pre-existing bugs unrelated to #2598's original 5 items:

1 Critical

• Tied lm_head/embed_tokens embeddings got corrupted after checkpoint save→reload. Root cause: refresh_quant_tensor_refs visited every module hosting the shared QuantTensor and did last-write-wins re-binding, picking up a stale placeholder object instead of the reloaded one. Verified end-to-end (diff was ~2.26e18 before fix).

4 Major

• refresh_quant_tensor_refs silently left placeholder (zero) weights if a checkpoint was missing expected keys instead of failing. Fixed to fail closed with a RuntimeError, requiring all mandatory buffers (qweight+scales, +qzeros if asymmetric) to individually show complete load evidence — an initial round-1 fix used OR logic across buffers, which a targeted re-review (code + critical reviewers) caught as a false negative on partial buffer loads; fixed in round 2 with AND logic.
• ModelBuilder ignored regex overrides for per-layer bits/group_size (only did a plain dict lookup). Fixed to use match_override.
• QuantEmbeddingNbit had no ORT block_size validation; GatherBlockQuantized requires power-of-2 and >=16. Added validation to both QuantEmbeddingNbit and QuantLinearNbit.
• QuantEmbeddingNbit.from_quant_tensor had a scales/qzeros shape mismatch; fixed to reshape like the Linear equivalent.

Also fixed: a torch.equal override for QuantTensor — transformers 5.4's tie_weights() calls torch.equal on tied meta-device params before Olive's postprocess_model hook runs, which previously crashed.

All fixed across two rounds (comprehensive fix, then a targeted fix for the partial-buffer false negative found in re-review), each verified with fail-before/pass-after tests. Added 2 new test files (test/common/hf/test_quant.py, test/common/quant/test_state_dict.py) plus a real end-to-end tied-embedding save→reload regression test in test/passes/pytorch/test_rtn.py.

Current status: 448 tests passing across the affected test modules, lintrunner clean, working tree clean. Commits 3ba5b683 (item 1/4-strengthening/5-doc from #2598) and 6baed50e (this review's fixes) pushed.

Two Minor findings remain deferred as non-blocking follow-ups (not filed as issues yet): torch.equal override implements encoding-equality rather than value-equality for quantized tensors; named_modules(remove_duplicate=True) could theoretically hide a dual-registered module in an edge case not currently exercised.