Spaces:
Running on Zero
Running on Zero
| """ | |
| OrbitQuant weight-only quantization backend (orbit2 / orbit3 / orbit4 qtypes). | |
| Implements the weight half of "OrbitQuant: Data-Agnostic Quantization for Image and | |
| Video Diffusion Transformers" (arXiv:2607.02461) as an OstrisQuantizer backend (see | |
| toolkit/util/ostris_quant.py). Each linear weight is rotated with a randomized | |
| permuted block-Hadamard (RPBH) rotation shared per input dimension, split into | |
| per-row norms and unit directions, and the directions are quantized with a Lloyd-Max | |
| codebook fit to the fixed post-rotation coordinate marginal N(0, 1/d). | |
| No calibration data is needed. | |
| At runtime the forward rotation is applied to the activations instead of un-rotating | |
| the weight, so the two cancel in the matmul (rotate = multiply by P): | |
| W' = W P^T, y = dequant(W') (P x) ~= W x | |
| Activations are not quantized. This is meant for holding a frozen base model at low | |
| bit-width while training adapters on top; activation quantization would only add | |
| noise without fused low-bit kernels. | |
| Quantized state attached to each module: | |
| orbit_packed packed codebook indices (uint8 bitstream) | |
| orbit_row_norms per output row l2 norms of the rotated weight (original dtype) | |
| orbit_codebook Lloyd-Max centroids for N(0, 1/in_features) (float32) | |
| orbit_perm / orbit_inv_perm / orbit_signs shared RPBH rotation for in_features | |
| orbit_bits / orbit_block bit-width and Hadamard block 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 | |
| # qtype name -> bits per weight coordinate | |
| ORBIT_QTYPES = {"orbit2": 2, "orbit3": 3, "orbit4": 4} | |
| # below this Hadamard block size the rotated coordinates are not gaussian enough for | |
| # the shared N(0, 1/d) codebook to be valid, so conversion is skipped for the layer | |
| MIN_HADAMARD_BLOCK = 32 | |
| _normal_codebook_cache: Dict[int, torch.Tensor] = {} | |
| _rotation_cache: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} | |
| _skip_warned = set() | |
| def gaussian_lloyd_max(bits: int, iters: int = 200) -> torch.Tensor: | |
| """MSE-optimal (Lloyd-Max) centroids for the standard normal, 2**bits levels, | |
| returned in ascending order as float32. Cached per bit-width.""" | |
| if bits in _normal_codebook_cache: | |
| return _normal_codebook_cache[bits] | |
| levels = 2 ** bits | |
| # init at the gaussian quantile midpoints | |
| q = (torch.arange(levels, dtype=torch.float64) + 0.5) / levels | |
| c = math.sqrt(2.0) * torch.erfinv(2.0 * q - 1.0) | |
| inf = torch.tensor([math.inf], dtype=torch.float64) | |
| for _ in range(iters): | |
| edges = (c[:-1] + c[1:]) / 2.0 | |
| lo = torch.cat([-inf, edges]) | |
| hi = torch.cat([edges, inf]) | |
| # centroid update: E[X | lo < X < hi] = (phi(lo) - phi(hi)) / (Phi(hi) - Phi(lo)) | |
| phi_lo = torch.exp(-lo * lo / 2.0) / math.sqrt(2.0 * math.pi) | |
| phi_hi = torch.exp(-hi * hi / 2.0) / math.sqrt(2.0 * math.pi) | |
| cdf_lo = 0.5 * (1.0 + torch.erf(lo / math.sqrt(2.0))) | |
| cdf_hi = 0.5 * (1.0 + torch.erf(hi / math.sqrt(2.0))) | |
| c = (phi_lo - phi_hi) / (cdf_hi - cdf_lo) | |
| c = c.to(torch.float32) | |
| _normal_codebook_cache[bits] = c | |
| return c | |
| def hadamard_block_size(d: int) -> int: | |
| # largest power of two dividing d | |
| return d & (-d) | |
| def rpbh_params(d: int) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Permutation (int64) and Rademacher signs (int8) of the RPBH rotation for input | |
| dimension d, on cpu. Sampled once per dimension with a seed derived from d so the | |
| rotation is identical across layers, runs, and resumes.""" | |
| if d in _rotation_cache: | |
| return _rotation_cache[d] | |
| g = torch.Generator().manual_seed(0x0EB17 + d) | |
| perm = torch.randperm(d, generator=g) | |
| signs = torch.randint(0, 2, (d,), generator=g, dtype=torch.int8) * 2 - 1 | |
| _rotation_cache[d] = (perm, signs) | |
| return perm, signs | |
| def _fwht(x: torch.Tensor, h: int) -> torch.Tensor: | |
| """Orthonormal fast Walsh-Hadamard transform applied to each contiguous block of h | |
| coordinates along the last dimension (h must be a power of two dividing the dim).""" | |
| shape = x.shape | |
| x = x.reshape(-1, h) | |
| m = x.shape[0] | |
| step = 1 | |
| while step < h: | |
| y = x.view(m, h // (2 * step), 2, step) | |
| x = torch.stack((y[:, :, 0] + y[:, :, 1], y[:, :, 0] - y[:, :, 1]), dim=2).view(m, h) | |
| step *= 2 | |
| return (x * h ** -0.5).view(shape) | |
| def rpbh_forward(x: torch.Tensor, perm: torch.Tensor, signs: torch.Tensor, h: int) -> torch.Tensor: | |
| """y = blkdiag(H_h D) P x applied to the last dimension of x.""" | |
| y = torch.index_select(x, -1, perm) * signs.to(x.dtype) | |
| return _fwht(y, h) | |
| def rpbh_inverse(y: torch.Tensor, inv_perm: torch.Tensor, signs: torch.Tensor, h: int) -> torch.Tensor: | |
| """Inverse of rpbh_forward (the rotation is orthogonal, H is self-inverse).""" | |
| z = _fwht(y, h) * signs.to(y.dtype) | |
| return torch.index_select(z, -1, inv_perm) | |
| def pack_codes(codes: torch.Tensor, bits: int) -> torch.Tensor: | |
| """Pack integer codes (values < 2**bits, any shape) into a flat uint8 bitstream.""" | |
| flat = codes.flatten().to(torch.uint8) | |
| 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.uint8) | |
| bit_mat = (flat.unsqueeze(-1) >> shifts) & 1 # (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_codes(packed: torch.Tensor, bits: int, numel: int) -> torch.Tensor: | |
| """Inverse of pack_codes. Returns a flat uint8 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) | |
| weights = torch.tensor([1 << i for i in range(bits - 1, -1, -1)], device=packed.device, dtype=torch.uint8) | |
| codes = (bit_mat * weights).sum(-1, dtype=torch.uint8) | |
| return codes[:numel] | |
| def _quantize_rows( | |
| w_fp32: torch.Tensor, | |
| perm: torch.Tensor, | |
| signs: torch.Tensor, | |
| h: int, | |
| codebook: torch.Tensor, | |
| bits: int, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Rotate weight rows into the RPBH basis and quantize their unit directions | |
| against the codebook. Returns (packed codes, float32 row norms).""" | |
| w_rot = rpbh_forward(w_fp32, perm, signs, h) | |
| row_norms = w_rot.norm(dim=1) | |
| unit = w_rot / (row_norms + 1e-10).unsqueeze(1) | |
| edges = (codebook[:-1] + codebook[1:]) / 2 | |
| codes = torch.bucketize(unit, edges, out_int32=True).to(torch.uint8) | |
| return pack_codes(codes, bits), row_norms | |
| class OrbitQuantizer(OstrisQuantizer): | |
| """OrbitQuant backend. One instance per bit-width, shareable across modules.""" | |
| def __init__(self, bits: int): | |
| self.bits = bits | |
| 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"OrbitQuant: skipping linears with in_features={d} " | |
| f"(power-of-two block {h} is too small for the rotation)" | |
| ) | |
| return False | |
| return True | |
| def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None: | |
| d = module.in_features | |
| h = hadamard_block_size(d) | |
| 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) | |
| codebook = (gaussian_lloyd_max(self.bits) * d ** -0.5).to(device) | |
| packed, row_norms = _quantize_rows(weight_fp32, perm, signs, h, codebook, self.bits) | |
| module.register_buffer("orbit_packed", packed, persistent=False) | |
| module.register_buffer("orbit_row_norms", row_norms.to(module.weight.dtype), persistent=False) | |
| # stored as a uint8 byte view so module.to(dtype=...) can't cast the fp32 | |
| # codebook (nn.Module._apply dtype-casts every floating buffer) | |
| module.register_buffer("orbit_codebook", codebook.view(torch.uint8), persistent=False) | |
| module.register_buffer("orbit_perm", perm, persistent=False) | |
| module.register_buffer("orbit_inv_perm", inv_perm, persistent=False) | |
| module.register_buffer("orbit_signs", signs, persistent=False) | |
| module.orbit_bits = self.bits | |
| module.orbit_block = h | |
| def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor: | |
| """Materialize the rotated-basis weight W' = W P^T in the given dtype.""" | |
| numel = module.out_features * module.in_features | |
| codes = unpack_codes(module.orbit_packed, module.orbit_bits, numel) | |
| codebook = module.orbit_codebook.view(torch.float32) | |
| w = torch.index_select(codebook.to(dtype), 0, codes.to(torch.int32)) | |
| w = w.view(module.out_features, module.in_features) | |
| return w * module.orbit_row_norms.to(dtype).unsqueeze(1) | |
| def dequantize(self, module) -> torch.Tensor: | |
| w = self._dequantize_rotated(module, torch.float32) | |
| return rpbh_inverse(w, module.orbit_inv_perm, module.orbit_signs, module.orbit_block) | |
| def requantize_(self, module, fp_weight: torch.Tensor) -> None: | |
| w = fp_weight.to(device=module.orbit_packed.device, dtype=torch.float32) | |
| packed, row_norms = _quantize_rows( | |
| w, module.orbit_perm, module.orbit_signs, module.orbit_block, | |
| module.orbit_codebook.view(torch.float32), module.orbit_bits, | |
| ) | |
| module.orbit_packed = packed | |
| module.orbit_row_norms = row_norms.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.orbit_perm, module.orbit_signs, module.orbit_block) | |
| return F.linear(x_rot, w, module.bias) | |