diff --git a/scripts/quantize_and_compare_perplexity.py b/scripts/quantize_and_compare_perplexity.py new file mode 100644 index 0000000000..79b3d46b7b --- /dev/null +++ b/scripts/quantize_and_compare_perplexity.py @@ -0,0 +1,492 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Manual verification helper: compare wikitext perplexity before/after quantizing a real HF model. + +This is a local, one-off validation tool -- it is NOT wired into any Olive workflow config or +CI. It exists to answer "does quantizing this real (downloaded) checkpoint produce a sane +quality regression, or does it blow up / silently corrupt the model?" -- the same class of +question that surfaced real bugs in #2584 (RTN) after synthetic-model unit tests had already +passed. + +Usage: + python scripts/quantize_and_compare_perplexity.py \ + --model_id ibm-granite/granite-3.0-1b-a400m-base \ + --pass_name Gptq \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' + +By default this evaluates perplexity over the *entire* wikitext-2 test split (use +``--num_samples`` to restrict to a prefix for a quicker/dirtier check). The model can be any HF +model id or local path; the pass can be any registered Olive pytorch *quantization* pass (Gptq, +Rtn, KQuant, AutoAWQQuantizer, GptqQuantizer, GptqModel, ...) with any config -- nothing here is +hardcoded to GPTQ or to MoE. Non-quantization passes (e.g. SparseGPT, which prunes rather than +quantizes) can technically be loaded too, but the size/perplexity comparison this script prints is +framed around quantization and may be misleading for a pruning pass. + +Fairness notes (read before trusting a delta from this script): + * Baseline and quantized models are loaded with the *same* ``--dtype`` (default "auto", i.e. + each checkpoint's native dtype) and the same ``--experts_implementation``, so a measured + perplexity delta is attributable to quantization and not to an incidental dtype/backend + mismatch between the two loads. + * "Weights size" reports two *different* metrics, both labeled explicitly in the summary: an + in-memory figure (parameter count x element size at the loaded dtype), computed identically + for baseline and quantized so that pair is directly comparable, and a separately-labeled + on-disk figure (actual saved weight-file bytes, quantized side only) that reflects real + storage compression. Do not compare the in-memory baseline number against the on-disk + quantized number -- they measure different things. + * Calibration sample/token counts are captured by intercepting the *actual* dataset the pass + builds internally (not recomputed separately by this script), so they are structurally + guaranteed to match what was really used to calibrate -- and are skipped/reported as n/a for + data-free passes (e.g. RTN) that don't consume a data_config at all. + * "Quantization time" is wall-clock for the whole ``pass.run()`` call (load + calibrate + + quantize + save), not a pure inference/perf benchmark -- a data-free pass like RTN being much + faster than a calibration-based pass like GPTQ is an expected algorithmic trade-off, not a + regression. + * ``--device`` controls where the baseline model loads and where perplexity is evaluated for + both models. It does NOT control where quantization/calibration itself runs -- some passes + (e.g. Gptq's layerwise calibration) pick their own device internally (cuda when available, + else cpu) independent of this flag. On a single-GPU machine this is usually moot; on a + multi-GPU machine, quantization may run on a different GPU than the one named here. + * A single run against one (possibly non-random) slice of one dataset, on one machine/GPU, is a + smoke test, not a statistically rigorous benchmark -- treat deltas smaller than the + run-to-run/sample-to-sample noise floor with caution, especially for small ``--num_samples`` + overrides, and do not treat cited timing numbers as reproducible to more than roughly + plus-or-minus 20% run-to-run without re-measuring on your own hardware. + * Calibration data defaults to Olive's own built-in default (wikitext-2 train), matching what a + real user gets out of the box -- this is what the primary results table should be based on. + Wikitext-2 train and the (also wikitext-2, by default) eval split are disjoint but same-domain, + so as a *secondary*, opt-in cross-domain sanity check (not the headline number), evaluate on + C4 instead: ``--dataset allenai/c4 --dataset_config en --split validation --eval_streaming + --num_samples 200`` (C4 is sharded across 1000+ files, so non-streaming slicing is + prohibitively slow -- ``--eval_streaming`` is required for it). +""" + +import argparse +import gc +import json +import shutil +import statistics +import tempfile +import time +from pathlib import Path + +import torch + +# ruff: noqa: T201 # this is a CLI tool; print() output is the point + +_WEIGHT_FILE_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth") + + +def weights_size_gb(path: Path) -> float: + """On-disk size of only the model-weight files under ``path``, in GiB. + + Excludes config/tokenizer/vocab/metadata files so this is comparable across checkpoints that + may differ in how much non-weight metadata they carry. + """ + return sum(f.stat().st_size for f in path.rglob("*") if f.is_file() and f.suffix in _WEIGHT_FILE_SUFFIXES) / ( + 1024**3 + ) + + +def param_memory_gb(model: torch.nn.Module) -> float: + """In-memory footprint of a model's parameters at their current dtype, in GiB.""" + return sum(p.numel() * p.element_size() for p in model.parameters()) / (1024**3) + + +def capture_moe_fallback_counts() -> tuple[list, callable]: + """Monkeypatch ``CoverageReport.add`` to capture per-layer MoE fallback (RTN) stats. + + ``Gptq``'s MoE calibration only *logs* its coverage/fallback report -- it doesn't return it + to the caller. Since the fallback experts (too few calibration tokens for a useful Hessian) + are exactly the "which experts did NOT get real GPTQ treatment" signal this script wants, + intercept the structured ``LayerCoverage`` objects before they're formatted into a log + string, instead of parsing log text. + + Returns (captured_layers, restore_fn). Stays empty for non-MoE passes: the module-level patch + is applied unconditionally (import is cheap and side-effect-free until ``CoverageReport.add`` + is actually called), but data-free/non-MoE passes never call it, so ``captured`` stays empty. + """ + from olive.passes.pytorch.moe_calib import CoverageReport + + captured = [] + original_add = CoverageReport.add + + def patched_add(self, coverage): + captured.append(coverage) + return original_add(self, coverage) + + CoverageReport.add = patched_add + + def restore(): + CoverageReport.add = original_add + + return captured, restore + + +def capture_calibration_dataset() -> tuple[dict, callable]: + """Monkeypatch the calibration-dataset builder actually called by ``run_layerwise_quantization``. + + ``olive.passes.pytorch.quant_utils`` does ``from .train_utils import get_calibration_dataset``, + binding its own module-level name -- patching ``train_utils.get_calibration_dataset`` would + not affect that call site, so this patches ``quant_utils.get_calibration_dataset`` directly. + This captures the *exact* dataset the pass consumes (not a separately-recomputed copy), so + the reported sample/token counts are structurally guaranteed to match, and for data-free + passes (e.g. RTN) that never call this function, ``captured`` stays empty. + """ + import olive.passes.pytorch.quant_utils as quant_utils_mod + + captured: dict = {} + original = quant_utils_mod.get_calibration_dataset + + def patched(model, data_config): + dataset = original(model, data_config) + captured["dataset"] = dataset + return dataset + + quant_utils_mod.get_calibration_dataset = patched + + def restore(): + quant_utils_mod.get_calibration_dataset = original + + return captured, restore + + +def compute_perplexity( + model, tokenizer, text: str, device: str, stride: int = 512, max_len: int = 2048 +) -> tuple[float, int]: + """Sliding-window perplexity over ``text``, per the standard HF perplexity recipe. + + Uses overlapping windows (``max_len`` context, ``stride`` step) so every scored token has a + long enough context, without requiring the whole document to fit in one forward pass. + ``max_len`` is a fixed evaluation-window size (matching the common HF perplexity recipe, + which typically also uses a fixed window rather than each model's full context length) -- + it is applied identically to baseline and quantized so the comparison stays fair even though + it may be smaller than a given model's actual ``max_position_embeddings``. Override via + ``--max_len`` if you need a specific window size for your model. + + Returns ``(perplexity, n_scored_tokens)``. ``n_scored_tokens`` is the number of positions + that actually contributed to the loss -- not ``trg_len``, since a causal-LM's internal + label-shift means the first position of every window's target has no prediction to score. + """ + model.eval() + input_ids = tokenizer(text, return_tensors="pt").input_ids + seq_len = input_ids.size(1) + if seq_len < 2: + raise ValueError( + f"Eval text tokenized to only {seq_len} token(s) -- need at least 2 for a scoreable " + "causal-LM window. Check --dataset/--dataset_config/--split/--num_samples." + ) + + nlls = [] + n_tokens = 0 + prev_end = 0 + for begin in range(0, seq_len, stride): + end = min(begin + max_len, seq_len) + trg_len = end - prev_end + ids = input_ids[:, begin:end].to(device) + target_ids = ids.clone() + target_ids[:, :-trg_len] = -100 # ignore_index: don't double-count the overlapped prefix + with torch.no_grad(): + loss = model(ids, labels=target_ids).loss + # the model shifts labels internally (logits[:-1] vs labels[1:]), so one more position + # than "trg_len" is unscored; weight by what was actually scored, not the raw window size. + n_valid = (target_ids[:, 1:] != -100).sum().item() + nlls.append(loss * n_valid) + n_tokens += n_valid + prev_end = end + if end == seq_len: + break + return torch.exp(torch.stack(nlls).sum() / n_tokens).item(), n_tokens + + +def load_eval_text( + dataset: str, dataset_config: str, split: str, num_samples: int | None, streaming: bool = False +) -> str: + """Load and concatenate non-empty rows of ``split`` for perplexity evaluation. + + ``num_samples=None`` (the default) uses the *entire* split -- wikitext-2's test split has + many blank/heading-only rows, so a small fixed prefix (e.g. the first 200 rows) can end up + representing only a small, non-random, unrepresentative slice of the actual text. + + ``streaming=True`` is required for datasets sharded across many files (e.g. ``allenai/c4``, + used as an optional cross-domain eval set): a non-streaming row-slice like ``split[:50]`` + still triggers enumerating/downloading every shard's metadata to compute slice boundaries, + which is far slower than just streaming the first ``num_samples`` rows directly. Streaming + mode requires ``num_samples`` to be set (no notion of "the entire split" for an unbounded + stream). + """ + from datasets import load_dataset + + if streaming: + if num_samples is None: + raise ValueError("--num_samples is required when using a streaming eval dataset (e.g. C4).") + ds = load_dataset(dataset, dataset_config, split=split, streaming=True) + rows = [] + for row in ds: + if row["text"].strip(): + rows.append(row["text"]) + if len(rows) >= num_samples: + break + return "\n\n".join(rows) + + ds = load_dataset(dataset, dataset_config, split=split) + rows = [row for row in ds["text"] if row.strip()] + if num_samples is not None: + rows = rows[:num_samples] + return "\n\n".join(rows) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model_id", required=True, help="HF model id or local model path.") + parser.add_argument("--pass_name", default="Gptq", help="Olive pytorch pass class name, e.g. Gptq, Rtn, KQuant.") + parser.add_argument( + "--pass_config", + default='{"bits": 4, "group_size": 128, "sym": true}', + help="JSON dict of pass config kwargs (e.g. bits/group_size/sym/moe/...). Pass an explicit " + "'data_config' key here to override the calibration dataset (see --calib_* flags for a shortcut).", + ) + parser.add_argument("--dataset", default="wikitext", help="HF dataset id for the perplexity eval text.") + parser.add_argument("--dataset_config", default="wikitext-2-raw-v1", help="HF dataset config name.") + parser.add_argument("--split", default="test", help="Dataset split to use for the eval text.") + parser.add_argument( + "--num_samples", + type=int, + default=None, + help="Restrict the eval split to the first N (non-empty) rows. Default: use the entire split. " + "Required (must be set) when --eval_streaming is used.", + ) + parser.add_argument( + "--eval_streaming", + action="store_true", + help="Stream the eval dataset instead of a plain row-slice load. Required for datasets sharded across " + "many files, e.g. '--dataset allenai/c4 --dataset_config en --split validation --eval_streaming " + "--num_samples 200' to use C4 as a cross-domain eval set (checking whether same-domain " + "wikitext-calibration-on-wikitext-eval inflates quality vs. a genuinely held-out domain).", + ) + parser.add_argument( + "--dtype", + default="auto", + help="torch_dtype used to load BOTH the baseline and the (pre-quantization) input model, so a measured " + "perplexity delta is attributable to quantization and not to a dtype mismatch between the two loads. " + "'auto' (default) preserves each checkpoint's native dtype, matching Olive's own default load behavior.", + ) + parser.add_argument( + "--calib_dataset", + default=None, + help="HF dataset id for calibration data, e.g. 'allenai/c4' for a cross-domain calibration set (the " + "convention used by the GPTQ/AWQ papers) instead of Olive's wikitext-2 default, which shares a domain " + "with the default --dataset/--dataset_config eval text above (train/test are disjoint, but same-domain " + "calibration can still inflate quality on a same-domain eval set relative to true generalization). " + "Default: use Olive's built-in default (Salesforce/wikitext, wikitext-2-raw-v1). Ignored if --pass_config " + "already specifies 'data_config', or if the target pass doesn't declare a data_config parameter at all " + "(e.g. Rtn, which is data-free).", + ) + parser.add_argument("--calib_dataset_config", default=None, help="HF dataset config name for --calib_dataset.") + parser.add_argument( + "--calib_split", + default=None, + help="Dataset split to use for calibration data. Default: None, meaning don't override -- inherit " + "whatever Olive's own current default is (get_calibration_data_config's own 'split' default), so this " + "script always reflects real Olive behavior even if that default changes in the future, instead of " + "silently drifting out of sync with a value hardcoded here.", + ) + parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu") + parser.add_argument( + "--experts_implementation", + default="eager", + help="Set via model.set_experts_implementation(...) on BOTH the baseline and the quantized model before " + "evaluating a MoE model, so the two loads use the same MoE compute backend. Quantized MoE weights are " + "storage-only QuantTensors that only the eager experts loop can consume.", + ) + parser.add_argument( + "--keep_output", + action="store_true", + help="Keep the quantized model directory instead of deleting it when the script exits.", + ) + parser.add_argument( + "--max_len", + type=int, + default=2048, + help="Sliding-window context length used for perplexity evaluation (applied identically to " + "baseline and quantized). Default 2048, matching the common HF perplexity recipe -- override " + "if you specifically need a window matching a given model's max_position_embeddings.", + ) + args = parser.parse_args() + + if args.num_samples is not None and args.num_samples <= 0: + parser.error("--num_samples must be a positive integer.") + + pass_config = {} + try: + pass_config = json.loads(args.pass_config) + except json.JSONDecodeError as e: + parser.error(f"--pass_config is not valid JSON: {e}") + if not isinstance(pass_config, dict): + parser.error("--pass_config must be a JSON object (dict), e.g. '{\"bits\": 4}'.") + + from olive.hardware import DEFAULT_CPU_ACCELERATOR + from olive.model import HfModelHandler + from olive.package_config import OlivePackageConfig + from olive.passes.olive_pass import create_pass_from_dict + from olive.passes.pytorch.train_utils import get_calibration_data_config + + # Resolve the pass class through Olive's own package registry (olive_config.json) rather than + # guessing a module path from the class name: several registered passes (e.g. AutoAWQQuantizer + # in autoawq.py, GptqQuantizer in autogptq.py) do not live in a module named after the + # lowercased class name, so a naive `olive.passes.pytorch.{name.lower()}` import silently + # breaks for them. + pass_cls = OlivePackageConfig.load_default_config().import_pass_module(args.pass_name) + pass_accepts_data_config = "data_config" in pass_cls.default_config(DEFAULT_CPU_ACCELERATOR) + + print( + f"=== Loading eval text: {args.dataset}/{args.dataset_config} [{args.split}][:{args.num_samples}] " + f"(streaming={args.eval_streaming}) ===" + ) + text = load_eval_text(args.dataset, args.dataset_config, args.split, args.num_samples, args.eval_streaming) + + dtype_kwargs = {} if args.dtype == "none" else {"torch_dtype": args.dtype} + + print(f"=== Loading baseline model: {args.model_id} (dtype={args.dtype}) ===") + from transformers import AutoModelForCausalLM, AutoTokenizer + + 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"): + baseline.set_experts_implementation(args.experts_implementation) + baseline_size_gb = param_memory_gb(baseline) + baseline_dtype = next(baseline.parameters()).dtype + + print("=== Baseline perplexity ===") + baseline_ppl, baseline_n_tokens = compute_perplexity(baseline, tokenizer, text, args.device, max_len=args.max_len) + print(f"Baseline perplexity: {baseline_ppl:.4f} (scored {baseline_n_tokens} tokens)") + del baseline + gc.collect() + torch.cuda.empty_cache() + + # Load the pre-quantization input model with the SAME dtype kwargs as the baseline above, so the + # only difference between "baseline" and "quantized" is the quantization itself. + input_model = HfModelHandler(model_path=args.model_id, load_kwargs=dtype_kwargs) + + if pass_accepts_data_config and "data_config" not in pass_config: + calib_dataset_kwargs = {} + if args.calib_split: + calib_dataset_kwargs["split"] = args.calib_split + if args.calib_dataset: + calib_dataset_kwargs["data_name"] = args.calib_dataset + if args.calib_dataset_config: + calib_dataset_kwargs["subset"] = args.calib_dataset_config + data_config = get_calibration_data_config( + args.model_id, + trust_remote_code=input_model.get_load_kwargs().get("trust_remote_code", False), + **calib_dataset_kwargs, + ) + pass_config = {**pass_config, "data_config": data_config} + elif not pass_accepts_data_config and (args.calib_dataset or "data_config" in pass_config): + print( + f"NOTE: {args.pass_name} does not declare a 'data_config' parameter (data-free quantization); " + "--calib_dataset / pass_config['data_config'] will be ignored." + ) + + # Capture the calibration dataset actually built and consumed *inside* run_layerwise_quantization + # (rather than recomputing a separate copy here), so the reported sample/token counts are + # structurally guaranteed to match what was really used -- and stay empty for data-free passes. + captured_calib, restore_calib_capture = capture_calibration_dataset() + fallback_layers, restore_coverage_capture = capture_moe_fallback_counts() + + printable_pass_config = {k: v for k, v in pass_config.items() if k != "data_config"} + out_dir = Path(tempfile.mkdtemp(prefix="olive_quant_demo_")) + try: + print(f"=== Running {args.pass_name} pass with config={printable_pass_config} ===") + quant_pass = create_pass_from_dict(pass_cls, pass_config, disable_search=True) + quant_start = time.time() + output_model = quant_pass.run(input_model, str(out_dir)) + quant_duration_s = time.time() - quant_start + quantized_weights_size_gb = weights_size_gb(out_dir) + + if "dataset" in captured_calib: + calib_num_samples = len(captured_calib["dataset"]) + # Sum tokens across the full batch dimension of every row (not just row["input_ids"][0]) + # so this stays correct if a --calib_* override or a future default ever uses batch_size > 1. + calib_num_tokens = sum(row["input_ids"].numel() for row in captured_calib["dataset"]) + calib_summary = f"{calib_num_samples} samples, {calib_num_tokens} tokens" + else: + calib_summary = "n/a (data-free pass)" + + print("=== Loading quantized model ===") + quantized = output_model.load_model() + if hasattr(quantized, "set_experts_implementation"): + quantized.set_experts_implementation(args.experts_implementation) + quantized = quantized.to(args.device) + + print("=== Quantized perplexity ===") + quant_ppl, quant_n_tokens = compute_perplexity(quantized, tokenizer, text, args.device, max_len=args.max_len) + print(f"Quantized perplexity: {quant_ppl:.4f} (scored {quant_n_tokens} tokens)") + + # Apples-to-apples with baseline_size_gb: same measurement method (in-memory param + # bytes at the loaded dtype) on both sides. weights_size_gb(out_dir) above is a + # DIFFERENT metric (on-disk artifact size, includes any save-time packing/compression) + # and is reported separately -- conflating the two was Major finding #3 from review. + quantized_size_gb = param_memory_gb(quantized) + + total_experts = sum(lc.num_experts for lc in fallback_layers) + total_starved = sum(lc.starved for lc in fallback_layers) + total_unseen = sum(lc.unseen for lc in fallback_layers) + + print("\n=== SUMMARY ===") + print(f"Model: {args.model_id}") + print(f"Pass: {args.pass_name}({printable_pass_config})") + print(f"Dtype (baseline & input): {baseline_dtype} (--dtype={args.dtype})") + print(f"Experts implementation: {args.experts_implementation} (applied to both models)") + print(f"Calibration set: {calib_summary}") + print( + f"Quantization time: {quant_duration_s:.1f}s " + "(end-to-end: load + calibrate + quantize + save, NOT a pure inference benchmark)" + ) + print(f"Baseline weights size: {baseline_size_gb:.3f} GiB (in-memory params, {baseline_dtype})") + quantized_dtype = next(quantized.parameters()).dtype + print(f"Quantized weights size: {quantized_size_gb:.3f} GiB (in-memory params, {quantized_dtype})") + print( + f"Quantized on-disk size: {quantized_weights_size_gb:.3f} GiB (saved weight files only, different metric)" + ) + print( + f"Eval text: {args.dataset}/{args.dataset_config}[{args.split}]" + f"{f'[:{args.num_samples}]' if args.num_samples is not None else ' (full split)'}" + ) + print(f"Baseline perplexity: {baseline_ppl:.4f} ({baseline_n_tokens} tokens scored)") + print(f"Quantized perplexity: {quant_ppl:.4f} ({quant_n_tokens} tokens scored)") + print(f"Delta: {quant_ppl - baseline_ppl:+.4f}") + if fallback_layers: + fallback_label = ( + f"{args.pass_name} + RTN fallback" + if total_starved + total_unseen + else f"{args.pass_name} (no fallback)" + ) + print( + f"MoE fallback experts: {total_starved + total_unseen}/{total_experts} " + f"({total_starved} starved, {total_unseen} unseen) across {len(fallback_layers)} MoE layers " + f"-- quantized with RTN instead of {args.pass_name} due to insufficient calibration coverage " + f"[label: {fallback_label}]" + ) + print("MoE per-layer token/expert coverage:") + for lc in fallback_layers: + counts = sorted(lc.token_counts) + print( + f" {lc.layer_name}: {lc.num_experts} experts, tokens/expert " + f"min={counts[0]} median={statistics.median(counts):.0f} max={counts[-1]}, " + f"{lc.starved} starved, {lc.unseen} unseen" + ) + else: + print("MoE fallback experts: n/a (not a MoE calibration run)") + finally: + restore_coverage_capture() + restore_calib_capture() + if args.keep_output: + print(f"Quantized model kept at: {out_dir}") + else: + shutil.rmtree(out_dir, ignore_errors=True) + + +if __name__ == "__main__": + main() diff --git a/skills/olive/SKILL.md b/skills/olive/SKILL.md index f3659f1700..5b0aa8e803 100644 --- a/skills/olive/SKILL.md +++ b/skills/olive/SKILL.md @@ -197,6 +197,21 @@ Structural validation cannot prove that remote models are accessible, local data the model is supported by every pass, or the target hardware has enough memory. Surface those constraints instead of presenting validation as execution success. +## Deep dives for quantization work + +For PyTorch/Hugging Face weight quantization (RTN, GPTQ, and related passes), read these before +making nontrivial changes or answering detailed questions: + +- [Quantization onboarding](references/quantization-onboarding.md) -- pass overview, shared config + surface, RTN vs. GPTQ trade-offs, calibration data and split-hygiene notes. +- [MoE GPTQ onboarding](references/moe-gptq.md) -- per-expert calibration mechanics, the K-last + layout requirement and architecture allow-list, the dual fallback-threshold design (routing + skew vs. statistical sufficiency), and what a three-model benchmark showed about fallback rates + and quantization wall-time. +- [Profiling/benchmark example](references/profiling-benchmark-example.md) -- how to use + `scripts/quantize_and_compare_perplexity.py` to compare a pass against baseline on a real model, + with a worked three-model example table. + ## Dependency and hardware rules - Reuse the user's active environment when it already contains the required Olive and runtime packages. diff --git a/skills/olive/references/moe-gptq.md b/skills/olive/references/moe-gptq.md new file mode 100644 index 0000000000..174e780c50 --- /dev/null +++ b/skills/olive/references/moe-gptq.md @@ -0,0 +1,242 @@ +# MoE GPTQ Onboarding + +How Olive's `Gptq` pass (`olive/passes/pytorch/gptq.py`) quantizes Mixture-of-Experts (MoE) +models, why it needs a dedicated calibration path (`olive/passes/pytorch/moe_calib.py`), the +dual fallback-threshold design, and what the empirical validation from a three-model benchmark +showed. Read [`quantization-onboarding.md`](quantization-onboarding.md) first for the general +RTN/GPTQ context if you haven't already. For how to reproduce the benchmark numbers cited below, +see [`profiling-benchmark-example.md`](profiling-benchmark-example.md). + +## Why MoE needs its own calibration path + +Standard (dense-model) GPTQ calibration hooks a single `nn.Linear` per layer and accumulates one +Hessian from whatever activations flow through it. Fused MoE architectures don't have a per-expert +`nn.Linear` — routing happens *inside* a single "experts" module's forward call, so a plain +`register_forward_hook` on that module sees one undifferentiated activation batch with no way to +attribute individual rows to the expert that actually processed them. + +`moe_calib.py` solves this by hooking into transformers' own experts-implementation registry +(`ALL_EXPERTS_FUNCTIONS` / `@use_experts_implementation`, requires `transformers >= 5.0`): +Olive registers one generic recording implementation and points the model at it for the duration +of calibration. Every decorated experts class dispatches through it uniformly, so this file +contains no per-architecture branching for the recording logic itself — only an allow-list (see +below) of which architectures it has been verified safe for. + +## The K-last layout requirement + +GPTQ's `(K, K)` Hessian math assumes the weight's contraction dimension (`K`, the input-feature +dimension) is the last dimension of the fused weight tensor: `(num_experts, out_features, K)`. +Olive's MoE support is **allow-listed**, not best-effort, because getting this wrong silently +mis-quantizes the model rather than erroring: + +- `SUPPORTED_MOE_EXPERTS_CLASSES` in `moe_calib.py` lists the exact experts class name (not just + `model_type`) verified to (a) use `(num_experts, out, in)` — K-last — fused weight storage, and + (b) carry the `@use_experts_implementation` decorator. Currently verified: DeepSeek-V3, + GraniteMoE, Jamba, Mixtral, OLMoE, Phi-MoE, Qwen2-MoE, Qwen3-MoE (against transformers 5.14.1). +- The class name is checked *in addition to* `model_type` specifically so a spoofed or mismatched + `config.model_type` string cannot smuggle an unverified experts module past the allow-list. +- **Transposed-layout architectures** (`gpt_oss`, `llama4`, `aria`) store `(num_experts, in, out)` + — K is not last. These are refused with a clear error rather than silently mis-quantized. + Supporting them requires a layout-normalization step that is deliberately out of scope for the + current implementation — if you need this, that normalization (transpose before Hessian + accumulation, transpose back before save) is the right place to start, not extending the + allow-list to include them as-is. + +If you're adding support for a new MoE architecture: verify its fused weight layout and decorator +support against the currently-installed `transformers` version before adding it to +`SUPPORTED_MOE_EXPERTS_CLASSES` — do not assume a new architecture matches by analogy to an +existing allow-listed one. + +## Per-expert fallback: two independent conditions + +Not every expert in a calibration run receives enough routed tokens to compute a well-formed +Hessian. When an expert's calibration signal is inadequate, `Gptq` falls back to quantizing that +expert with RTN instead of GPTQ. The design doc's rationale (`gptq_moe_design.md`) is that this +should never be *worse* than plain RTN for that expert, since GPTQ's correction degenerates +toward the damping prior once the Hessian is under-determined — but note this is a working design +policy, not a formally proven guarantee: the design doc itself flags that a damped, rank-deficient +Hessian does not mathematically reduce to exactly RTN (damping reweights the null directions +rather than eliminating their influence, and the layerwise corrections stay coupled across +columns). Treat "never worse than RTN" as the intended, empirically-motivated behavior rather than +a theorem. + +An expert falls back if it fails **either** of two independent conditions (an OR-gate, +implemented as `LayerCoverage.threshold = max(skew_threshold, k_threshold)` in `moe_calib.py` +and mirrored as `N < skew_threshold or N < sufficiency_threshold` in `Gptq._process_moe_param` in +`gptq.py`): + +1. **Routing skew** (`moe_fallback_threshold`, default `0.005` = 0.5%, matching GPTQModel's + default): is this expert under-served *relative to its peers* in this same calibration run? + Computed as `fallback_threshold * tokens_seen_by_this_layer`. This condition is + **scale-invariant** — if you 10x the calibration set, this expert's absolute token count also + scales ~10x, so the skew ratio is roughly unchanged. It catches routing imbalance but not + under-calibration of the whole run. Note this threshold is only ever binding when + `fallback_threshold * tokens_seen < K` for the parameter being checked — at Olive's default + 262,144-token calibration budget, that means the skew condition can only fire for `K <~ 1311`; + for larger `K` values it is structurally dominated by the sufficiency condition below (see the + OLMoE example) and never actually gates anything on its own. +2. **Statistical sufficiency** (`moe_fallback_min_k_multiple`, default `2.0`): does this expert + have *enough absolute samples* to produce a well-formed `(K, K)` Hessian at all? An expert's + Hessian `H = sum(x xT)` accumulated from `N` routed tokens satisfies `rank(H) <= N`, so + `N < K` is *necessary* for `H` to be singular but not by itself sufficient to prove it (`N` + linearly-dependent samples at `N >= K` can still leave `H` rank-deficient) — practically, + `N < K` (`k=1`) is only the bare-minimum-rank floor, not the point at which the Hessian is + actually well-conditioned. A full-team review of this pass measured the real RTN-vs-GPTQ MSE + crossover under anisotropic activations at 4-bit to sit closer to `N` in the `1.5x-2x K` range: + at the `k=1` floor, GPTQ was measured ~4-6% *worse* than RTN on average for experts in that + band. `moe_fallback_min_k_multiple` was therefore raised from the original `k=1` default to + `k=2` (`N = 2K`) as a conservative, empirically-motivated policy choice past that crossover — + not a proven sufficient threshold. No external MoE-quantization literature search turned up a + specific "safe multiple of K" recommendation, only qualitative guidance; `k=2` is the repo's + own measured choice. This condition is **absolute**: more calibration data always helps it + directly, unlike the routing-skew condition. + +Note also that fallback is decided **per (expert, parameter)**, not per expert as a whole: a fused +MoE layer typically has two quantizable parameters (`gate_up_proj`, `down_proj`), and they can +have different `K` (input-feature dimension), so the *same* expert can pass the sufficiency check +for one parameter and fail it for the other. "N/M fallback experts" in a coverage summary counts +starved-or-unseen occurrences for the parameter(s) actually tracked in the coverage report, not +necessarily every quantizable parameter for that expert — check the per-layer log for the +specific parameter dimension if you need to know exactly which weights fell back. + +### Why both conditions, and why this was a design-vs-implementation gap + +The original design doc (`gptq_moe_design.md`) settled on K-multiple sufficiency as the *sole* +fallback trigger, treating the routing-skew percentage as diagnostic/reporting-only ("skew = +diagnostic (report); sufficiency = trigger (fallback)"). The first shipped implementation, +however, only implemented the skew-percentage condition (`DEFAULT_MOE_FALLBACK_THRESHOLD = +0.005`) from its very first commit — the design's settled K-multiple decision was never actually +implemented. This was not a later regression; it was a design decision that was reconciled only +during later benchmark validation work, which added `moe_fallback_min_k_multiple` and made +*both* conditions gate the fallback via an OR-gate — a stricter combination than the design doc's +own final call, which gated on sufficiency alone and used skew only for reporting. + +**Empirical evidence the two conditions are not redundant** (OLMoE-1B-7B-0924, full WikiText-2 +`train` calibration, 262,144 calibration tokens reaching the layer, `k=1` i.e. +`moe_fallback_min_k_multiple=1.0`, the value in effect when this specific case was first found): +layer 2, expert 5 received `N=1331` tokens for its `gate_up_proj` parameter (K=2048). +- Skew threshold: `0.005 * 262,144 = 1310.7` — `1331 > 1310.7`, so the **skew-only** check would + **not** flag this expert. This is not because the expert is genuinely well-served: OLMoE routes + top-8-of-64, so a "fair share" of tokens for one expert is `262,144 * 8/64 = 32,768` — the + 1,331 tokens this expert saw is only ~4% of fair share (~24x under-routed). The skew threshold + simply doesn't bind at this `K`, for the structural reason noted above (0.5% of 262,144 = + 1310.7 < K=2048), not because the expert looks reasonably calibrated. +- Sufficiency threshold at `k=1`: `1.0 * K = 1.0 * 2048 = 2048.0` — `1331 < 2048.0`, so the + **sufficiency** check **does** flag it: this expert's Hessian is necessarily rank-deficient + (`N < K`). + +This confirms the two conditions are not redundant in general (there exist real experts the +sufficiency check catches that skew alone would not), though this specific benchmark only produced +a case in the "sufficiency fires, skew doesn't" direction — it does not exercise a case where skew +would fire but sufficiency wouldn't, so it validates the OR-gate's usefulness over +skew-alone, but does not by itself demonstrate a case where dropping the sufficiency-alone design +(skew as report-only) would have been wrong in the opposite direction. + +With the default later raised to `k=2` (`moe_fallback_min_k_multiple=2.0`, see the note above), +this same OLMoE layer/expert is caught even more clearly: the sufficiency threshold becomes +`2.0 * 2048 = 4096.0`, so `1331 < 4096.0` by a wider margin, and the same rerun at `k=2` also +newly catches several additional experts across other layers whose `N` fell in the `1x-2x K` +"no man's land" (e.g. layer 8 expert 18/48/55 with `N` in the 2445-3772 range against +`K=2048` — these pass at `k=1` but fail at `k=2`). See "What the three-model benchmark showed" +below for the aggregate before/after `k=1` vs `k=2` comparison. + +## Coverage logging and diagnostics + +Every MoE layer's calibration coverage is logged via `CoverageReport.log_summary()` (built from +per-layer `LayerCoverage.format()` strings; see `moe_calib.py`), which reports, per layer: number +of experts covered/starved/unseen, the combined effective threshold (`max(skew, sufficiency)`), +and the min/median/max observed token count per expert. Read this log when investigating an +unexpectedly high fallback count for a given model — it will tell you whether the fallback is +driven by routing imbalance across experts or raw insufficiency (the whole calibration set too +small for this model's K), and at which specific `K` (i.e. which parameter) the threshold bound. + +## What the three-model benchmark showed about fallback rates + +Fallback rate does **not** scale with the total number of experts in a model. From a benchmark +across three MoE models (bits=4, group_size=128, sym=true, full WikiText-2 `train` calibration, +`k=2` i.e. `moe_fallback_min_k_multiple=2.0` — see `profiling-benchmark-example.md` for the full +table and reproduction steps): + +| Model | Experts/layer x layers | Total experts | Fallback experts (k=2) | +| --- | --- | --- | --- | +| granite-3.0-1b-a400m-base | 32 x 24 | 768 | 3 (0.4%) | +| OLMoE-1B-7B-0924 | 64 x 16 | 1024 | 21 (2.1%) | +| Qwen1.5-MoE-A2.7B | 60 x 24 | 1440 | 0 (0.0%) | + +Qwen1.5-MoE has the *most* total experts (1440, driven by having more layers, not more +experts-per-layer than OLMoE) but *zero* fallbacks, while OLMoE has fewer total experts but the +*highest* fallback rate. Since calibration tokens are split **per layer** among that layer's +experts, the relevant comparison is experts-per-layer (~60-64, comparable across both models), +not the total. The proximate cause is genuine routing imbalance for the affected OLMoE experts +(e.g. the layer-2-expert-5 case above, ~24x under fair share), but note the *fallback condition +that actually fires* for these K=2048 cases is sufficiency, not skew (skew is structurally inert +at this K, per the note above) — "routing skew" describes the underlying router behavior, while +"sufficiency" is the specific check that catches it in the current implementation. Qwen1.5-MoE's +per-layer logs showed min tokens/expert consistently well above the `k=2` sufficiency threshold +across all 24 layers, with no layer ever showing a significantly under-routed expert. + +Raising the sufficiency multiplier from `k=1` to `k=2` roughly doubled the fallback count for +granite (2->3) and OLMoE (10->21) as expected — more experts in the `1x-2x K` "no man's land" are +now caught — while Qwen1.5-MoE's fallback count stayed at 0 in both runs (its router is well +enough balanced on this calibration set that no expert falls below even the `k=2` threshold). +Despite the higher fallback count, **GPTQ perplexity did not get measurably worse at `k=2` vs +`k=1`** (see `profiling-benchmark-example.md`'s `k=1` vs `k=2` comparison table) — consistent with +the `1x-2x K` band being a region where GPTQ wasn't reliably beating RTN in the first place, so +routing those experts to RTN removes downside risk at effectively no aggregate cost. + +**Takeaway**: read fallback count as a diagnostic of *that specific model's router balance on +that specific calibration domain*, not as a function of how many experts the model has. Two +models with similar experts-per-layer counts can have very different fallback rates purely +because their trained routers distribute tokens differently (this is plausibly related to +architecture differences like shared-expert-plus-top-k gating and the load-balancing auxiliary +loss used at pretraining time, but this benchmark did not directly inspect router logits to +confirm that hypothesis). + +## Quantization wall-time: calibration-set size matters less than you'd expect + +A common intuition is "bigger calibration set → proportionally longer GPTQ quantization." For +MoE models this is **not** the dominant effect. GPTQ per-layer cost splits into two phases with +very different scaling behavior: + +1. **Hessian accumulation** (forward passes over calibration data) — scales ~linearly with the + number of calibration tokens. +2. **Per-expert GPTQ solve** (Cholesky decomposition + blockwise quantization on each expert's + `(K, K)` Hessian) — a **fixed cost** independent of how many tokens built that Hessian; it + depends only on `K` and the number of (expert, parameter) pairs to solve. + +On granite-3.0-1b-a400m-base, increasing the calibration set from a stale `train[:1000]` slice +(~61,816 tokens, a single unreplicated prior run, not independently re-verified) to the full +`train` split (262,144 tokens, ~4.24x more) only increased quantization wall-time by ~18% +(658.7s at `k=1` vs. ~558s in that single comparison) — consistent with phase 2 (1,536 +fixed-cost Cholesky/quantize solves for this model: 24 layers x 32 experts x 2 params) dominating +total time for a model with a modest active-parameter count and many experts. This ratio is +model-specific and based on a single before/after comparison, not repeated runs — for a model +with a much larger active-parameter count (slower forward pass) or fewer experts (solve phase +less dominant), the balance could shift back toward calibration-size-linear scaling. Treat the +two-phase *mechanism* as the durable takeaway; treat the specific "~18% for ~4.24x more data" +ratio as one illustrative data point, not a general formula. + +The three-model table's wall-times at `k=1` (658.7s / 1499.3s / 2475.6s for 1,536 / 2,048 / 2,880 +total solves) are directionally consistent with solve-count-driven scaling, but with only three +data points where solve count, `K`, and total parameter volume all increase together, this +benchmark cannot cleanly separate "cost driven by number of solves" from "cost driven by total +active parameter volume" as the dominant covariate — both plausibly contribute. Don't treat +`num_layers x num_experts` as a validated predictive formula from this data alone; treat it as the +mechanistically-motivated hypothesis the two-phase model above suggests. Rerunning the same three +models at `k=2` (653.4s / 1484.9s / 2588.9s) showed wall-time essentially unchanged from `k=1` +despite the higher fallback count — consistent with the fixed-cost-per-solve model, since raising +the sufficiency multiplier slightly *reduces* the number of GPTQ solves actually performed (more +experts skip straight to RTN), so total time trends flat-to-down rather than up. + +## Before touching MoE GPTQ code + +1. Read `moe_calib.py`'s module docstring and `SUPPORTED_MOE_EXPERTS_CLASSES` first — the + allow-list is a deliberate fail-closed safety mechanism, not an oversight to work around. +2. If you're changing fallback-threshold behavior, keep both conditions in mind — they measure + different things (relative skew vs. absolute sufficiency) and are not interchangeable. +3. Check `test/passes/pytorch/test_gptq.py` and any MoE-specific test file for the existing test + conventions (parametrization over architectures, fixture patterns) before adding new tests. +4. Validate any threshold or calibration change against a real model, not just synthetic/tiny + test fixtures — the dual-condition disagreement case above only showed up on a real model + (OLMoE) with real routing behavior; tiny random-weight test models are unlikely to reproduce + realistic routing skew. diff --git a/skills/olive/references/profiling-benchmark-example.md b/skills/olive/references/profiling-benchmark-example.md new file mode 100644 index 0000000000..0d8d20bf69 --- /dev/null +++ b/skills/olive/references/profiling-benchmark-example.md @@ -0,0 +1,161 @@ +# Profiling / Benchmark Example: `quantize_and_compare_perplexity.py` + +A worked example for `scripts/quantize_and_compare_perplexity.py`, a local validation script that +quantizes a real Hugging Face model with a given Olive pass and reports the perplexity +regression, quantization wall-time, model size, and (for MoE calibration) per-expert fallback +coverage. Read [`quantization-onboarding.md`](quantization-onboarding.md) for background on the +passes being compared and [`moe-gptq.md`](moe-gptq.md) for the MoE-specific fallback mechanism +whose coverage output this script surfaces. + +## What the script is (and isn't) + +This is a **local, one-off validation tool** — it is not wired into any Olive workflow config or +CI. It exists to answer "does quantizing this real (downloaded) checkpoint produce a sane quality +regression, or does it blow up / silently corrupt the model?", the same class of question that +has previously surfaced real bugs after synthetic-model unit tests had already passed. It is not +a substitute for unit tests, and a single run is a smoke test, not a statistically rigorous +benchmark — see the "Fairness notes" in the script's own module docstring before trusting a +delta from it. + +The script is generic: `--pass_name` can be any Olive PyTorch quantization pass (`Gptq`, `Rtn`, +`KQuant`, ...) with any `--pass_config`, and `--model_id` can be any HF model id or local path — +nothing is hardcoded to GPTQ or to MoE. + +## Basic usage + +```shell +python scripts/quantize_and_compare_perplexity.py \ + --model_id ibm-granite/granite-3.0-1b-a400m-base \ + --pass_name Gptq \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' \ + --device cuda:0 +``` + +By default this evaluates perplexity over the **entire** WikiText-2 `test` split and calibrates +GPTQ on Olive's default WikiText-2 `train` split (use `--num_samples` to restrict the eval split +to a prefix for a quicker/dirtier check; see `--help` for calibration-side overrides such as +`--calib_dataset`). + +Key flags: + +| Flag | Purpose | +| --- | --- | +| `--model_id` | HF model id or local path (required). | +| `--pass_name` | Olive pass class name, e.g. `Gptq`, `Rtn` (default `Gptq`). | +| `--pass_config` | JSON dict of pass config overrides. | +| `--dataset` / `--dataset_config` / `--split` | Eval dataset (default WikiText-2 `test`). | +| `--calib_dataset` / `--calib_dataset_config` | Calibration dataset override for calibrated passes. | +| `--dtype` | Load dtype for both baseline and quantized models (default `auto`, i.e. each checkpoint's native dtype). | +| `--experts_implementation` | Forces a specific MoE experts implementation on both loads for a fair comparison. | +| `--device` | Device for both quantization and eval. | + +## Fairness guarantees baked into the script (read before trusting a delta) + +- Baseline and quantized models are loaded with the *same* `--dtype` and + `--experts_implementation`, so a measured perplexity delta is attributable to quantization, not + to an incidental dtype/backend mismatch between the two loads. +- Reported "weights size" includes **both** an in-memory (parameter-count x dtype-size) figure — + computed identically for baseline and quantized so the two are apples-to-apples — and a + separately labeled on-disk size (actual saved weight-file bytes) for the quantized model. Do + not compare the in-memory baseline figure against the on-disk quantized figure; they are + different metrics reported side-by-side, not the same metric. +- Calibration sample/token counts are captured by intercepting the actual dataset the pass builds + internally, not recomputed separately by the script, so they are guaranteed to match what was + really used to calibrate — and are reported as `n/a` for data-free passes (e.g. RTN). +- "Quantization time" is wall-clock for the entire pass run (load + calibrate + quantize + save), + not a pure inference benchmark — RTN being much faster than GPTQ is an expected algorithmic + trade-off, not a regression. + +## Worked example: three-way MoE model comparison + +The following table was produced by running the script once per (model, pass) pair — six runs +total — with `bits=4, group_size=128, sym=true, moe=true`, full WikiText-2 `train` calibration +(Olive's default post-#2609), and the full WikiText-2 `test` split for eval: + +```shell +# Baseline + RTN +python scripts/quantize_and_compare_perplexity.py \ + --model_id --pass_name Rtn \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' \ + --device cuda:0 + +# Baseline + GPTQ +python scripts/quantize_and_compare_perplexity.py \ + --model_id --pass_name Gptq \ + --pass_config '{"bits": 4, "group_size": 128, "sym": true, "moe": true}' \ + --device cuda:0 +``` + +| Model | Baseline PPL | RTN PPL (Δ) | GPTQ PPL (Δ) | Quant time RTN / GPTQ | Fallback experts | +| --- | --- | --- | --- | --- | --- | +| ibm-granite/granite-3.0-1b-a400m-base | 6.2877 (354,564 tok) | 7.5861 (+1.2984) | 6.9492 (+0.6615) | 8.0s / 653.4s | 3/768 (0.4%) | +| allenai/OLMoE-1B-7B-0924 | 6.6182 (288,720 tok) | 7.1091 (+0.4909) | 6.8937 (+0.2755) | 52.6s / 1484.9s | 21/1024 (2.1%) | +| Qwen/Qwen1.5-MoE-A2.7B | 6.4246 (298,937 tok) | 6.9251 (+0.5005) | 6.6117 (+0.1872) | 100.0s / 2588.9s | 0/1440 (0.0%) | + +Calibration set for all GPTQ runs: 128 samples / 262,144 tokens (full WikiText-2 `train` split). +`moe_fallback_min_k_multiple=2.0` (`k=2`, the current default; see `moe-gptq.md`). All +baseline/quantized in-memory weight sizes are equal within each model (fake-quantization +dequantizes back to the original dtype for `transformers` compatibility) — only the *on-disk* +saved size actually shrinks; see the script's own summary output for per-run on-disk figures. + +**What this table demonstrates**: + +- GPTQ beat RTN on perplexity delta for every model tested (not just on average) — e.g. granite: + +0.6615 vs. +1.2984; Qwen1.5-MoE: +0.1872 vs. +0.5005 — at the cost of substantially longer + quantization time (minutes vs. seconds). +- MoE fallback rate is **not** proportional to total expert count — see `moe-gptq.md` for the + detailed explanation (Qwen1.5-MoE has the most total experts of the three but zero fallbacks; + the driver is per-model routing skew, not raw expert count). +- Quantization time scales roughly with `num_layers x num_experts` (the per-expert Cholesky solve + count), not with calibration token count — see `moe-gptq.md` for why a ~4x increase in + calibration tokens (from a stale `train[:1000]` slice to the full `train` split) only produced + an ~18% wall-time increase on granite, rather than the naively-expected ~4x. + +### `k=1` vs `k=2`: effect of the sufficiency-threshold multiplier + +`moe_fallback_min_k_multiple` was raised from `k=1` (the original shipped default, `N=K`) to +`k=2` (`N=2K`, the current default) after a full-team review found the real RTN-vs-GPTQ MSE +crossover under anisotropic activations at 4-bit sits closer to `1.5x-2x K`, not `1x K` (see +`moe-gptq.md`). All three models were rerun with the identical methodology at both settings to +quantify the actual effect: + +| Model | GPTQ PPL (Δ) at k=1 | GPTQ PPL (Δ) at k=2 | Fallback experts k=1 | Fallback experts k=2 | GPTQ time k=1 | GPTQ time k=2 | +| --- | --- | --- | --- | --- | --- | --- | +| granite-3.0-1b-a400m-base | 6.9560 (+0.6683) | 6.9492 (+0.6615) | 2/768 (0.3%) | 3/768 (0.4%) | 658.7s | 653.4s | +| OLMoE-1B-7B-0924 | 6.8966 (+0.2784) | 6.8937 (+0.2755) | 10/1024 (1.0%) | 21/1024 (2.1%) | 1499.3s | 1484.9s | +| Qwen1.5-MoE-A2.7B | 6.6117 (+0.1872) | 6.6117 (+0.1872) | 0/1440 (0.0%) | 0/1440 (0.0%) | 2475.6s | 2588.9s | + +Baseline and RTN-only numbers are identical between the two runs (RTN never reads +`moe_fallback_min_k_multiple`), confirming no other environment drift between the two benchmark +sessions. Despite roughly doubling the fallback count for granite and OLMoE, **perplexity did not +get measurably worse at `k=2` — it stayed flat or improved slightly**, and quantization time did +not increase (granite/OLMoE were slightly faster; Qwen1.5-MoE's small increase is within normal +run-to-run variance for a ~2,500s job, and it has 0 fallback at both settings so there is no +solve-count difference to explain it). The likely explanation: experts in the `1x-2x K` band have +technically-full-rank but severely ill-conditioned Hessians, so GPTQ's correction there was +already dominated by the damping prior rather than real signal — routing those borderline experts +to RTN at `k=2` removes the risk of an unlucky bad correction without giving up much upside, so +raising the threshold is close to a free win on the models tested here. + + +## Interpreting a run's console output + +Each run prints a `=== SUMMARY ===` block. For a MoE calibration run, also read the per-layer +`MoE coverage [...]` log lines emitted during quantization — each reports the effective fallback +threshold (`max(skew, sufficiency)`), how many experts were starved/unseen, and the observed +min/median/max token counts per expert for that layer. This is the fastest way to tell whether an +unexpectedly high fallback count is caused by routing imbalance in a specific layer or by an +under-sized calibration set overall. + +## Tips for running your own comparison + +- Start with a small/tiny model (e.g. a `*-tiny-random` HF checkpoint) to smoke-test your + `--pass_config` before committing GPU time to a multi-billion-parameter model — GPTQ + quantization time for a large MoE model can run into the tens of minutes. +- Pin `--dtype` explicitly if you need bit-for-bit reproducible baseline numbers across machines; + `auto` (the default) picks each checkpoint's native dtype, which is normally what you want for + a realistic baseline but can differ between environments with different default dtype handling. +- If you see a `transformers` warning about token sequence length exceeding + `max_position_embeddings` during eval-text tokenization, this is expected and harmless — the + perplexity computation uses a sliding window (`stride`/`max_len`), not a single full-sequence + forward pass; it does not indicate truncation or a scoring bug. diff --git a/skills/olive/references/quantization-onboarding.md b/skills/olive/references/quantization-onboarding.md new file mode 100644 index 0000000000..7b65a48799 --- /dev/null +++ b/skills/olive/references/quantization-onboarding.md @@ -0,0 +1,111 @@ +# Olive Quantization Onboarding + +An orientation to Olive's weight-quantization passes for PyTorch/Hugging Face models: what each +pass does, when to reach for it, the config knobs they share, and where to look in the codebase +before making changes. For MoE-specific GPTQ details (per-expert Hessians, fallback thresholds), +see [`moe-gptq.md`](moe-gptq.md). For a worked benchmark example comparing passes end-to-end, see +[`profiling-benchmark-example.md`](profiling-benchmark-example.md). + +## Where quantization passes live + +PyTorch/Hugging Face weight-quantization passes live in `olive/passes/pytorch/`: + +| Pass | Real class / registry name | Module | Calibration data? | Notes | +| --- | --- | --- | --- | --- | +| `Rtn` | `Rtn` | `rtn.py` | No | Round-to-nearest; fastest, data-free, weakest accuracy recovery. | +| `Gptq` | `Gptq` | `gptq.py` | Yes | Layerwise, Hessian-based weight correction using calibration data. | +| AutoGPTQ wrapper | `GptqQuantizer` | `autogptq.py` | Yes | Thin wrapper delegating to the third-party `auto-gptq` library. | +| GPTQModel wrapper | `GptqModel` | `gptqmodel.py` | Yes | Thin wrapper delegating to the third-party `gptqmodel` library. | +| AutoAWQ wrapper | `AutoAWQQuantizer` | `autoawq.py` | Yes | Wraps the third-party `autoawq` library (activation-aware weight quantization). | +| K-quant | `KQuant` | `kquant.py` | Varies | K-quant style block quantization. | + +The registry/class name (not a guessed lowercase-of-class-name module path) is what you must pass +to `--pass_name` for `scripts/quantize_and_compare_perplexity.py`, to `olive_config.json`'s +`"passes"` map, and to workflow configs — check `olive/olive_config.json` if you're ever unsure of +the exact registered name for a pass. + +ONNX-side quantization passes (post-export, operate on ONNX graphs rather than PyTorch modules) +live separately in `olive/passes/onnx/` (e.g. `rtn_quantization.py`, `hqq_quantization.py`, +`nvmo_quantization.py`, `inc_quantization.py`) — those are a different code path and are out of +scope for this doc, which focuses on the PyTorch-side `Rtn`/`Gptq` family. + +Shared logic (quantizer construction, model wrapping/unwrapping, layerwise iteration, save/load) +lives in `olive/passes/pytorch/quant_utils.py`. Read `get_quantizer_config()`, +`prepare_model()`, `run_layerwise_quantization()`, and `finalize()` there before modifying any +pass — almost every pass calls into these four functions and duplicating their logic in a new +pass is very rarely the right move. + +## Shared config surface + +Every weight-quantization pass accepts a common set of parameters from +`get_quantizer_config()` in `quant_utils.py`: + +| Param | Meaning | +| --- | --- | +| `bits` | Quantization bit-width (`PrecisionBits.BITS2/4/8`). | +| `group_size` | Block size for per-group scale/zero-point (`-1` means per-channel/whole-row). | +| `sym` | Symmetric (zero-point fixed at the bit-width's midpoint) vs. asymmetric quantization. | +| `lm_head` | Whether to also quantize the language-model head. | +| `embeds` | Whether to also quantize input embeddings (`Rtn` and `KQuant` only). | +| `overrides` | Per-module overrides for any of the above, keyed by module name pattern. | + +`Gptq` additionally exposes `damp_percent` (Hessian damping factor), `desc_act` (activation-order +column permutation, only valid for `group_size=-1`), and `data_config` (calibration dataset — see +below). When `moe=True` it also exposes `moe_fallback_threshold` and +`moe_fallback_min_k_multiple` — see `moe-gptq.md`. + +## RTN vs. GPTQ: when to use which + +- **RTN** is data-free and much faster than GPTQ (single-digit seconds for a ~1.3B active-param + model, tens of seconds to ~1-2 minutes for larger multi-billion-parameter MoE models in + practice — see `profiling-benchmark-example.md` for exact figures), but has the largest + accuracy regression of the two. Use it as a fast baseline, for environments without a + calibration dataset, or when the accuracy loss is acceptable for the target use case. +- **GPTQ** uses a calibration dataset to compute per-layer Hessians (`H = sum(x xT)` over + observed activations) and applies second-order error correction while quantizing each column, + substantially reducing the accuracy regression relative to RTN at the cost of a much longer + quantization pass (calibration forward passes + per-layer Cholesky/blockwise-quantize solves). + +Empirically (see `profiling-benchmark-example.md` for the full three-model MoE comparison), GPTQ +consistently produced a smaller perplexity regression than RTN on every model tested — but at +roughly 30-80x the quantization wall-time of RTN on the models measured there (single-run, +single-machine timings; treat the exact ratio as illustrative, not a guaranteed multiplier for +every model/hardware combination). Choose RTN when turnaround time matters more than the last bit +of accuracy; choose GPTQ when accuracy matters more and you can afford a longer one-time +quantization pass. + +## Calibration data (`data_config`) + +`Gptq` (and the other calibration-based passes) accept a `data_config` pointing at a Hugging +Face dataset config, or fall back to Olive's default WikiText-2 calibration set for +`HfModelHandler` inputs if none is given. See `olive/data/config.py` and +`configure-workflows/how-to-configure-data` in the Sphinx docs for how to point at a custom +dataset. + +Two split-hygiene facts worth internalizing: + +- WikiText-2's `train`/`validation`/`test` are official, pre-defined, non-overlapping Hugging + Face dataset splits (split by source Wikipedia article) — Olive does not construct or dedupe + them itself. A handful of near-identical rows across splits (e.g. repeated section-header + strings like `" = = Career = = "`) are not real leakage; they are trivial boilerplate that + recurs across unrelated articles. +- Use the `train` split (or a data-appropriate calibration set) for GPTQ calibration and a + disjoint split (e.g. `test`) for perplexity evaluation. Do not calibrate and evaluate on the + same rows. + +## Before modifying a quantization pass + +1. Read `quant_utils.py` fully — most behavior you might think belongs in a specific pass file is + actually implemented once, shared across passes. +2. Check `test/passes/pytorch/` for the existing test pattern for the pass you're touching (e.g. + `test_gptq.py`, `test_rtn.py`) before writing new tests; follow the existing + parametrization/fixture conventions rather than introducing a new test style. +3. If your change affects calibration or fallback behavior for Mixture-of-Experts models + specifically, read `moe-gptq.md` first — MoE calibration has its own module + (`moe_calib.py`) and several subtleties (per-expert Hessians, routing skew vs. statistical + sufficiency) that are easy to get wrong by analogy with the dense-model code path. + +If you're adding a brand-new pass rather than modifying an existing one, see the Sphinx guide +`docs/source/how-to/extending/how-to-add-optimization-pass.md` for how to register it in +`olive/olive_config.json` and wire it into the pass-discovery system — that step is not covered +here.