File size: 1,734 Bytes
8b8e59d | 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 | """Our own Q8 quantization (GGUF Q8_0-style): int8 storage, fp32 compute.
torch's built-in dynamic quantization has no kernels on this ARM build
('unknown architecure'), so we implement the standard scheme used by GGUF
Q8_0: per-channel symmetric int8 weights with fp32 scales. Storage drops to
~1/4 of fp32; compute dequantizes on load, so behavior is near-lossless.
Usage:
qstate = quantize_q8(model) # {key: {"scale": [out], "q": int8}}
load_q8(model, qstate) # dequant into float weights in place
"""
import torch
@torch.no_grad()
def quantize_tensor(t: torch.Tensor) -> dict:
"""Per-channel symmetric int8 quantization of a [out, in] linear weight."""
t = t.float().contiguous()
out_dim = t.shape[0]
amax = t.abs().amax(dim=1, keepdim=True).clamp(min=1e-8)
scale = (amax / 127.0).squeeze(1)
q = torch.round(t / scale.view(-1, 1)).clamp(-127, 127).to(torch.int8)
return {"q": q, "scale": scale}
@torch.no_grad()
def dequantize_tensor(qstate: dict) -> torch.Tensor:
return qstate["q"].float() * qstate["scale"].view(-1, 1)
def quantize_q8(model: torch.nn.Module) -> dict:
qs = {}
for name, mod in model.named_modules():
if isinstance(mod, torch.nn.Linear):
prefix = name + ".weight"
qs[prefix] = quantize_tensor(mod.weight.detach())
return qs
def load_q8(model: torch.nn.Module, qstate: dict):
"""Dequantize Q8 states into the model's float weights (in place)."""
with torch.no_grad():
for name, mod in model.named_modules():
if isinstance(mod, torch.nn.Linear) and name + ".weight" in qstate:
mod.weight.copy_(dequantize_tensor(qstate[name + ".weight"]))
return model
|