code/flashsim/quantization.py

"""Independent reimplementation of FlashOptim's companded 8-bit optimizer-state
quantization (Algorithms 2-3 of arXiv:2602.23349v2), plus the linear
(non-companded) baseline the paper ablates in Figs. 4-5.

Written from the paper's pseudocode only; cross-validated against the authors'
Triton `quantize`/`dequantize` on GPU (see gpu job scripts).

Momentum (Alg. 2, signed int8, group G=32, softsign companding):
    s_g = max |m_g| ;  m' = m/s ;  m'' = 2m'/(1+|m'|) ;  q = round(m''*127)
    inverse: m'' = q/127 ; m' = m''/(2-|m''|) ; m = m' * s
Variance (Alg. 3, unsigned uint8, sqrt companding):
    v' = sqrt(v) ; s_g = max v'_g ; q = round((v'/s)*255)
    inverse: v' = (q/255)*s ; v = v'^2

Scales are stored in fp16 (paper §3.2: FP16 scale per group of 32 -> 2/32
bytes/param overhead).
"""

from __future__ import annotations

import torch

GROUP = 32


def _grouped(x: torch.Tensor) -> tuple[torch.Tensor, int]:
    """Flatten and pad x to a multiple of GROUP; return (groups, orig_numel)."""
    flat = x.reshape(-1)
    n = flat.numel()
    pad = (-n) % GROUP
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=flat.dtype, device=flat.device)])
    return flat.reshape(-1, GROUP), n


def quantize(x: torch.Tensor, signed: bool = True, sqrt: bool = False,
             softsign: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
    """Group-32 absmax INT8/UINT8 quantization with optional companding."""
    g, n = _grouped(x.to(torch.float32))
    if sqrt:
        g = torch.sqrt(g)
    if signed:
        scales = g.abs().amax(dim=1, keepdim=True)
    else:
        scales = g.amax(dim=1, keepdim=True)
    scales_f16 = scales.to(torch.float16)
    safe = torch.where(scales_f16 == 0, torch.ones_like(scales_f16), scales_f16).to(torch.float32)
    y = g / safe
    if softsign:
        y = 2.0 * y / (1.0 + y.abs())
    if signed:
        q = torch.clamp(torch.round(y * 127.0), -127, 127).to(torch.int8)
    else:
        q = torch.clamp(torch.round(y * 255.0), 0, 255).to(torch.uint8)
    return q.reshape(-1)[:n], scales_f16.reshape(-1)


def dequantize(q: torch.Tensor, scales: torch.Tensor, signed: bool = True,
               sqrt: bool = False, softsign: bool = True) -> torch.Tensor:
    g, n = _grouped(q.to(torch.float32))
    s = scales.to(torch.float32).reshape(-1, 1)
    y = g / (127.0 if signed else 255.0)
    if softsign:
        y = y / (2.0 - y.abs())
    out = y * s
    if sqrt:
        out = out * out
    return out.reshape(-1)[:n]


def roundtrip(x: torch.Tensor, signed: bool = True, sqrt: bool = False,
              softsign: bool = True) -> torch.Tensor:
    q, s = quantize(x, signed=signed, sqrt=sqrt, softsign=softsign)
    return dequantize(q, s, signed=signed, sqrt=sqrt, softsign=softsign).reshape(x.shape)


# Named configurations used in experiments
def momentum_companded(x):
    return roundtrip(x, signed=True, sqrt=False, softsign=True)


def momentum_linear(x):
    return roundtrip(x, signed=True, sqrt=False, softsign=False)


def variance_companded(x):
    return roundtrip(x, signed=False, sqrt=True, softsign=False)


def variance_linear(x):
    return roundtrip(x, signed=False, sqrt=False, softsign=False)


def nmse(orig: torch.Tensor, recon: torch.Tensor) -> float:
    """Normalized MSE, as in paper §4.5."""
    denom = float((orig.double() ** 2).mean())
    if denom == 0:
        return 0.0
    return float(((orig.double() - recon.double()) ** 2).mean()) / denom