Spaces:
Paused
Paused
File size: 3,096 Bytes
6f75f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """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)
|