Add MoE support to PyTorch KQuant - #2618
Merged
Merged
Conversation
PR #2584 merged RTN MoE quantization (moe=True) support without a layout check on fused-expert weights. RTN's WeightQuantizer groups unconditionally along a tensor's last dimension, which is only correct when the fused expert weight is stored (num_experts, out_features, in_features) -- K last. Architectures such as gpt-oss store the transposed (num_experts, in, out) layout instead, and would be silently mis-quantized (wrong axis grouped, no error) if run through moe=True today. Add olive/passes/pytorch/moe_support.py with check_moe_layout_support, gated into RTN's _run_for_config after prepare_model() and before finalize(). The check trusts transformers' own is_transposed attribute (set by the use_experts_implementation decorator, not derived from config/checkpoint data) directly: - Accept only experts modules that report is_transposed is False. - Reject anything where is_transposed is missing or not a bool (covers older transformers releases, undecorated architectures such as llama4/aria, and unrecognized implementations). - Reject is_transposed=True (e.g. gpt-oss). - Exempt classic per-expert nn.ModuleList experts (e.g. Mixtral/PhiMoE on older transformers) that carry no direct 3D parameter, since Olive's quantizer selection only ever groups a module's own 3D parameters and a bare ModuleList structurally cannot hold one. A trust_remote_code custom experts implementation that misreports its own is_transposed is out of scope: that is treated as user-introduced misuse of an explicitly opted-in trust boundary, not a layout Olive can independently verify. Also documents the moe flag's layout contract in docs/source/features/quantization.md and adds unit/integration tests covering the accept/reject/exempt paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
…vocation's request prepare_model ORs a pre-existing checkpoint's moe flag into the merged qcfg.moe (quant_utils.prepare_model), so re-running RTN with moe=False on an already MoE-quantized checkpoint incorrectly re-triggered the fused-experts layout safety check. Gate on the current invocation's own config.moe instead. Addresses the automated Copilot review comment on PR #2616. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
… invocation's request prepare_model ORs a pre-existing checkpoint's moe flag into the merged qcfg.moe (quant_utils.prepare_model), so re-running KQuant with moe=False on an already MoE-quantized checkpoint incorrectly re-triggered the fused-experts layout safety check. Gate on the current invocation's own config.moe instead, mirroring the same fix applied to RTN in PR #2616. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
Contributor
There was a problem hiding this comment.
Pull request overview
Extends Olive’s PyTorch-native KQuant pass to safely support Mixture-of-Experts (MoE) fused expert weights by generalizing qparam computation to N-D tensors and gating MoE quantization behind the shared fused-expert layout guard (check_moe_layout_support) so unsupported/transposed layouts fail closed.
Changes:
- Generalize
kquant_find_qparamsfrom 2D to N-D tensors (grouping always along the last dimension), enabling direct quantization of fused expert weights shaped like(E, OUT, K). - Add
moe=Truesupport toKQuantviaallow_moe, MoE layout gating withcheck_moe_layout_support, and correct parameter discovery using_iter_quant_info_params. - Add unit tests covering 3D qparam behavior, MoE layout rejection/exemptions, MoE round-trip sanity, and regression coverage for the
config.moevsqcfg.moegating bug; document KQuant MoE behavior.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| test/passes/pytorch/test_kquant.py | Adds N-D/MoE-focused unit tests (layout gate behavior, 3D qparams parity, round-trip sanity, gating regression). |
| olive/passes/pytorch/kquant.py | Enables MoE support, validates fused-expert layout when requested, and quantizes all selected parameters via _iter_quant_info_params; updates qparam logic to N-D. |
| docs/source/features/quantization.md | Documents KQuant capabilities and MoE layout safety behavior, including failure modes and a config example. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Addresses Copilot review feedback: the 'missing' layout case in test_kquant_moe_rejects_unsafe_or_unverifiable_layout previously called del experts.is_transposed unconditionally, which would raise AttributeError if the attribute were already absent (e.g. a different/ older transformers implementation), failing the test before it could exercise the intended missing-attribute code path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
This was referenced Aug 12, 2026
Jambay Kinley (jambayk)
approved these changes
Aug 13, 2026
Ti-Tai Wang (titaiwangms)
added a commit
that referenced
this pull request
Aug 13, 2026
…ized inference (#2620) ## Describe your changes Documents a limitation surfaced while validating KQuant/RTN MoE quantization on a real (locally-constructed) `Qwen3MoeForCausalLM` model: `transformers` may auto-select the `"grouped_mm"` experts implementation at inference time (even on CPU), which internally calls `weight.transpose(-2, -1)` on the fused-experts weight before its matmul kernel. Olive's 3D fused-expert `QuantTensor` is storage-only and cannot represent a transpose without a lossy dequantize/re-quantize round trip, so this raises a `RuntimeError` at inference time -- even for architectures whose checkpoint layout (`is_transposed=False`) is already fully supported for quantization by `Rtn`/`Gptq`/`KQuant`. This reproduces identically for both `Rtn` (merged, #2616) and `KQuant` (#2618), confirming it's a shared `QuantTensor` limitation rather than a pass-specific bug. Adds a short doc section next to the existing "`moe` and ONNX export" note, documenting the workaround (`model.set_experts_implementation("eager")` before running inference) and linking the follow-up issue. Follow-up issue: #2619 (tracks whether the deferred "transposed layout" (`is_transposed=True`) `QuantTensor` design work could also resolve this as a side effect, or whether it needs separate design). Doc-only change, no code/test changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
Ti-Tai Wang (titaiwangms)
added a commit
that referenced
this pull request
Aug 14, 2026
## Describe your changes Stacked on top of #2610 (this PR targets `b1-gptq-moe`, not `main`). Adds a manual validation script for comparing perplexity/size before and after quantizing a real (downloaded) HF checkpoint, plus three onboarding docs under `skills/olive/references/` for quantization work in this repo: - `scripts/quantize_and_compare_perplexity.py` — generic, pass-agnostic script (works for any registered Olive PyTorch quantization pass, not just GPTQ/MoE) that loads a real model, quantizes it, and reports weights-size and WikiText-2 perplexity deltas. This is the same style of real-model validation tool that surfaced real RTN bugs in #2584 after synthetic-model unit tests had already passed. - `skills/olive/references/quantization-onboarding.md` — general RTN/GPTQ pass onboarding: shared config surface, when to use RTN vs. GPTQ, calibration split hygiene. - `skills/olive/references/moe-gptq.md` — MoE-GPTQ-specific onboarding: why MoE needs its own calibration path, the K-last layout allow-list, the dual fallback-threshold design (#2610), and what real-model benchmarking showed about fallback rates and quantization wall-time. - `skills/olive/references/profiling-benchmark-example.md` — worked example of running the benchmark script and interpreting its output. ### Three-model benchmark (bits=4, group_size=128, sym=true, full WikiText-2 `train` calibration, full `test` eval) | Model | Baseline PPL | RTN PPL (Δ, time) | GPTQ PPL (Δ, time) | KQuant PPL (Δ, time) | Fallback experts | | --- | --- | --- | --- | --- | --- | | granite-3.0-1b-a400m-base | 6.2877 | 7.5861 (+1.2984, 8.0s) | 6.9560 (+0.6683, 658.7s) | 7.5162 (+1.2286, 12.1s) | 2/768 (0.3%) | | OLMoE-1B-7B-0924 | 6.6182 | 7.1091 (+0.4909, 52.6s) | 6.8966 (+0.2784, 1499.3s) | 7.0507 (+0.4325, 71.8s) | 10/1024 (1.0%) | | Qwen1.5-MoE-A2.7B | 6.4246 | 6.9251 (+0.5005, 85.8s) | 6.6117 (+0.1872, 2475.6s) | 6.9318 (+0.5072, 148.2s) | 0/1440 (0.0%) | GPTQ consistently beats RTN on perplexity delta across all three models, at a real (but model-size/expert-count-correlated, not cleanly separable) wall-time cost. See `moe-gptq.md` for the full discussion, including the OLMoE layer-2/expert-5 case that empirically validates the dual fallback-threshold design from #2610. KQuant (#2618) numbers added for comparison: KQuant is data-free (no calibration set, no per-expert fallback concept — the "Fallback experts" column doesn't apply to it) and its quantization time is close to RTN's (both are cheap, uncalibrated passes), but its perplexity delta tracks RTN's rather than GPTQ's on all three models. All three KQuant runs used `moe=true` and forced `experts_implementation="eager"` at inference (`grouped_mm` cannot run against `QuantTensor`-wrapped experts; see #2619). ### Notes - This PR depends on `b1-gptq-moe` (#2610): `capture_moe_fallback_counts()` in the script unconditionally imports `olive.passes.pytorch.moe_calib`, which only exists on that branch. Please review/merge #2610 first. - Went through a full internal review pass (readability/correctness/adversarial/spec-adherence/ cross-module) before opening; findings incorporated include: fixing pass-name resolution to use the actual pass registry (`OlivePackageConfig.import_pass_module`) instead of guessing module paths, several docstring/arithmetic corrections in the reference docs, and hedging a couple of causal claims that the 3-data-point benchmark can't fully support. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe your changes
Extends the PyTorch
KQuantpass to support quantizing fused MoE expertweights, mirroring the layout-safety approach already applied to RTN in #2616:
kquant_find_qparamsto N-D tensors so fused expert weights ofshape
(E, OUT, K)can be quantized directly.allow_moe/moeconfig flag, gated behind the sharedcheck_moe_layout_supportguard frommoe_support.pyso quantization failsclosed on transposed or unverifiable expert layouts instead of silently
producing wrong results.
_iter_quant_info_params(was silentlyskipping non-
weight-named MoE params before).config.moerequestrather than the merged
qcfg.moe(same bug independently found by theCopilot automated reviewer on Guard RTN MoE quantization by expert layout #2616 and fixed there; KQuant had copied the
same buggy pattern).
Based on
moe-layout-guard-fix2(#2616) since this only depends onmoe_support.py, not on any GPTQ-specific work in #2610/#2612.Real-model perplexity numbers for this pass (granite-3.0-1b-a400m-base,
OLMoE-1B-7B-0924, Qwen1.5-MoE-A2.7B) are in the "KQuant PPL (Δ, time)" column
of the three-model benchmark table in #2612's PR description, alongside the
existing RTN/GPTQ results for the same models.
Checklist before requesting a review
lintrunner -a(Optional) Issue link