Model Quantization skill for Codex

Quantization is done when the quantized model meets an agreed quality budget against the full-precision baseline on the tasks that matter (for example, no more than 1 point drop on the main eval and no failure on a fixed set of critical prompts), and the memory and latency gains are measured on the target hardware and serving stack, not assumed.

Inspect first

1. Read AGENTS.md; identify the model (architecture, size, dtype), the serving runtime (vLLM, Transformers, TGI, llama.cpp, ONNX Runtime, mobile) and the hardware (GPU generation, CPU, Apple silicon).
2. Find or build the evaluation: a task eval with a pass threshold plus 50-200 representative prompts. Without a baseline score, quantization quality is a guess.
3. Record the baseline: eval scores, tokens/s or requests/s, p95 latency, peak memory.

Choose the scheme

Goal / hardware	Default	Why
LLM on NVIDIA Hopper/Ada/Blackwell, serving with vLLM	FP8 weights and activations (FP8_DYNAMIC)	Near-lossless, no calibration data, real speedups
LLM must fit on smaller GPUs, memory-bound decoding	INT4 weight-only (W4A16) via GPTQ or AWQ	~4x smaller weights; small quality loss when calibrated well
High-throughput LLM serving on Ampere+	INT8 weights and activations (W8A8, SmoothQuant + GPTQ)	Compute speedup at large batch
Quick experiment or QLoRA fine-tuning in Transformers	bitsandbytes NF4 4-bit	One-line load; slower inference than GPTQ/AWQ kernels
CPU or laptop inference	GGUF k-quants (see llama-cpp-gguf)	Best CPU kernels
Classic CNN/transformer encoder on CPU	ONNX Runtime INT8 dynamic (static for CNNs)	Easy 2-4x on CPU
PyTorch-native model in a PyTorch service	torchao weight-only INT8/INT4	Stays in PyTorch, works with torch.compile

Rules of thumb:

• Weight-only helps when decoding is memory-bound (small batch, LLM generation). Weight-and-activation helps when compute-bound (large batch, prefill, encoders).
• Smaller models lose more from 4-bit than large ones; below ~3B parameters, test INT8 or FP8 before INT4.
• Keep lm_head and embeddings in higher precision unless you have measured otherwise.

LLMs with llm-compressor (output loads in vLLM)

Needs a GPU, pip install llmcompressor (network) and the model weights locally.

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier

MODEL = "path/or/hub-id"
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="auto")
tok = AutoTokenizer.from_pretrained(MODEL)

# FP8 dynamic: no calibration data
fp8 = QuantizationModifier(targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"])

# INT4 weight-only with GPTQ: calibrate on data that looks like production prompts
ds = load_dataset("json", data_files="calib.jsonl", split="train").shuffle(seed=0).select(range(512))
ds = ds.map(lambda r: tok(tok.apply_chat_template(r["messages"], tokenize=False),
                          max_length=2048, truncation=True, add_special_tokens=False),
            remove_columns=ds.column_names)
w4 = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])

oneshot(model=model, dataset=ds, recipe=w4, max_seq_length=2048, num_calibration_samples=512)
model.save_pretrained("out/model-w4a16", save_compressed=True)
tok.save_pretrained("out/model-w4a16")

• AWQ is available as AWQModifier in llm-compressor; the standalone AutoAWQ and AutoGPTQ projects are no longer maintained, so do not start new work on them.
• Calibration data: 256-512 samples from your real prompt distribution (chat template applied) beat generic web text. For multilingual or code models, include those domains.
• Serve with vllm serve out/model-w4a16; vLLM reads the quantization config from the checkpoint (see vllm-serving).

Quick 4-bit in Transformers (bitsandbytes)

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                         bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
model = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=bnb, device_map="auto")

Good for fitting a model to experiment or for QLoRA (see lora-fine-tuning). For production serving throughput, export GPTQ/AWQ/FP8 instead.

Classic models

ONNX Runtime dynamic INT8 (CPU, transformers encoders, MLPs):

from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic("model.onnx", "model.int8.onnx", weight_type=QuantType.QInt8)

Static INT8 (CNNs) needs a CalibrationDataReader that yields ~100-500 real inputs and quantize_static(...); run ONNX Runtime’s pre-processing (python -m onnxruntime.quantization.preprocess --input model.onnx --output model.pre.onnx) first.

torchao (PyTorch-native):

from torchao.quantization import quantize_, Int8WeightOnlyConfig
quantize_(model, Int8WeightOnlyConfig())     # in place; then torch.compile(model)

The older torch.ao.quantization eager APIs are being phased out in favour of torchao.

Measuring quality loss

Run the same evals on baseline and quantized checkpoints, same prompts, same decoding (greedy or fixed seed):

lm_eval --model vllm --model_args pretrained=out/model-w4a16,max_model_len=4096 \
  --tasks gsm8k,mmlu,hellaswag --batch_size auto --output_path results/w4a16
lm_eval --model vllm --model_args pretrained=$BASELINE,max_model_len=4096 \
  --tasks gsm8k,mmlu,hellaswag --batch_size auto --output_path results/baseline

• Public benchmarks are a smoke test. The decision rests on your task eval (see llm-evals) and the critical-prompt set; quantization damage concentrates in math, long context, code and low-resource languages.
• Also compare perplexity on a held-out slice of your own text; a jump of more than a few percent signals a bad calibration set or a layer that should stay unquantized.
• Diff outputs on the critical prompts side by side; look for broken JSON, repetition and language switching.

Measuring the gain

On the target hardware and server: peak GPU memory, max concurrent sequences or KV cache size, tokens/s at the production batch size, time to first token and p95 latency. A 4-bit model that is smaller but slower at your batch size is not a win.

Workflow

• Read AGENTS.md; record model, runtime, hardware
• Baseline evals, critical prompts, latency and memory
• Choose scheme from the table; prepare calibration data from real prompts
• Quantize; save with config and the exact library versions
• Re-run evals and benchmarks; diff critical prompts
• Decide against the quality budget; if it fails, try a higher-precision scheme or leave sensitive layers unquantized before trying more calibration data

Output template

## Quantization report: <model>
Scheme: W4A16 GPTQ (llm-compressor <ver>), calibration: 512 prod prompts, lm_head kept BF16
| Metric | BF16 | W4A16 | Delta |
|---|---|---|---|
| Task eval (pass rate) | 0.842 | 0.835 | -0.7 pt |
| GSM8K | 0.781 | 0.762 | -1.9 pt |
| Peak GPU memory | 31.2 GB | 10.4 GB | -67% |
| Tokens/s at batch 32 | 2,410 | 3,050 | +27% |
Critical prompts: 48/50 unchanged, 2 minor wording changes, 0 failures
Decision: ship / do not ship, reason

Gotchas

• Licence terms of the base model still apply to quantized derivatives; check before publishing a checkpoint, and get user confirmation before uploading anywhere.
• A quantized checkpoint is tied to kernels: a format the runtime cannot accelerate falls back to slow paths or fails to load. Test loading in the exact serving version.
• Calibrating on the eval set inflates scores; keep calibration and eval data disjoint.
• KV cache, not weights, often dominates memory at long context; consider FP8 KV cache in the server before squeezing weights further.