Add MoE GPTQ benchmark script and quantization onboarding docs - #2612
Merged
Ti-Tai Wang (titaiwangms) merged 4 commits intoAug 14, 2026
Merged
Conversation
Ti-Tai Wang (titaiwangms)
marked this pull request as ready for review
August 7, 2026 17:54
Ti-Tai Wang (titaiwangms)
requested review from
Jambay Kinley (jambayk) and
Xiaoyu (xiaoyu-work)
August 7, 2026 17:56
Contributor
There was a problem hiding this comment.
Pull request overview
Adds developer-facing onboarding references for PyTorch/Hugging Face weight quantization (RTN/GPTQ, incl. MoE) and introduces a standalone script to manually validate real-checkpoint perplexity/size deltas before vs. after running an Olive quantization pass.
Changes:
- Add
scripts/quantize_and_compare_perplexity.pyto quantize a real HF model with a chosen Olive PyTorch quantization pass and report WikiText-2 perplexity + size deltas (plus MoE fallback coverage when applicable). - Add three new reference docs under
skills/olive/references/covering quantization onboarding, MoE GPTQ specifics, and an end-to-end benchmark walkthrough. - Link the new quantization “deep dive” docs from
skills/olive/SKILL.md.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| skills/olive/SKILL.md | Adds pointers to deeper quantization reference docs for contributors. |
| skills/olive/references/quantization-onboarding.md | General RTN/GPTQ onboarding + shared config surface + where to look in code/tests. |
| skills/olive/references/moe-gptq.md | MoE GPTQ calibration/fallback design explanation and benchmark-backed guidance. |
| skills/olive/references/profiling-benchmark-example.md | Worked example of running and interpreting the benchmark script output. |
| scripts/quantize_and_compare_perplexity.py | New manual validation script for pass-agnostic, real-model perplexity/size comparisons. |
Suppressed comments (3)
scripts/quantize_and_compare_perplexity.py:355
- Baseline tokenizer/model loading should pass
trust_remote_codethrough (matching--trust_remote_codeand the HfModelHandler load). Without this, the script can fail early on models that Olive could otherwise quantize, and it can also make baseline vs quantized behavior diverge.
tokenizer = AutoTokenizer.from_pretrained(args.model_id)
baseline = AutoModelForCausalLM.from_pretrained(args.model_id, **dtype_kwargs).to(args.device)
if hasattr(baseline, "set_experts_implementation"):
scripts/quantize_and_compare_perplexity.py:369
HfModelHandlershould be constructed with the sametrust_remote_codesetting as the baseline load. Otherwise, baseline might load successfully (or fail) under one policy while the quantization input model uses another, which defeats the script's stated fairness guarantees.
input_model = HfModelHandler(model_path=args.model_id, load_kwargs=dtype_kwargs)
scripts/quantize_and_compare_perplexity.py:127
- The return type annotation uses
callable, which is a built-in function, not a typing construct. This makes the annotation misleading and can confuse type-checkers/readers. Prefer either a properCallable[...]annotation (with an import) or drop the return annotation here since this is a CLI helper.
def capture_calibration_dataset() -> tuple[dict, callable]:
Ti-Tai Wang (titaiwangms)
force-pushed
the
feat/moe-gptq-benchmark-and-onboarding
branch
from
August 7, 2026 23:05
38fdfb8 to
71e586c
Compare
5 tasks
Ti-Tai Wang (titaiwangms)
added a commit
that referenced
this pull request
Aug 13, 2026
## Describe your changes Extends the PyTorch `KQuant` pass to support quantizing fused MoE expert weights, mirroring the layout-safety approach already applied to RTN in #2616: - Generalizes `kquant_find_qparams` to N-D tensors so fused expert weights of shape `(E, OUT, K)` can be quantized directly. - Adds an `allow_moe`/`moe` config flag, gated behind the shared `check_moe_layout_support` guard from `moe_support.py` so quantization fails closed on transposed or unverifiable expert layouts instead of silently producing wrong results. - Fixes the discovery loop to use `_iter_quant_info_params` (was silently skipping non-`weight`-named MoE params before). - Fixes the MoE gate to key off this invocation's own `config.moe` request rather than the merged `qcfg.moe` (same bug independently found by the Copilot automated reviewer on #2616 and fixed there; KQuant had copied the same buggy pattern). Based on `moe-layout-guard-fix2` (#2616) since this only depends on `moe_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 - [x] Add unit tests for this change. - [x] Make sure all tests can pass. - [ ] Update documents if necessary. - [x] 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. ## (Optional) Issue link --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
Ti-Tai Wang (titaiwangms)
force-pushed
the
feat/moe-gptq-benchmark-and-onboarding
branch
from
August 14, 2026 18:34
df231cb to
2664cdc
Compare
Ti-Tai Wang (titaiwangms)
force-pushed
the
feat/moe-gptq-benchmark-and-onboarding
branch
from
August 14, 2026 18:58
c1af463 to
bb4beaa
Compare
Add scripts/quantize_and_compare_perplexity.py, a local validation tool that quantizes a real HF model with a given Olive pass and reports the perplexity regression, quantization wall-time, model size, and (for MoE calibration) per-expert fallback coverage against a baseline. Add three skill reference docs under skills/olive/references/: - quantization-onboarding.md: overview of Olive's PyTorch quantization passes (RTN, GPTQ, and related), shared config surface, RTN vs. GPTQ trade-offs, and calibration/eval split hygiene notes. - moe-gptq.md: MoE-specific GPTQ calibration mechanics (per-expert Hessians, K-last layout requirement and architecture allow-list), the dual fallback-threshold design, and empirical findings from a three-model benchmark (fallback rate vs. expert count, quantization time vs. calibration set size). - profiling-benchmark-example.md: worked example of using the benchmark script, including a three-model (granite/OLMoE/Qwen1.5-MoE) results table. Link the new references from SKILL.md.
…uracy - scripts/quantize_and_compare_perplexity.py: resolve pass classes via OlivePackageConfig.import_pass_module instead of guessing module names from lowercased class names (broke for AutoAWQQuantizer/GptqQuantizer); remove unused dir_size_gb(); fix stale docstrings; add --num_samples, --pass_config, and seq_len validation; make calibration token counting batch-size robust; add --max_len override. - skills/olive/references/quantization-onboarding.md: fix pass-name table to match olive_config.json's real registry, fix RTN-timing and GPTQ/RTN-ratio claims to match cited data, fix embeds support note, cross-link how-to-add-optimization-pass.md. - skills/olive/references/moe-gptq.md: correct the design-doc characterization (sufficiency was the final gate, not skew), fix solve-count arithmetic, hedge causal scaling and 'never worse than RTN' claims, correct the OLMoE 'looks fair' framing and OR-gate validation claim, fix LayerCoverage method names, note per-(expert,parameter) fallback granularity. - skills/olive/references/profiling-benchmark-example.md: fix KQuant capitalization. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: e10674c1-6909-4b09-9f5a-d41b28c89d2d
#2610's full-team review raised the default sufficiency multiplier from k=1 (N=K) to k=2 (N=2K) after measuring the real RTN-vs-GPTQ crossover sits closer to 1.5x-2x K. Reran all three benchmark models (granite-3.0-1b-a400m-base, OLMoE-1B-7B-0924, Qwen1.5-MoE-A2.7B) with identical methodology under the new default and updated: - moe-gptq.md: default value, OLMoE empirical example recomputed at k=1 (historical) and cross-referenced at k=2, fallback-rate table, wall-time section. - profiling-benchmark-example.md: main results table now reflects k=2, plus a new k=1 vs k=2 side-by-side comparison table and analysis of why the higher fallback rate did not measurably hurt perplexity or wall-time on these three models. Numbers labeled explicitly as k=1/k=2 (not "old/new") throughout, per review convention, since both configurations remain independently reproducible via --pass_config.
CodeQL flagged 'pass_config' as potentially used before initialization at the isinstance() check, since it cannot statically prove parser.error() (which calls sys.exit()) never returns. Initialize pass_config to an empty dict before the try block so the variable is always bound regardless of the flagged control-flow path.
Ti-Tai Wang (titaiwangms)
force-pushed
the
feat/moe-gptq-benchmark-and-onboarding
branch
from
August 14, 2026 20:40
bb4beaa to
14e42f4
Compare
Jambay Kinley (jambayk)
approved these changes
Aug 14, 2026
Ti-Tai Wang (titaiwangms)
deleted the
feat/moe-gptq-benchmark-and-onboarding
branch
August 14, 2026 21:36
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
Stacked on top of #2610 (this PR targets
b1-gptq-moe, notmain).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/forquantization work in this repo:
scripts/quantize_and_compare_perplexity.py— generic, pass-agnostic script (works for anyregistered 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 Extend PyTorch RTN weight quantization to MoE experts #2584 after synthetic-model unit tests
had already passed.
skills/olive/references/quantization-onboarding.md— general RTN/GPTQ pass onboarding: sharedconfig surface, when to use RTN vs. GPTQ, calibration split hygiene.
skills/olive/references/moe-gptq.md— MoE-GPTQ-specific onboarding: why MoE needs its owncalibration path, the K-last layout allow-list, the dual fallback-threshold design (Add GPTQ quantization support for K-last MoE architectures #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 thebenchmark script and interpreting its output.
Three-model benchmark (bits=4, group_size=128, sym=true, full WikiText-2
traincalibration, fulltesteval)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.mdforthe 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=trueand forcedexperts_implementation="eager"at inference (grouped_mmcannot runagainst
QuantTensor-wrapped experts; see #2619).Notes
b1-gptq-moe(Add GPTQ quantization support for K-last MoE architectures #2610):capture_moe_fallback_counts()in the scriptunconditionally imports
olive.passes.pytorch.moe_calib, which only exists on that branch.Please review/merge Add GPTQ quantization support for K-last MoE architectures #2610 first.
cross-module) before opening; findings incorporated include: fixing pass-name resolution to use
the actual pass registry (
OlivePackageConfig.import_pass_module) instead of guessing modulepaths, 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.