Skip to content

Add MoE GPTQ benchmark script and quantization onboarding docs - #2612

Merged
Ti-Tai Wang (titaiwangms) merged 4 commits into
mainfrom
feat/moe-gptq-benchmark-and-onboarding
Aug 14, 2026
Merged

Add MoE GPTQ benchmark script and quantization onboarding docs#2612
Ti-Tai Wang (titaiwangms) merged 4 commits into
mainfrom
feat/moe-gptq-benchmark-and-onboarding

Conversation

@titaiwangms

@titaiwangms Ti-Tai Wang (titaiwangms) commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

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 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: 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 (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 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 (Add GPTQ quantization support for K-last MoE architectures #2610): capture_moe_fallback_counts() in the script
    unconditionally 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.
  • 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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.py to 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_code through (matching --trust_remote_code and 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

  • HfModelHandler should be constructed with the same trust_remote_code setting 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 proper Callable[...] annotation (with an import) or drop the return annotation here since this is a CLI helper.
def capture_calibration_dataset() -> tuple[dict, callable]:

Comment thread scripts/quantize_and_compare_perplexity.py
Comment thread scripts/quantize_and_compare_perplexity.py
Comment thread scripts/quantize_and_compare_perplexity.py
@titaiwangms
Ti-Tai Wang (titaiwangms) force-pushed the feat/moe-gptq-benchmark-and-onboarding branch from 38fdfb8 to 71e586c Compare August 7, 2026 23:05
Comment thread scripts/quantize_and_compare_perplexity.py Fixed
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
@titaiwangms
Ti-Tai Wang (titaiwangms) force-pushed the feat/moe-gptq-benchmark-and-onboarding branch from df231cb to 2664cdc Compare August 14, 2026 18:34
@titaiwangms
Ti-Tai Wang (titaiwangms) force-pushed the feat/moe-gptq-benchmark-and-onboarding branch from c1af463 to bb4beaa Compare August 14, 2026 18:58
Base automatically changed from b1-gptq-moe to main August 14, 2026 20:40
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.
@titaiwangms
Ti-Tai Wang (titaiwangms) force-pushed the feat/moe-gptq-benchmark-and-onboarding branch from bb4beaa to 14e42f4 Compare August 14, 2026 20:40
@titaiwangms
Ti-Tai Wang (titaiwangms) merged commit 0cbcdd9 into main Aug 14, 2026
10 checks passed
@titaiwangms
Ti-Tai Wang (titaiwangms) deleted the feat/moe-gptq-benchmark-and-onboarding branch August 14, 2026 21:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants