Spaces:
Paused
Paused
| """Unbiased low-bit activation quantization — the activation-memory lever for the counter update. | |
| The counter update needs grad_w = Delta^T X. It does NOT need X exactly: if X is replaced by an | |
| *unbiased* low-bit quantization Q(X) with E[Q(X) | X] = X, then | |
| E[ Delta^T Q(X) | X, Delta ] = Delta^T X, | |
| so the update stays unbiased (it only gains variance). Storing Q(X) at b bits instead of fp X | |
| shrinks the saved-activation memory the backward needs. This module provides per-row symmetric | |
| stochastic quantization (the unbiasedness comes from stochastic rounding) and an effective-bits | |
| accounting that includes the per-row scale overhead. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| __all__ = ["stochastic_quantize", "quantize_codes", "dequantize_codes", "effective_bits", | |
| "pack_int4", "unpack_int4"] | |
| def pack_int4(codes: torch.Tensor) -> torch.Tensor: | |
| """Pack signed 4-bit codes (values in [-8,7]) two per byte -> uint8, ~half the bytes. | |
| Returns a flat uint8 tensor; the caller keeps the element count to unpack.""" | |
| flat = codes.reshape(-1).to(torch.int64) | |
| nib = (flat + 8).clamp_(0, 15) # signed [-8,7] -> [0,15] | |
| if nib.numel() % 2: | |
| nib = torch.cat([nib, nib.new_zeros(1)]) | |
| lo = nib[0::2] | |
| hi = nib[1::2] | |
| return (lo | (hi << 4)).to(torch.uint8) | |
| def unpack_int4(packed: torch.Tensor, n: int) -> torch.Tensor: | |
| """Inverse of pack_int4. Returns a flat int8 tensor of length n with values in [-8,7].""" | |
| p = packed.to(torch.int64) | |
| lo = (p & 0x0F) - 8 | |
| hi = ((p >> 4) & 0x0F) - 8 | |
| out = torch.stack([lo, hi], dim=1).reshape(-1)[:n] | |
| return out.to(torch.int8) | |
| def _levels(bits: int) -> int: | |
| # signed symmetric: codes in [-levels, levels], levels = 2^(bits-1) - 1 | |
| return (1 << (bits - 1)) - 1 | |
| def quantize_codes(x: torch.Tensor, bits: int, dim: int = -1): | |
| """Per-row symmetric stochastic quantization. Returns (codes int, scale fp) with | |
| codes in [-levels, levels] and E[codes * scale | x] = x. Storing codes (b bits) + one | |
| scale per row is the low-bit saved activation.""" | |
| levels = _levels(bits) | |
| amax = x.abs().amax(dim=dim, keepdim=True).clamp_min(1e-12) | |
| scale = amax / levels | |
| y = x / scale | |
| floor = torch.floor(y) | |
| codes = floor + (torch.rand_like(y) < (y - floor)).to(y.dtype) # stochastic round (unbiased) | |
| codes = codes.clamp_(-levels, levels) | |
| return codes, scale | |
| def dequantize_codes(codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: | |
| return codes * scale | |
| def stochastic_quantize(x: torch.Tensor, bits: int, dim: int = -1) -> torch.Tensor: | |
| """Convenience: round-trip x through unbiased b-bit quantization, returning the dequantized | |
| tensor (same shape/dtype as x). E[stochastic_quantize(x) | x] = x.""" | |
| codes, scale = quantize_codes(x, bits, dim) | |
| return dequantize_codes(codes, scale).to(x.dtype) | |
| def effective_bits(bits: int, row_len: int, scale_bits: int = 16) -> float: | |
| """Bits per element including the amortized per-row scale (one scale_bits value per row).""" | |
| return bits + scale_bits / float(row_len) | |