| """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) |
| |
| 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] |
|
|
|
|
| |
| 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) |
| 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) |
| 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) |
| return W.t().contiguous() |
|
|
