File size: 5,701 Bytes
c675b8d | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Quantizer core (torch, CPU or GPU).
Grid: standard asymmetric GPTQ grid per (output row, input group): range includes 0,
scale = (max - min) / (2^b - 1), integer zero-point z = round(-min / scale), q = clamp(round(w/scale) + z, 0, 2^b - 1),
w_hat = (q - z) * scale. Grids are fixed from the original weights ("static groups").
Zero-points are kept >= 1 so the checkpoint can use the GPTQ v1 storage convention (stores z - 1).
Rounding: sequential error feedback that minimises tr(R H R^T) for a design matrix H (in x in), i.e. the
OBQ/GPTQ recursion with lazy block updates. H = None gives round-to-nearest. With H = I it is identical to RTN.
"""
import torch
def group_params(W, bits, group):
"""W: (out, in). Returns scale, zero of shape (out, in // group)."""
out, inn = W.shape
assert inn % group == 0, f"input dim {inn} not divisible by group size {group}"
maxq = 2 ** bits - 1
Wg = W.reshape(out, inn // group, group)
xmin = torch.clamp(Wg.amin(-1), max=0.0)
xmax = torch.clamp(Wg.amax(-1), min=0.0)
both0 = (xmin == 0) & (xmax == 0)
xmin, xmax = torch.where(both0, -torch.ones_like(xmin), xmin), torch.where(both0, torch.ones_like(xmax), xmax)
scale = (xmax - xmin) / maxq
zero = torch.round(-xmin / scale)
# keep z >= 1 (GPTQ v1 storage): if the group has no negative part, use z = 1 and rescale so max still fits
low = zero < 1
scale = torch.where(low, xmax / (maxq - 1), scale)
zero = torch.where(low, torch.ones_like(zero), zero)
zero = torch.clamp(zero, 1, maxq)
return scale, zero
def _q(w, scale, zero, maxq):
return torch.clamp(torch.round(w / scale) + zero, 0, maxq)
def quantize(W, bits, group, H=None, block=128, dynamic_groups=True):
"""W: (out, in) float tensor (torch Linear layout). Returns (q uint8 (out, in), scale (out, G), zero (out, G)).
H: (in, in) positive-definite design matrix, or None for round-to-nearest.
dynamic_groups: each group's grid is fitted to the error-fed (updated) weights when the group starts,
as in standard GPTQ, so feedback never pushes weights outside their grid. False = grids from the original W."""
W = W.to(torch.float32)
out, inn = W.shape
maxq = 2 ** bits - 1
scale, zero = group_params(W, bits, group)
gid = torch.arange(inn, device=W.device) // group
if H is None:
Q = _q(W, scale[:, gid], zero[:, gid], maxq)
return Q.to(torch.uint8), scale, zero
W = W.clone()
H = H.to(device=W.device, dtype=torch.float32)
L = torch.linalg.cholesky(H)
Hinv = torch.cholesky_inverse(L); del L
U = torch.linalg.cholesky(Hinv, upper=True); del Hinv
Q = torch.zeros((out, inn), dtype=torch.uint8, device=W.device)
for i1 in range(0, inn, block):
i2 = min(i1 + block, inn); n = i2 - i1
W1 = W[:, i1:i2].clone(); E1 = torch.zeros_like(W1); U1 = U[i1:i2, i1:i2]
for i in range(n):
c = i1 + i; g = c // group
if dynamic_groups and c % group == 0:
seg = W1[:, i:i + group] if i + group <= n else torch.cat([W1[:, i:], W[:, i2:c + group]], 1)
sg, zg = group_params(seg, bits, group)
scale[:, g], zero[:, g] = sg[:, 0], zg[:, 0]
w = W1[:, i]
q = _q(w, scale[:, g], zero[:, g], maxq)
Q[:, c] = q.to(torch.uint8)
err = (w - (q - zero[:, g]) * scale[:, g]) / U1[i, i]
W1[:, i:] -= err.unsqueeze(1) * U1[i, i:].unsqueeze(0)
E1[:, i] = err
if i2 < inn:
W[:, i2:] -= E1 @ U[i1:i2, i2:]
return Q, scale, zero
def dequantize(Q, scale, zero, group):
gid = torch.arange(Q.shape[1], device=Q.device) // group
return (Q.to(torch.float32) - zero[:, gid]) * scale[:, gid]
# --------------------------------------------------------------- GPTQ v1 packing (AutoGPTQ / GPTQModel "gptq" format)
def pack(Q, scale, zero, bits, group):
"""Returns dict(qweight int32 (in*bits/32, out), qzeros int32 (G, out*bits/32), scales fp16 (G, out), g_idx int32 (in,))."""
assert bits in (2, 4, 8)
per = 32 // bits
out, inn = Q.shape
q = Q.t().contiguous().to(torch.int64) # (in, out)
assert inn % per == 0 and out % per == 0
qw = torch.zeros((inn // per, out), dtype=torch.int64, device=Q.device)
for k in range(per):
qw |= q[k::per, :] << (bits * k)
z = (zero.t().contiguous().to(torch.int64) - 1) # v1 convention: store z - 1
qz = torch.zeros((z.shape[0], out // per), dtype=torch.int64, device=Q.device)
for k in range(per):
qz |= z[:, k::per] << (bits * k)
to32 = lambda t: (t & 0xFFFFFFFF).to(torch.int64).where(t < 2 ** 31, t - 2 ** 32).to(torch.int32)
return {"qweight": to32(qw), "qzeros": to32(qz), "scales": scale.t().contiguous().to(torch.float16),
"g_idx": (torch.arange(inn, device=Q.device) // group).to(torch.int32)}
def unpack(qweight, qzeros, scales, g_idx, bits):
"""Inverse of pack: dequantized weight in torch Linear layout (out, in), float32."""
per = 32 // bits; mask = 2 ** bits - 1
qw = qweight.to(torch.int64) & 0xFFFFFFFF
inn = qw.shape[0] * per
q = torch.empty((inn, qw.shape[1]), dtype=torch.int64, device=qw.device)
for k in range(per):
q[k::per] = (qw >> (bits * k)) & mask
qz = qzeros.to(torch.int64) & 0xFFFFFFFF
z = torch.empty((qz.shape[0], qz.shape[1] * per), dtype=torch.int64, device=qz.device)
for k in range(per):
z[:, k::per] = ((qz >> (bits * k)) & mask) + 1
g = g_idx.to(torch.int64)
W = (q.to(torch.float32) - z[g].to(torch.float32)) * scales[g].to(torch.float32) # (in, out)
return W.t().contiguous()
|