366 lines
16 KiB
Python
366 lines
16 KiB
Python
"""
|
|
OrbitVQ weight-only quantization backend (orbitvq2 / orbitvq3 / orbitvq4 qtypes).
|
|
|
|
Extends the OrbitQuant recipe (toolkit/util/orbit_quant.py) with three accuracy
|
|
upgrades aimed at low bit-widths, at the cost of no longer being the paper's pure
|
|
scalar method:
|
|
|
|
1. Lattice vector codebooks. Groups of coordinates are quantized jointly against a
|
|
codebook built from the densest lattice for the group dimension (E8 for 8-dim
|
|
groups at 2 bits, D4 for 4-dim groups at 3/4 bits) instead of one coordinate at
|
|
a time. A scalar codebook wastes indices on corner combinations that iid
|
|
gaussian coordinates never produce; a lattice codebook spends all its codewords
|
|
on the spherical shell where rotated weights actually live.
|
|
2. Per-group scales. One scale per GROUP_SIZE rotated coordinates (instead of one
|
|
norm per row) absorbs the finite-sample energy fluctuation between groups.
|
|
3. Least-squares scale refit. After codes are chosen the scale is refit to the
|
|
MSE-optimal value <w, c> / <c, c>.
|
|
|
|
Everything stays data-free and deterministic: the same seeded RPBH rotation as
|
|
OrbitQuant, codebooks that are fixed mathematical objects (lattice points sorted by
|
|
norm), and precomputed distortion-optimal lattice scales for the N(0,1) source.
|
|
|
|
Encoding uses the closed-form nearest-lattice-point algorithm plus a hash lookup
|
|
into the truncated codebook; the rare vectors whose nearest lattice point falls
|
|
outside the codebook fall back to an exact brute-force search.
|
|
|
|
Bits per param: index bits exactly (2/3/4) + 16-bit group scales / GROUP_SIZE
|
|
(+0.125 at the default 128).
|
|
|
|
Quantized state attached to each module:
|
|
ovq_packed packed codeword indices (uint8 bitstream, index_bits per vector)
|
|
ovq_scales per (row, group) least-squares scales (original dtype)
|
|
ovq_perm / ovq_inv_perm / ovq_signs shared RPBH rotation for in_features
|
|
ovq_block / ovq_group Hadamard block size and group size
|
|
"""
|
|
|
|
import math
|
|
from typing import Dict, Tuple
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from toolkit.print import print_acc
|
|
from toolkit.util.ostris_quant import OstrisQuantizer
|
|
from toolkit.util.orbit_quant import (
|
|
MIN_HADAMARD_BLOCK,
|
|
hadamard_block_size,
|
|
rpbh_forward,
|
|
rpbh_inverse,
|
|
rpbh_params,
|
|
)
|
|
|
|
# qtype name -> quantizer constructor kwargs
|
|
ORBIT_VQ_QTYPES = {
|
|
"orbitvq2": {"bits": 2, "vec_dim": 8, "lattice": "E8", "codebook_size": 2 ** 16},
|
|
"orbitvq3": {"bits": 3, "vec_dim": 4, "lattice": "D4", "codebook_size": 2 ** 12},
|
|
"orbitvq4": {"bits": 4, "vec_dim": 4, "lattice": "D4", "codebook_size": 2 ** 16},
|
|
}
|
|
|
|
# coordinates per stored scale (the group-norm granularity)
|
|
GROUP_SIZE = 128
|
|
|
|
# encode -> least-squares scale refit rounds. re-encoding with the refit scale
|
|
# (rounds=2) measured only 0.0005 lower relative error at 2-bit while doubling the
|
|
# encode cost, so one round is the default
|
|
LS_REFIT_ROUNDS = 1
|
|
|
|
# distortion-optimal lattice scale for a unit-variance gaussian source, per
|
|
# (lattice, codebook_size), found by an exact nearest-codeword sweep over a seeded
|
|
# 60k gaussian sample. per-coordinate MSE at these scales (vs scalar Lloyd-Max):
|
|
# E8-65536 2-bit: 0.0916 (0.1175) D4-4096 3-bit: 0.0301 (0.0345)
|
|
# D4-65536 4-bit: 0.0089 (0.0095)
|
|
# fixed constants so codebooks are bit-identical everywhere.
|
|
BETA = {
|
|
("E8", 2 ** 16): 0.9800,
|
|
("D4", 2 ** 12): 0.4722,
|
|
("D4", 2 ** 16): 0.2617,
|
|
}
|
|
|
|
# key packing for the hash lookup: doubled coordinates + _KEY_OFFSET must fit in
|
|
# _KEY_BITS bits per dimension (covers every codebook point of the sizes above)
|
|
_KEY_BITS = 6
|
|
_KEY_OFFSET = 32
|
|
|
|
_skip_warned = set()
|
|
_master_tables: Dict[Tuple[str, int], "_VQTables"] = {}
|
|
_device_tables: Dict[Tuple[str, int, str], "_VQTables"] = {}
|
|
|
|
|
|
def enumerate_lattice_codebook(lattice: str, size: int) -> torch.Tensor:
|
|
"""The `size` lattice points closest to the origin (ties broken lexicographically,
|
|
so the result is fully deterministic), as float32 (size, dim), in lattice units.
|
|
|
|
D4 = {x in Z^4 : sum(x) even} (densest 4-dim packing).
|
|
E8 = D8 union (D8 + 1/2) (densest 8-dim packing). In doubled coordinates both
|
|
become: uniform-parity integer vectors with sum divisible by 4.
|
|
"""
|
|
if lattice == "D4":
|
|
dim, reach = 4, 27 # doubled coords in [-26, 26], covers 65536 points
|
|
vals = torch.arange(-(reach - 1), reach, 2, dtype=torch.int32) # even only
|
|
parities = [vals]
|
|
elif lattice == "E8":
|
|
dim = 8
|
|
# doubled coords: even in [-6, 6] (integer points, norm^2 <= 12) and odd in
|
|
# [-5, 5] (half points); both ranges cover far more than 65536 points
|
|
parities = [
|
|
torch.arange(-6, 7, 2, dtype=torch.int32),
|
|
torch.arange(-5, 6, 2, dtype=torch.int32),
|
|
]
|
|
else:
|
|
raise ValueError(f"unknown lattice {lattice}")
|
|
|
|
kept = []
|
|
for vals in parities:
|
|
# chunk over the first coordinate to keep the cartesian product small
|
|
for v0 in vals.tolist():
|
|
rest = torch.cartesian_prod(*([vals] * (dim - 1)))
|
|
pts = torch.cat(
|
|
[torch.full((rest.shape[0], 1), v0, dtype=torch.int32), rest], dim=1
|
|
)
|
|
pts = pts[pts.sum(dim=1).remainder(4) == 0]
|
|
norm2 = (pts.to(torch.int64) ** 2).sum(dim=1)
|
|
# generous radius cut just to bound memory before the exact sort below
|
|
keep = norm2 <= (48 if lattice == "E8" else 26 ** 2 + 1)
|
|
kept.append(pts[keep])
|
|
pts = torch.cat(kept)
|
|
|
|
# sort by (norm, lexicographic key) and take the closest `size` points
|
|
norm2 = (pts.to(torch.int64) ** 2).sum(dim=1)
|
|
key = _point_keys(pts)
|
|
order = torch.argsort(norm2 * (1 << (_KEY_BITS * dim)) + key)
|
|
pts = pts[order[:size]]
|
|
if pts.shape[0] < size:
|
|
raise RuntimeError(f"lattice enumeration too small for {lattice}/{size}")
|
|
return pts.to(torch.float32) / 2.0 # back to lattice units
|
|
|
|
|
|
def _point_keys(doubled_pts: torch.Tensor) -> torch.Tensor:
|
|
"""int64 hash key per point from doubled integer coordinates."""
|
|
digits = doubled_pts.to(torch.int64) + _KEY_OFFSET
|
|
key = torch.zeros(doubled_pts.shape[0], dtype=torch.int64, device=doubled_pts.device)
|
|
for i in range(doubled_pts.shape[1]):
|
|
key = key | (digits[:, i] << (_KEY_BITS * i))
|
|
return key
|
|
|
|
|
|
def _round_Dn(x: torch.Tensor) -> torch.Tensor:
|
|
"""Nearest point of D_n (integer vectors with even coordinate sum) to each row."""
|
|
f = x.round()
|
|
odd = f.to(torch.int64).sum(dim=-1).remainder(2) != 0
|
|
err = x - f
|
|
idx = err.abs().argmax(dim=-1, keepdim=True)
|
|
step = torch.where(err.gather(-1, idx) >= 0, 1.0, -1.0).to(x.dtype)
|
|
adjusted = f.scatter(-1, idx, f.gather(-1, idx) + step)
|
|
return torch.where(odd.unsqueeze(-1), adjusted, f)
|
|
|
|
|
|
def _round_lattice(x: torch.Tensor, lattice: str) -> torch.Tensor:
|
|
"""Closed-form nearest lattice point (Conway & Sloane) to each row of x."""
|
|
a = _round_Dn(x)
|
|
if lattice == "D4":
|
|
return a
|
|
# E8 = D8 union (D8 + 1/2): take the closer of the two cosets
|
|
b = _round_Dn(x - 0.5) + 0.5
|
|
da = (x - a).square().sum(dim=-1)
|
|
db = (x - b).square().sum(dim=-1)
|
|
return torch.where((da <= db).unsqueeze(-1), a, b)
|
|
|
|
|
|
class _VQTables:
|
|
"""Codebook + hash tables for one (lattice, codebook_size), on one device."""
|
|
|
|
def __init__(self, lattice: str, size: int):
|
|
self.lattice = lattice
|
|
self.size = size
|
|
self.beta = BETA[(lattice, size)]
|
|
points = enumerate_lattice_codebook(lattice, size)
|
|
self.codebook = points * self.beta # (size, dim) float32, source units
|
|
keys = _point_keys((points * 2).to(torch.int32))
|
|
self.sorted_keys, order = torch.sort(keys)
|
|
self.key_to_index = order.to(torch.int32)
|
|
# for the brute-force fallback: argmax(z.c - |c|^2/2) == nearest codeword
|
|
self.half_sq_norms = self.codebook.square().sum(dim=1) / 2
|
|
self.codebook_t = self.codebook.T.contiguous()
|
|
|
|
def to(self, device: torch.device) -> "_VQTables":
|
|
out = object.__new__(_VQTables)
|
|
out.lattice, out.size, out.beta = self.lattice, self.size, self.beta
|
|
out.codebook = self.codebook.to(device)
|
|
out.sorted_keys = self.sorted_keys.to(device)
|
|
out.key_to_index = self.key_to_index.to(device)
|
|
out.half_sq_norms = self.half_sq_norms.to(device)
|
|
out.codebook_t = self.codebook_t.to(device)
|
|
return out
|
|
|
|
|
|
def get_vq_tables(lattice: str, size: int, device) -> _VQTables:
|
|
mkey = (lattice, size)
|
|
if mkey not in _master_tables:
|
|
_master_tables[mkey] = _VQTables(lattice, size)
|
|
dkey = (lattice, size, str(device))
|
|
if dkey not in _device_tables:
|
|
_device_tables[dkey] = _master_tables[mkey].to(torch.device(device))
|
|
return _device_tables[dkey]
|
|
|
|
|
|
def encode_vectors(z: torch.Tensor, tables: _VQTables) -> torch.Tensor:
|
|
"""Exact nearest-codeword indices (int32) for rows of z (float32, source units).
|
|
|
|
Closed-form lattice rounding + hash lookup; rows whose nearest lattice point is
|
|
outside the truncated codebook fall back to a brute-force search.
|
|
"""
|
|
dim = z.shape[-1]
|
|
p = _round_lattice(z / tables.beta, tables.lattice)
|
|
digits = (p * 2).round().to(torch.int64) + _KEY_OFFSET
|
|
in_range = ((digits >= 0) & (digits < (1 << _KEY_BITS))).all(dim=-1)
|
|
key = torch.zeros(z.shape[0], dtype=torch.int64, device=z.device)
|
|
for i in range(dim):
|
|
key = key | (digits[:, i].clamp(0, (1 << _KEY_BITS) - 1) << (_KEY_BITS * i))
|
|
pos = torch.searchsorted(tables.sorted_keys, key).clamp(max=tables.size - 1)
|
|
hit = in_range & (tables.sorted_keys.gather(0, pos) == key)
|
|
idx = tables.key_to_index.gather(0, pos.to(torch.int64))
|
|
|
|
miss = ~hit
|
|
n_miss = int(miss.sum())
|
|
if n_miss > 0:
|
|
# on cuda run the search in fp16 (tensor cores + half the score-matrix
|
|
# traffic); the score gap between competing codewords for these overload
|
|
# vectors is far above fp16 noise. cpu stays fp32.
|
|
dt = torch.float16 if z.device.type == "cuda" else torch.float32
|
|
z_miss = z[miss].to(dt)
|
|
cb_t = tables.codebook_t.to(dt)
|
|
half_norms = tables.half_sq_norms.to(dt)
|
|
found = torch.empty(n_miss, dtype=torch.int32, device=z.device)
|
|
# chunk rows so the score matrix stays ~256MB while each matmul stays large
|
|
# enough to saturate the gpu (overload vectors are ~15-20% of the total)
|
|
chunk = max(256, (2 ** 26) // tables.size)
|
|
for s in range(0, n_miss, chunk):
|
|
scores = z_miss[s:s + chunk] @ cb_t - half_norms
|
|
found[s:s + chunk] = scores.argmax(dim=1).to(torch.int32)
|
|
idx[miss] = found
|
|
return idx
|
|
|
|
|
|
def pack_indices(idx: torch.Tensor, bits: int) -> torch.Tensor:
|
|
"""Pack integer indices (< 2**bits, bits <= 16) into a flat uint8 bitstream."""
|
|
flat = idx.flatten().to(torch.int32)
|
|
pad = (-flat.numel()) % 8
|
|
if pad:
|
|
flat = torch.cat([flat, flat.new_zeros(pad)])
|
|
shifts = torch.arange(bits - 1, -1, -1, device=flat.device, dtype=torch.int32)
|
|
bit_mat = ((flat.unsqueeze(-1) >> shifts) & 1).to(torch.uint8) # (n, bits)
|
|
byte_mat = bit_mat.view(-1, 8)
|
|
weights = torch.tensor([1 << i for i in range(7, -1, -1)], device=flat.device, dtype=torch.uint8)
|
|
return (byte_mat * weights).sum(-1, dtype=torch.uint8)
|
|
|
|
|
|
def unpack_indices(packed: torch.Tensor, bits: int, numel: int) -> torch.Tensor:
|
|
"""Inverse of pack_indices. Returns a flat int32 tensor of length numel."""
|
|
shifts = torch.arange(7, -1, -1, device=packed.device, dtype=torch.uint8)
|
|
bit_mat = ((packed.unsqueeze(-1) >> shifts) & 1).view(-1, bits) # uint8
|
|
idx = None
|
|
# accumulate in <=8-bit chunks so intermediates stay uint8
|
|
for c0 in range(0, bits, 8):
|
|
cw = min(8, bits - c0)
|
|
w = torch.tensor([1 << i for i in range(cw - 1, -1, -1)], device=packed.device, dtype=torch.uint8)
|
|
part = (bit_mat[:, c0:c0 + cw] * w).sum(-1, dtype=torch.uint8).to(torch.int32)
|
|
part = part << (bits - c0 - cw)
|
|
idx = part if idx is None else idx | part
|
|
return idx[:numel]
|
|
|
|
|
|
class OrbitVQQuantizer(OstrisQuantizer):
|
|
"""RPBH rotation + lattice vector codebook + per-group least-squares scales.
|
|
One instance per qtype, shareable across modules."""
|
|
|
|
def __init__(self, bits: int, vec_dim: int, lattice: str, codebook_size: int,
|
|
group_size: int = GROUP_SIZE):
|
|
self.bits = bits
|
|
self.vec_dim = vec_dim
|
|
self.lattice = lattice
|
|
self.codebook_size = codebook_size
|
|
self.group_size = group_size
|
|
self.index_bits = bits * vec_dim
|
|
|
|
def can_quantize(self, module: torch.nn.Linear) -> bool:
|
|
d = module.in_features
|
|
h = hadamard_block_size(d)
|
|
if h < MIN_HADAMARD_BLOCK:
|
|
if d not in _skip_warned:
|
|
_skip_warned.add(d)
|
|
print_acc(
|
|
f"OrbitVQ: skipping linears with in_features={d} "
|
|
f"(power-of-two block {h} is too small for the rotation)"
|
|
)
|
|
return False
|
|
return True
|
|
|
|
def _encode_rotated(self, w_rot: torch.Tensor, g: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""Quantize a rotated weight (m, d) in float32. Returns (packed indices,
|
|
float32 (m, d//g) group scales)."""
|
|
tables = get_vq_tables(self.lattice, self.codebook_size, w_rot.device)
|
|
m, d = w_rot.shape
|
|
u = w_rot.view(m, d // g, g)
|
|
# initial scale: per-group rms, so standardized coords are ~N(0, 1)
|
|
scale = u.norm(dim=-1, keepdim=True) / g ** 0.5 + 1e-12
|
|
idx = None
|
|
for _ in range(LS_REFIT_ROUNDS):
|
|
z = (u / scale).reshape(-1, self.vec_dim)
|
|
idx = encode_vectors(z, tables)
|
|
c = tables.codebook.index_select(0, idx).view(m, d // g, g)
|
|
# least-squares optimal scale given the chosen codewords
|
|
num = (u * c).sum(dim=-1, keepdim=True)
|
|
den = c.square().sum(dim=-1, keepdim=True) + 1e-12
|
|
scale = num / den
|
|
return pack_indices(idx, self.index_bits), scale.view(m, d // g)
|
|
|
|
def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None:
|
|
d = module.in_features
|
|
h = hadamard_block_size(d)
|
|
g = min(self.group_size, h)
|
|
device = weight_fp32.device
|
|
perm_cpu, signs_cpu = rpbh_params(d)
|
|
perm = perm_cpu.to(device=device, dtype=torch.int32)
|
|
inv_perm = torch.argsort(perm_cpu).to(device=device, dtype=torch.int32)
|
|
signs = signs_cpu.to(device)
|
|
w_rot = rpbh_forward(weight_fp32, perm, signs, h)
|
|
packed, scales = self._encode_rotated(w_rot, g)
|
|
module.register_buffer("ovq_packed", packed, persistent=False)
|
|
module.register_buffer("ovq_scales", scales.to(module.weight.dtype), persistent=False)
|
|
module.register_buffer("ovq_perm", perm, persistent=False)
|
|
module.register_buffer("ovq_inv_perm", inv_perm, persistent=False)
|
|
module.register_buffer("ovq_signs", signs, persistent=False)
|
|
module.ovq_block = h
|
|
module.ovq_group = g
|
|
|
|
def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
|
|
"""Materialize the rotated-basis weight W' = W P^T in the given dtype."""
|
|
tables = get_vq_tables(self.lattice, self.codebook_size, module.ovq_packed.device)
|
|
m, d = module.out_features, module.in_features
|
|
g = module.ovq_group
|
|
idx = unpack_indices(module.ovq_packed, self.index_bits, m * d // self.vec_dim)
|
|
w = tables.codebook.to(dtype).index_select(0, idx).view(m, d // g, g)
|
|
w = w * module.ovq_scales.to(dtype).unsqueeze(-1)
|
|
return w.view(m, d)
|
|
|
|
def dequantize(self, module) -> torch.Tensor:
|
|
w = self._dequantize_rotated(module, torch.float32)
|
|
return rpbh_inverse(w, module.ovq_inv_perm, module.ovq_signs, module.ovq_block)
|
|
|
|
def requantize_(self, module, fp_weight: torch.Tensor) -> None:
|
|
w = fp_weight.to(device=module.ovq_packed.device, dtype=torch.float32)
|
|
w_rot = rpbh_forward(w, module.ovq_perm, module.ovq_signs, module.ovq_block)
|
|
packed, scales = self._encode_rotated(w_rot, module.ovq_group)
|
|
module.ovq_packed = packed
|
|
module.ovq_scales = scales.to(module.ostris_orig_dtype)
|
|
|
|
def forward(self, module, x: torch.Tensor) -> torch.Tensor:
|
|
# rotate the activation instead of un-rotating the weight; the rotations
|
|
# cancel in the matmul. the weight is frozen, so build it outside autograd;
|
|
# gradients still flow to x through the rotation and the matmul
|
|
with torch.no_grad():
|
|
w = self._dequantize_rotated(module, x.dtype)
|
|
x_rot = rpbh_forward(x, module.ovq_perm, module.ovq_signs, module.ovq_block)
|
|
return F.linear(x_rot, w, module.bias)
|