Add GPTQ quantization support for K-last MoE architectures

[Ti-Tai Wang (titaiwangms)]

Describe your changes

Adds GPTQ support for Mixture-of-Experts (MoE) models whose fused expert weights are stored
natively as (num_experts, out_features, in_features) — i.e. K (the reduction dim) is already
the last dimension, matching GPTQ's existing (OUT, K) assumption with no layout transpose
required ("K-last" architectures). This builds on #2584 (RTN MoE quantization / QuantTensor /
ModelWrapper infrastructure, not yet merged) and targets this exact PR's branch as its base.

Layout gating (updated): rather than a fixed model_type allow-list, GPTQ now delegates
fused-expert layout validation to the same shared, strict is_transposed-based guard used by
RTN (olive/passes/pytorch/moe_support.py::check_moe_layout_support) — trusting the metadata
transformers' own @use_experts_implementation decorator attaches, rather than a static
per-architecture list. Any experts module reporting (or defaulting to) a transposed
(E, K, OUT) layout, or missing/unverifiable is_transposed metadata (e.g. gpt_oss,
llama4, aria, or any architecture whose experts haven't adopted the fused-experts
decorator), fails closed with an actionable error message — those require a
layout-normalization step before GPTQ's Hessian accumulation and are intentionally out of
scope here (planned as a follow-up PR). GPTQ additionally requires the experts module to
be forward-interceptable (.config present, not a bare nn.ModuleList), since — unlike
RTN — it must record per-expert Hessians via a forward-hook swap.

Key design points

• Per-expert Hessian collection uses transformers' ALL_EXPERTS_FUNCTIONS registry
  (transformers >= 5.0), not monkey-patching: a single generic calibration function is
  registered once and swapped in per-model via set_experts_implementation, so every
  decorated Experts module routes through it uniformly — no per-architecture branching. Each
  expert gets its own independent (K, K) Hessian (no cross-expert pooling), and an explicit
  record on/off switch prevents double-recording on GPTQ's true-sequential second pass.
• Per-expert RTN fallback for cold/low-coverage experts, gated by a percentage-of-
  calibration-set threshold (moe_fallback_threshold, default 0.5%), following GPTQModel's
  convention. Guarantees GPTQ+fallback is never worse than plain RTN for a given expert,
  including the zero-sample case (no Hessian exists at all).
• Routing coverage report, logged per layer and at a run summary, derived from the
  forward call's own routing-index argument (never re-implements top-k routing, which would be
  actively wrong for architectures like DeepSeek-V3 that apply grouped, bias-corrected scoring
  before top-k).
• MoE routers are now unconditionally excluded from quantization (any form, including bare
  nn.Linear routers such as Jamba's), matching GPTQ Model/AWQ/vLLM convention. This is a
  shared-infrastructure change in iter_quant_targets and affects the existing RTN MoE path
  too, not just GPTQ.
• LayerWrapper (olive/common/hf/wrapper.py) gained MLP/router mapping entries for
  granitemoe (block_sparse_moe) and jamba (feed_forward), fixing a pre-existing crash
  for these two architectures that affected RTN as well as GPTQ.
• Fail-closed gating: transformers version + shared is_transposed-based layout guard
  (see above) + a forward-interceptability check (.config presence, rejecting bare
  nn.ModuleList experts) — verified before any weight mutation.
• A preflight warning estimates per-layer Hessian memory from the experts module's actual
  tensor shapes and warns above a threshold (DeepSeek-V3-scale configs can require tens of GB
  per layer); this is a known v1 limitation, not solved here.

This PR went through a 5-reviewer fan-out (readability, correctness, adversarial/critical, deep
spec-adherence, cross-module integration) plus a QA pass that executed concrete repros for the
highest-risk claims (exception safety, layout-gate bypass, RTN-fallback boundary). All Critical/
Major findings from that round were fixed in a follow-up commit, including:
exception-safe calibration lifecycle (state restoration on error), a re-entrancy guard and
registry-identity check on the experts-implementation swap, scoping the
get_attention_inputs() partial-resolution change back to only the call sites that need it (to
avoid silently breaking rotate.py's positional QKV assumption), the Hessian-memory preflight
warning, on-grid RTN-fallback weights before the true-sequential re-run, and coverage counts no
longer hardcoded to a specific parameter name. The allow-list + class-identity cross-check
originally used for layout gating was later replaced with the shared is_transposed-based
guard described above, to align with RTN's approach.

A subsequent full-PR review pass (same 5-reviewer + QA fan-out, but against the whole diff
rather than incrementally) found a few additional pre-existing issues predating this round —
most notably that get_mlp_inputs/get_mlp_outputs need the same partial_ok opt-in treatment
get_attention_inputs already received, to avoid rotate.py silently skipping MoE MLP rotation.
Fixes for these are being applied as a follow-up commit.

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.

This is user-facing (new Gptq pass capability for MoE models via allow_moe/moe=True +
moe_fallback_threshold) — release note: "Added GPTQ quantization support for Mixture-of-
Experts models with native (K-last) expert weight layouts: Qwen2-MoE, Qwen3-MoE, Phi-MoE,
Mixtral, DeepSeek-V3, Granite-MoE, OLMoE, and Jamba."

(Optional) Issue link

Builds on #2584. Related to #2599.

[Ti-Tai Wang (titaiwangms)]

CI failure diagnosis: test_gptq_moe_accepts_model_type_outside_old_allow_list — AttributeError: 'Qwen3MoeForCausalLM' object has no attribute 'get_experts_implementation'

Root cause: MoeCalibrationSession._start_owned() calls self.model.get_experts_implementation() to
snapshot the model's experts implementation before calibration. That method was only added to
transformers.PreTrainedModel in 5.14.0. test/requirements-test.txt currently pins
transformers<5.4.0, so CI installs a version that predates it entirely — hence the
AttributeError.

It goes deeper than just this one method, though: even with a polyfill for
get_experts_implementation (straightforward — config._experts_implementation /
config.sub_configs have existed since 5.0.0), set_experts_implementation()'s own validation
only became registry-aware (i.e. accepts a custom-registered implementation like our
"olive_moe_calib") starting at transformers 5.7.0; before that it hard-rejects anything
outside a fixed ["eager", "grouped_mm", "batched_mm", ...] list. So calibrated MoE quantization
(moe=True) is fundamentally incompatible with any transformers < 5.7.0, regardless of any
Olive-side workaround — this isn't a bug we can code around on our end below that floor.

Given #2622 already unpins test/requirements-test.txt's transformers<5.4.0 constraint entirely
(to resolve a security advisory), the practical fix here is to let that land and rebase/merge onto
main afterward, rather than adding compatibility shims for the 5.0–5.13 range in this PR. Once
#2622 merges, CI will install the latest transformers (currently 5.15.0), which has both
get_experts_implementation (5.14.0+) and the registry-aware set_experts_implementation
validation (5.7.0+), and this failure should resolve on its own.

No code change needed in this PR for this specific failure; tracking it against #2622 landing.

[Ti-Tai Wang (titaiwangms)]

cc Jambay Kinley (@jambayk) CI passes