File size: 2,781 Bytes
3495881 | 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 76 77 78 79 80 81 82 83 84 | """
Adaptive numerical precision as a modular execution feature.
We implement *fake quantisation* (quantise -> dequantise) with a straight-through
estimator, so the FP32 reference and the reduced-precision path share one code
path and can be compared numerically. Precision is assigned per named site
(e.g. 'attn', 'mlp'), which makes it a routable resource rather than a fixed
compile-time property.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
import torch
LEVELS = {"int2": 2, "int3": 3, "int4": 4, "int8": 8, "fp16": 16, "bf16": 16, "fp32": 32}
def quantise(x: torch.Tensor, bits: int, per: str = "tensor") -> torch.Tensor:
"""Symmetric uniform fake-quant with straight-through gradient."""
if bits >= 32:
return x
qmax = 2 ** (bits - 1) - 1
if per == "row":
s = x.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / qmax
else:
s = x.abs().amax().clamp_min(1e-8) / qmax
q = torch.clamp(torch.round(x / s), -qmax - 1, qmax) * s
return x + (q - x).detach()
def cast_float(x: torch.Tensor, kind: str) -> torch.Tensor:
if kind == "fp16":
return x.half().float()
if kind == "bf16":
return x.bfloat16().float()
return x
@dataclass
class QuantPolicy:
"""Maps a computation site -> precision level.
`default` applies to any site not explicitly listed. `sensitive` sites are
kept at higher precision. A policy is data-independent unless
`difficulty_fn` is supplied, in which case precision may vary per position.
"""
default: str = "fp32"
sites: dict = field(default_factory=dict)
per: str = "row"
enabled: bool = True
def level_for(self, site: str) -> str:
return self.sites.get(site, self.default)
def describe(self) -> dict:
return {"default": self.default, "sites": dict(self.sites), "per": self.per}
def maybe_quant(x: torch.Tensor, qp: Optional[QuantPolicy], site: str) -> torch.Tensor:
if qp is None or not qp.enabled:
return x
lvl = qp.level_for(site)
if lvl in ("fp16", "bf16"):
return cast_float(x, lvl)
if lvl == "fp32":
return x
return quantise(x, LEVELS[lvl], qp.per)
def quantise_model_weights(model, bits: int, skip=("norm", "mem0", "step_emb")) -> dict:
"""In-place fake-quantise weights. Returns per-tensor relative error."""
err = {}
with torch.no_grad():
for name, p in model.named_parameters():
if any(s in name for s in skip) or p.dim() < 2:
continue
ref = p.detach().clone()
q = quantise(p.data, bits, "row")
p.data.copy_(q)
err[name] = float((q - ref).norm() / ref.norm().clamp_min(1e-9))
return err
|