Hugging Face's logo

dharun2049
/
MAE

Text Generation
Transformers
qwen3
quantization
low-bit
quinary
five-state
pq5
experimental
mmlu
conversational
Model card Files
xet
 Community
MAE / tools /pq5_quantize.py
dharun2049's picture
dharun2049
Add PQ5 reference quantizer
51e69f8 verified
Raw History Blame Contribute Delete
23.1 kB
#!/usr/bin/env python3
"""
pq5_quantize.py
Reference Pentary / Quinary post-training quantizer for Hugging Face safetensors models.
Representation
--------------
Per quantized tensor:
q in {-alpha, -1, 0, +1, +alpha}
w_hat = scale_g * q
where one FP16 scale is shared by each group of `group_size` weights and alpha
is learned once per tensor.
Packing
-------
Three base-5 digits form one symbol in [0, 124]:
symbol = d0 + 5*d1 + 25*d2
Eight 7-bit symbols are packed into seven bytes.
This yields 7/3 = 2.333... physical bits/weight for codes, before scales.
With FP16 scales and group_size=128:
7/3 + 16/128 ~= 2.4583 bits/weight
plus negligible tensor metadata / alpha overhead.
This is a storage + quality-reference implementation, NOT yet a fast inference
kernel. A custom Triton/CUDA/llama.cpp kernel should consume packed weights
directly for speed.
Example:
pip install torch safetensors huggingface_hub
python pq5_quantize.py \
--model Qwen/Qwen3-0.6B \
--output ./Qwen3-0.6B-PQ5 \
--device cuda \
--group-size 128 \
--alpha-grid 1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0 \
--verify
The script can also take a local model directory with --model /path/to/model.
"""
from __future__ import annotations
import argparse
import gc
import json
import math
import os
import re
import shutil
from pathlib import Path
from typing import Dict, Iterable, List, Tuple
import torch
from safetensors import safe_open
from safetensors.torch import save_file
# -----------------------------
# Utilities
# -----------------------------
def parse_size(s: str) -> int:
s = s.strip().upper()
m = re.fullmatch(r"([0-9]*\.?[0-9]+)\s*(B|KB|MB|GB|KIB|MIB|GIB)?", s)
if not m:
raise ValueError(f"Invalid size: {s}")
val = float(m.group(1))
unit = m.group(2) or "B"
mult = {
"B": 1,
"KB": 1000,
"MB": 1000**2,
"GB": 1000**3,
"KIB": 1024,
"MIB": 1024**2,
"GIB": 1024**3,
}[unit]
return int(val * mult)
def tensor_bytes(t: torch.Tensor) -> int:
return t.numel() * t.element_size()
def resolve_model(model: str, cache_dir: str | None = None) -> Path:
p = Path(model)
if p.exists():
return p.resolve()
from huggingface_hub import snapshot_download
local = snapshot_download(
repo_id=model,
cache_dir=cache_dir,
allow_patterns=[
"*.json",
"*.safetensors",
"*.model",
"*.txt",
"*.tiktoken",
"tokenizer*",
"merges.txt",
"vocab.json",
"*.py",
],
)
return Path(local)
def copy_non_weight_files(src: Path, dst: Path) -> None:
dst.mkdir(parents=True, exist_ok=True)
for p in src.iterdir():
if not p.is_file():
continue
if p.suffix == ".safetensors":
continue
if p.name.endswith(".safetensors.index.json"):
continue
# Avoid copying giant unrelated binaries.
if p.stat().st_size > 64 * 1024 * 1024:
continue
shutil.copy2(p, dst / p.name)
def list_safetensors(model_dir: Path) -> List[Path]:
files = sorted(model_dir.glob("*.safetensors"))
if not files:
raise FileNotFoundError(f"No .safetensors files found in {model_dir}")
return files
def iter_safetensor_items(files: Iterable[Path]):
for sf in files:
with safe_open(str(sf), framework="pt", device="cpu") as f:
for key in f.keys():
yield sf.name, key, f.get_tensor(key)
# -----------------------------
# PQ5 codebook optimization
# -----------------------------
@torch.no_grad()
def _assign_levels(x: torch.Tensor, scale: torch.Tensor, alpha: float):
"""
x: [n_groups, group_size], fp32
scale: [n_groups, 1], fp32
Returns:
idx: uint8/int64 code index in [0..4]
q: fp32 quantized normalized levels
"""
eps = 1e-12
z = x / scale.clamp_min(eps)
# Boundaries are exact midpoints between ordered levels
# [-a, -1, 0, 1, a].
b0 = -(alpha + 1.0) * 0.5
b1 = -0.5
b2 = 0.5
b3 = (alpha + 1.0) * 0.5
idx = torch.empty_like(z, dtype=torch.uint8)
idx[z < b0] = 0
idx[(z >= b0) & (z < b1)] = 1
idx[(z >= b1) & (z < b2)] = 2
idx[(z >= b2) & (z < b3)] = 3
idx[z >= b3] = 4
q = torch.zeros_like(z)
q[idx == 0] = -alpha
q[idx == 1] = -1.0
q[idx == 2] = 0.0
q[idx == 3] = 1.0
q[idx == 4] = alpha
return idx, q
@torch.no_grad()
def _fit_scale_for_alpha(
x: torch.Tensor,
alpha: float,
iters: int = 4,
) -> Tuple[torch.Tensor, torch.Tensor, float]:
"""
Alternating nearest-code assignment and least-squares group-scale fitting.
x is [groups, group_size] float32.
"""
# Robust initial scale. 2*mean(|x|)/(1+alpha) roughly centers both
# non-zero magnitudes around the observed absolute-weight mass.
mean_abs = x.abs().mean(dim=1, keepdim=True)
scale = (2.0 * mean_abs / (1.0 + alpha)).clamp_min(1e-8)
idx = None
q = None
for _ in range(iters):
idx, q = _assign_levels(x, scale, alpha)
num = (x * q).sum(dim=1, keepdim=True)
den = (q * q).sum(dim=1, keepdim=True).clamp_min(1e-8)
scale = (num / den).abs().clamp_min(1e-8)
idx, q = _assign_levels(x, scale, alpha)
recon = q * scale
mse = torch.mean((x - recon) ** 2).item()
return scale, idx, mse
@torch.no_grad()
def choose_alpha(
flat: torch.Tensor,
group_size: int,
alpha_grid: List[float],
device: str,
sample_weights: int = 1_048_576,
iters: int = 4,
) -> Tuple[float, Dict[float, float]]:
"""
Pick one alpha per tensor using a deterministic sample of groups.
"""
n = flat.numel()
n_groups = math.ceil(n / group_size)
# Deterministic evenly-spaced group sample.
max_groups = max(1, sample_weights // group_size)
take_groups = min(n_groups, max_groups)
if take_groups == n_groups:
group_ids = torch.arange(n_groups)
else:
group_ids = torch.linspace(0, n_groups - 1, take_groups).round().long().unique()
# Build sampled groups without padding the entire tensor on device.
cpu = flat.detach().float().cpu()
chunks = []
for gid in group_ids.tolist():
start = gid * group_size
g = cpu[start : start + group_size]
if g.numel() < group_size:
g = torch.nn.functional.pad(g, (0, group_size - g.numel()))
chunks.append(g)
sample = torch.stack(chunks, dim=0).to(device=device, dtype=torch.float32)
scores = {}
for a in alpha_grid:
_, _, mse = _fit_scale_for_alpha(sample, float(a), iters=iters)
scores[float(a)] = float(mse)
best = min(scores, key=scores.get)
del sample, chunks
if device.startswith("cuda"):
torch.cuda.empty_cache()
return float(best), scores
@torch.no_grad()
def quantize_tensor_pq5(
w: torch.Tensor,
group_size: int,
alpha_grid: List[float],
device: str,
group_chunk: int = 8192,
sample_weights: int = 1_048_576,
iters: int = 4,
):
"""
Quantize a floating tensor into:
packed uint8 byte stream
FP16 group scales
FP16 scalar alpha
Returns dict and stats.
"""
original_shape = list(w.shape)
original_numel = w.numel()
flat_cpu = w.detach().float().reshape(-1).cpu()
alpha, alpha_scores = choose_alpha(
flat_cpu,
group_size=group_size,
alpha_grid=alpha_grid,
device=device,
sample_weights=sample_weights,
iters=iters,
)
padded_numel = math.ceil(original_numel / group_size) * group_size
pad = padded_numel - original_numel
if pad:
flat_cpu = torch.nn.functional.pad(flat_cpu, (0, pad))
groups_cpu = flat_cpu.view(-1, group_size)
all_codes = []
all_scales = []
sqerr = 0.0
signal = 0.0
for start in range(0, groups_cpu.shape[0], group_chunk):
x = groups_cpu[start : start + group_chunk].to(device=device, dtype=torch.float32)
scales, idx, _ = _fit_scale_for_alpha(x, alpha, iters=iters)
# Stats
level_values = torch.tensor(
[-alpha, -1.0, 0.0, 1.0, alpha],
device=device,
dtype=torch.float32,
)
q = level_values[idx.long()]
recon = q * scales
sqerr += torch.sum((x - recon) ** 2).item()
signal += torch.sum(x ** 2).item()
all_codes.append(idx.reshape(-1).cpu())
all_scales.append(scales.reshape(-1).to(torch.float16).cpu())
del x, scales, idx, q, recon
if device.startswith("cuda"):
torch.cuda.empty_cache()
codes = torch.cat(all_codes, dim=0)
scales = torch.cat(all_scales, dim=0)
packed, code_pad = pack_base5_codes(codes)
rmse_rel = math.sqrt(sqerr / max(signal, 1e-30))
code_hist = torch.bincount(codes.long(), minlength=5).tolist()
result = {
"packed": packed.contiguous(),
"scales": scales.contiguous(),
"alpha": torch.tensor([alpha], dtype=torch.float16),
}
stats = {
"shape": original_shape,
"numel": original_numel,
"padded_numel": padded_numel,
"group_size": group_size,
"alpha": alpha,
"alpha_mse": alpha_scores,
"relative_rmse": rmse_rel,
"code_histogram": code_hist,
"code_pad": code_pad,
"packed_bytes": packed.numel(),
"scale_bytes": tensor_bytes(scales),
"effective_bpw_tensor": 8.0 * (packed.numel() + tensor_bytes(scales) + 2) / original_numel,
}
del flat_cpu, groups_cpu, all_codes, all_scales, codes
return result, stats
# -----------------------------
# Exact 5-ary packing
# -----------------------------
@torch.no_grad()
def pack_base5_codes(codes: torch.Tensor) -> Tuple[torch.Tensor, int]:
"""
codes in {0,1,2,3,4}.
3 digits -> one 7-bit symbol; 8 symbols -> 7 bytes.
Returns (packed_bytes_uint8, number_of_padding_codes).
"""
c = codes.to(torch.uint8).reshape(-1).cpu()
if c.numel() and int(c.max()) > 4:
raise ValueError("PQ5 codes must be in [0,4]")
code_pad = (-c.numel()) % 3
if code_pad:
# Pad with zero-level code 2, not with code 0.
c = torch.cat([c, torch.full((code_pad,), 2, dtype=torch.uint8)])
tri = c.view(-1, 3).to(torch.int16)
symbols = tri[:, 0] + 5 * tri[:, 1] + 25 * tri[:, 2] # [0,124]
symbols = symbols.to(torch.uint8)
sym_pad = (-symbols.numel()) % 8
if sym_pad:
# symbol representing three zero-level digits: 2 + 10 + 50 = 62
symbols = torch.cat([symbols, torch.full((sym_pad,), 62, dtype=torch.uint8)])
s = symbols.view(-1, 8).to(torch.int16)
out = torch.empty((s.shape[0], 7), dtype=torch.uint8)
out[:, 0] = (s[:, 0] | ((s[:, 1] & 0x01) << 7)).to(torch.uint8)
out[:, 1] = ((s[:, 1] >> 1) | ((s[:, 2] & 0x03) << 6)).to(torch.uint8)
out[:, 2] = ((s[:, 2] >> 2) | ((s[:, 3] & 0x07) << 5)).to(torch.uint8)
out[:, 3] = ((s[:, 3] >> 3) | ((s[:, 4] & 0x0F) << 4)).to(torch.uint8)
out[:, 4] = ((s[:, 4] >> 4) | ((s[:, 5] & 0x1F) << 3)).to(torch.uint8)
out[:, 5] = ((s[:, 5] >> 5) | ((s[:, 6] & 0x3F) << 2)).to(torch.uint8)
out[:, 6] = ((s[:, 6] >> 6) | (s[:, 7] << 1)).to(torch.uint8)
return out.reshape(-1), int(code_pad)
@torch.no_grad()
def unpack_base5_codes(
packed: torch.Tensor,
n_codes: int,
) -> torch.Tensor:
"""
Reference inverse of pack_base5_codes. Returns uint8 codes [0..4].
"""
b = packed.to(torch.uint8).reshape(-1).cpu()
if b.numel() % 7 != 0:
raise ValueError("Packed PQ5 stream length must be a multiple of 7 bytes.")
x = b.view(-1, 7).to(torch.int16)
s = torch.empty((x.shape[0], 8), dtype=torch.int16)
s[:, 0] = x[:, 0] & 0x7F
s[:, 1] = ((x[:, 0] >> 7) | ((x[:, 1] & 0x3F) << 1))
s[:, 2] = ((x[:, 1] >> 6) | ((x[:, 2] & 0x1F) << 2))
s[:, 3] = ((x[:, 2] >> 5) | ((x[:, 3] & 0x0F) << 3))
s[:, 4] = ((x[:, 3] >> 4) | ((x[:, 4] & 0x07) << 4))
s[:, 5] = ((x[:, 4] >> 3) | ((x[:, 5] & 0x03) << 5))
s[:, 6] = ((x[:, 5] >> 2) | ((x[:, 6] & 0x01) << 6))
s[:, 7] = x[:, 6] >> 1
symbols = s.reshape(-1)
d0 = symbols % 5
d1 = (symbols // 5) % 5
d2 = (symbols // 25) % 5
codes = torch.stack([d0, d1, d2], dim=1).reshape(-1).to(torch.uint8)
return codes[:n_codes]
@torch.no_grad()
def dequantize_pq5(
packed: torch.Tensor,
scales: torch.Tensor,
alpha: torch.Tensor | float,
shape: List[int] | Tuple[int, ...],
group_size: int = 128,
dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
"""
Reference dequantizer. Useful for verification / fake-quant evaluation.
"""
n = math.prod(shape)
padded_n = math.ceil(n / group_size) * group_size
codes = unpack_base5_codes(packed, padded_n).long()
a = float(alpha.item()) if isinstance(alpha, torch.Tensor) else float(alpha)
lut = torch.tensor([-a, -1.0, 0.0, 1.0, a], dtype=torch.float32)
q = lut[codes].view(-1, group_size)
s = scales.float().view(-1, 1)
w = (q * s).reshape(-1)[:n].reshape(shape)
return w.to(dtype)
# -----------------------------
# Sharded writer
# -----------------------------
class ShardWriter:
def __init__(self, out_dir: Path, max_shard_bytes: int):
self.out_dir = out_dir
self.max_shard_bytes = max_shard_bytes
self.pending: Dict[str, torch.Tensor] = {}
self.pending_bytes = 0
self.shards: List[str] = []
self.weight_map: Dict[str, str] = {}
self.shard_no = 0
def add(self, key: str, tensor: torch.Tensor):
nbytes = tensor_bytes(tensor)
if self.pending and self.pending_bytes + nbytes > self.max_shard_bytes:
self.flush()
self.pending[key] = tensor.contiguous().cpu()
self.pending_bytes += nbytes
def flush(self):
if not self.pending:
return
self.shard_no += 1
fname = f"model-pq5-{self.shard_no:05d}.safetensors"
save_file(self.pending, str(self.out_dir / fname))
for k in self.pending:
self.weight_map[k] = fname
self.shards.append(fname)
self.pending = {}
self.pending_bytes = 0
gc.collect()
def finalize(self):
self.flush()
# -----------------------------
# CLI
# -----------------------------
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True, help="HF repo id or local model directory")
ap.add_argument("--output", required=True)
ap.add_argument("--cache-dir", default=None)
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
ap.add_argument("--group-size", type=int, default=128)
ap.add_argument(
"--alpha-grid",
default="1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0",
help="Comma-separated outer-magnitude candidates. Must all be > 1.",
)
ap.add_argument("--alpha-iters", type=int, default=4)
ap.add_argument("--alpha-sample-weights", type=int, default=1_048_576)
ap.add_argument("--group-chunk", type=int, default=8192)
ap.add_argument("--max-shard-size", default="1GB")
ap.add_argument(
"--min-elements",
type=int,
default=4096,
help="Only quantize floating 2D+ tensors at least this large.",
)
ap.add_argument(
"--exclude-regex",
default=r"(?:^|\.)(?:norm|layernorm|ln_)(?:\.|$)",
help="Regex of tensor names to leave unchanged.",
)
ap.add_argument(
"--quantize-embeddings",
action=argparse.BooleanOptionalAction,
default=True,
)
ap.add_argument(
"--quantize-lm-head",
action=argparse.BooleanOptionalAction,
default=True,
)
ap.add_argument("--verify", action="store_true")
args = ap.parse_args()
alpha_grid = [float(x) for x in args.alpha_grid.split(",") if x.strip()]
if not alpha_grid or any(a <= 1.0 for a in alpha_grid):
raise ValueError("Every alpha candidate must be > 1.0")
if args.group_size <= 0:
raise ValueError("--group-size must be positive")
src = resolve_model(args.model, args.cache_dir)
out = Path(args.output)
out.mkdir(parents=True, exist_ok=True)
copy_non_weight_files(src, out)
files = list_safetensors(src)
writer = ShardWriter(out, parse_size(args.max_shard_size))
manifest = {
"format": "PQ5",
"version": 1,
"description": "5-state groupwise PTQ with tensorwise learned alpha and exact base-5 bit packing",
"codebook": "[-alpha, -1, 0, +1, +alpha]",
"group_size": args.group_size,
"packing": {
"digits_per_symbol": 3,
"symbol_bits": 7,
"symbols_per_block": 8,
"bytes_per_block": 7,
"code_bits_per_weight": 7.0 / 3.0,
},
"tensors": {},
}
total_orig_bytes = 0
total_stored_bytes = 0
total_quantized_numel = 0
total_passthrough_bytes = 0
exclude = re.compile(args.exclude_regex) if args.exclude_regex else None
print("=" * 88)
print("PQ5 :: Pentary / Quinary post-training quantizer")
print(f"source : {src}")
print(f"output : {out}")
print(f"device : {args.device}")
print(f"group size : {args.group_size}")
print(f"alpha grid : {alpha_grid}")
print("=" * 88)
tensor_idx = 0
for source_shard, name, t in iter_safetensor_items(files):
tensor_idx += 1
orig_bytes = tensor_bytes(t)
total_orig_bytes += orig_bytes
is_float = t.dtype in {
torch.float16,
torch.bfloat16,
torch.float32,
torch.float64,
}
quantizable = is_float and t.ndim >= 2 and t.numel() >= args.min_elements
if exclude and exclude.search(name):
quantizable = False
lname = name.lower()
if not args.quantize_embeddings and ("embed" in lname or "embedding" in lname):
quantizable = False
if not args.quantize_lm_head and ("lm_head" in lname or lname.endswith("output.weight")):
quantizable = False
if not quantizable:
writer.add(name, t)
manifest["tensors"][name] = {
"kind": "passthrough",
"shape": list(t.shape),
"dtype": str(t.dtype).replace("torch.", ""),
"source_shard": source_shard,
}
total_stored_bytes += orig_bytes
total_passthrough_bytes += orig_bytes
print(f"[{tensor_idx:04d}] PASS {name:70s} {list(t.shape)} {t.dtype}")
del t
continue
print(f"[{tensor_idx:04d}] PQ5 {name:70s} {list(t.shape)} {t.dtype}")
q, stats = quantize_tensor_pq5(
t,
group_size=args.group_size,
alpha_grid=alpha_grid,
device=args.device,
group_chunk=args.group_chunk,
sample_weights=args.alpha_sample_weights,
iters=args.alpha_iters,
)
k_packed = name + ".pq5_packed"
k_scales = name + ".pq5_scales"
k_alpha = name + ".pq5_alpha"
writer.add(k_packed, q["packed"])
writer.add(k_scales, q["scales"])
writer.add(k_alpha, q["alpha"])
stored = (
tensor_bytes(q["packed"])
+ tensor_bytes(q["scales"])
+ tensor_bytes(q["alpha"])
)
total_stored_bytes += stored
total_quantized_numel += t.numel()
manifest["tensors"][name] = {
"kind": "pq5",
"shape": list(t.shape),
"dtype": str(t.dtype).replace("torch.", ""),
"source_shard": source_shard,
"packed_key": k_packed,
"scales_key": k_scales,
"alpha_key": k_alpha,
"group_size": args.group_size,
"padded_numel": stats["padded_numel"],
"code_pad": stats["code_pad"],
"alpha": stats["alpha"],
"relative_rmse": stats["relative_rmse"],
"code_histogram": stats["code_histogram"],
"effective_bpw_tensor": stats["effective_bpw_tensor"],
}
print(
f" alpha={stats['alpha']:.3f} "
f"relRMSE={stats['relative_rmse']:.5f} "
f"bpw={stats['effective_bpw_tensor']:.4f} "
f"hist={stats['code_histogram']}"
)
if args.verify:
deq = dequantize_pq5(
q["packed"],
q["scales"],
q["alpha"],
shape=list(t.shape),
group_size=args.group_size,
dtype=torch.float32,
)
# Packing verification: quantization stats already measure distortion;
# this catches byte-packing/depacking corruption.
if not torch.isfinite(deq).all():
raise RuntimeError(f"Non-finite dequantized values in {name}")
print(
f" verify=OK dequant_range=[{deq.min().item():.5g}, {deq.max().item():.5g}]"
)
del deq
del t, q
gc.collect()
if args.device.startswith("cuda"):
torch.cuda.empty_cache()
writer.finalize()
manifest["summary"] = {
"original_tensor_bytes": total_orig_bytes,
"stored_tensor_bytes": total_stored_bytes,
"compression_ratio": total_orig_bytes / max(total_stored_bytes, 1),
"quantized_parameters": total_quantized_numel,
"passthrough_bytes": total_passthrough_bytes,
"overall_effective_bpw_vs_original_param_count": (
8.0 * total_stored_bytes / max(total_quantized_numel, 1)
),
}
manifest["weight_map"] = writer.weight_map
manifest["shards"] = writer.shards
with open(out / "pq5_manifest.json", "w", encoding="utf-8") as f:
json.dump(manifest, f, indent=2)
with open(out / "model.pq5.index.json", "w", encoding="utf-8") as f:
json.dump(
{
"metadata": {
"format": "PQ5",
"group_size": args.group_size,
"total_size": total_stored_bytes,
},
"weight_map": writer.weight_map,
},
f,
indent=2,
)
print("=" * 88)
print("DONE")
print(f"Original tensor bytes : {total_orig_bytes / 1e9:.3f} GB")
print(f"Stored tensor bytes : {total_stored_bytes / 1e9:.3f} GB")
print(f"Compression ratio : {total_orig_bytes / max(total_stored_bytes, 1):.3f}x")
if total_quantized_numel:
quant_only_ideal = 7.0 / 3.0 + 16.0 / args.group_size
print(f"PQ5 quantized ideal : {quant_only_ideal:.4f} bits/weight + tiny alpha/packing pad")
print(f"Manifest : {out / 'pq5_manifest.json'}")
print("=" * 88)
if __name__ == "__main__":
main()