File size: 2,593 Bytes
515b383 | 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 | """FP8 (e4m3) linear layers for Hopper via torch._scaled_mm.
Weights: per-output-channel scales, quantised once. Activations: per-token dynamic scales.
Under torch.compile the quantisation ops fuse into the surrounding elementwise work."""
import torch, torch.nn as nn
E4M3_MAX = 448.0
def _quant_rowwise(x):
s = x.abs().amax(dim=-1, keepdim=True).float().clamp(min=1e-12) / E4M3_MAX
return (x.float() / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn), s
class FP8Linear(nn.Module):
def __init__(self, lin: nn.Linear):
super().__init__()
wq, sw = _quant_rowwise(lin.weight.detach()) # [N,K] fp8, [N,1]
self.register_buffer("wq", wq.contiguous()) # [N,K]; passed as wq.t() -> [K,N] column-major, as _scaled_mm wants
self.register_buffer("sw_t", sw.t().contiguous()) # [1,N]
self.bias = None if lin.bias is None else nn.Parameter(lin.bias.detach().clone(), requires_grad=False)
self.in_features, self.out_features = lin.in_features, lin.out_features
self.out_dtype = lin.weight.dtype
def forward(self, x):
shp = x.shape[:-1]
x2 = x.reshape(-1, self.in_features)
xq, sx = _quant_rowwise(x2)
y = torch._scaled_mm(xq, self.wq.t(), scale_a=sx, scale_b=self.sw_t, bias=self.bias, out_dtype=self.out_dtype)
return y.reshape(*shp, self.out_features)
def convert_to_fp8(model, skip=("lm_head",), min_dim=1024):
"""Replace nn.Linear (with in/out >= min_dim) by FP8Linear in place. Returns count."""
n = 0
for name, mod in list(model.named_modules()):
for cname, child in list(mod.named_children()):
full = f"{name}.{cname}" if name else cname
if isinstance(child, nn.Linear) and not any(s in full for s in skip) and min(child.in_features, child.out_features) >= min_dim:
setattr(mod, cname, FP8Linear(child)); n += 1
return n
if __name__ == "__main__":
import time
lin = nn.Linear(2048, 6144, bias=False).cuda().to(torch.bfloat16)
f8 = FP8Linear(lin)
x = torch.randn(8192, 2048, device="cuda", dtype=torch.bfloat16)
ref = lin(x); got = f8(x)
print("rel err", ((ref.float() - got.float()).abs().mean() / ref.float().abs().mean()).item())
for f, name in [(lin, "bf16 linear"), (f8, "fp8 linear (eager)"), (torch.compile(f8), "fp8 linear (compiled)")]:
for _ in range(3): f(x)
torch.cuda.synchronize(); t = time.time()
for _ in range(20): f(x)
torch.cuda.synchronize(); print(f"{name:24s} {(time.time()-t)/20*1000:.3f} ms")
|