#!/usr/bin/env python3 """Quantise a checkpoint to four-bit AWQ with llm-compressor, for the non-GGUF row of the table. Purpose: produce the fifth file in the five-way comparison on the tracks that can run it - an activation-aware four-bit checkpoint that vLLM and transformers load directly, so the table contains one row that was not made by llama.cpp. AWQ is chosen over GPTQ here because its paper reports that it does no backpropagation or reconstruction and so generalises without overfitting the calibration set, which matters when the calibration data is somebody else's. llm-compressor is chosen over AutoAWQ because AutoAWQ's own README states that it "is officially deprecated and will no longer be maintained". Platform: spark, nvidia (llm-compressor targets CUDA and the vLLM deployment path). Track X can attempt it where a working PyTorch ROCm stack from Part 11 is present, and Track M should use the MLX path in the lab instead; mlx_lm.convert is the first-party quantiser there. Minimum memory: 24 GB for a 4B-class model at bfloat16 plus the calibration forward passes; 32 GB and above for an 8B-class model. Reduce --num-calibration-samples and --max-seq-length before reducing the model size. Assumes: a Python environment with torch, transformers and llmcompressor installed (`uv pip install llmcompressor`), a local model directory or Hugging Face model id, and network access for the calibration dataset unless one is already cached. Usage: python3 quantise-awq.py --model ~/models/runs/format-qwen3-4b-merged \ --out ~/models/awq/format-qwen3-4b-awq --labbook labbook.md python3 quantise-awq.py --model Qwen/Qwen3-4B --out ~/models/awq/Qwen3-4B-awq \ --num-calibration-samples 128 --max-seq-length 512 """ from __future__ import annotations import argparse import json import platform import sys import time from pathlib import Path def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--model", required=True, help="local model directory or Hugging Face model id to quantise") parser.add_argument("--out", required=True, help="directory to write the quantised model to") parser.add_argument("--dataset", default="perfectblend", help="calibration dataset, as named in the llm-compressor examples. To " "calibrate on your own text instead, follow llm-compressor's " "custom-dataset example; the recipe below does not change.") parser.add_argument("--split", default="train[:512]", help="dataset split expression") parser.add_argument("--num-calibration-samples", type=int, default=256) parser.add_argument("--max-seq-length", type=int, default=512) parser.add_argument("--scheme", default="W4A16_ASYM", help="quantisation scheme; W4A16_ASYM is four-bit weights with sixteen-bit " "activations, which is the configuration that helps single-user decode") parser.add_argument("--prompt", default="Hello my name is", help="one prompt run through the quantised model as a smoke test") parser.add_argument("--labbook", default=None) args = parser.parse_args() try: from transformers import AutoModelForCausalLM, AutoTokenizer from llmcompressor import oneshot from llmcompressor.modifiers.quantization import QuantizationModifier from llmcompressor.modifiers.transform.awq import AWQModifier except ImportError as exc: # pragma: no cover - depends on the reader's environment sys.exit( f"missing dependency: {exc}\n" "Install with: uv pip install llmcompressor\n" "Track M: use the MLX path in the lab instead; llm-compressor targets the CUDA and " "vLLM deployment path and Apple silicon is not one of its documented targets." ) print(f"==> Loading {args.model}") model = AutoModelForCausalLM.from_pretrained(args.model) tokenizer = AutoTokenizer.from_pretrained(args.model) # The recipe is two modifiers, in the order llm-compressor's own AWQ example uses. The first # computes the per-channel scales from calibration activations; the second does the actual # rounding. lm_head is left alone because quantising the output projection costs more than # it saves. recipe = [ AWQModifier(duo_scaling="both"), QuantizationModifier( ignore=["lm_head"], scheme=args.scheme, targets=["Linear"], ), ] print(f"==> Calibrating on {args.dataset} [{args.split}], " f"{args.num_calibration_samples} samples at {args.max_seq_length} tokens") print(" This is the step that decides which channels get protected. A calibration set") print(" that looks nothing like your traffic protects the wrong channels.") started = time.time() oneshot( model=model, dataset=args.dataset, splits=args.split, recipe=recipe, max_seq_length=args.max_seq_length, num_calibration_samples=args.num_calibration_samples, ) elapsed = time.time() - started print("\n========== SAMPLE GENERATION ==============") sample = tokenizer(args.prompt, return_tensors="pt") sample = {key: value.to(model.device) for key, value in sample.items()} output = model.generate(**sample, max_new_tokens=64) print(tokenizer.decode(output[0])) print("==========================================\n") print("If that is fluent, the quantisation did not break the model outright. It says nothing") print("about how much it cost; measure-quants.sh answers that.") out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) model.save_pretrained(str(out), save_compressed=True) tokenizer.save_pretrained(str(out)) print(f"\n==> Wrote {out}") total_bytes = sum(f.stat().st_size for f in out.rglob("*") if f.is_file()) record = { "lab": "part-16/lab-quantise-five-ways-and-measure/awq", "run_id": time.strftime("%Y%m%dT%H%M%S"), "source_model": args.model, "out": str(out), "scheme": args.scheme, "calibration_dataset": args.dataset, "calibration_split": args.split, "num_calibration_samples": args.num_calibration_samples, "max_seq_length": args.max_seq_length, "bytes": total_bytes, "seconds": round(elapsed, 1), "host": platform.platform(), "date": time.strftime("%Y-%m-%d"), } print(json.dumps(record, indent=2)) if args.labbook: with Path(args.labbook).open("a", encoding="utf-8") as handle: handle.write(json.dumps(record) + "\n") print(f"recorded in {args.labbook}") if __name__ == "__main__": main()