206 lines
8.8 KiB
Python
206 lines
8.8 KiB
Python
"""
|
|
Bit-exact reimplementation of torchao 0.10.0's UIntXWeightOnlyConfig weight
|
|
quantization as an OstrisQuantizer backend, so the uint2..uint7 qtypes keep
|
|
producing byte-identical weights after torchao drops uintx support. Existing
|
|
accuracy recovery adapters were trained against these exact quantized bases,
|
|
so every op below mirrors torchao's sequence (same order, same dtypes):
|
|
|
|
choose_qparams_affine (ASYMMETRIC, block_size (1, 64), preserve_zero=True,
|
|
eps=float32 eps, INT zero-point domain, scale in the weight dtype):
|
|
min/max per 64-wide group along in_features, extended to include 0
|
|
scale = (max_val_pos - min_val_neg) / (qmax - qmin), clamped to eps
|
|
zero_point = clamp(qmin - round(min_val_neg / scale), qmin, qmax) as int
|
|
quantize_affine:
|
|
q = clamp(round(w * (1.0 / scale)) + zero_point, qmin, qmax)
|
|
dequantize_affine:
|
|
w = ((q as int) - zero_point) cast to the weight dtype, times scale
|
|
|
|
The arithmetic is done in the original weight dtype (usually bfloat16) because
|
|
that is what torchao did; doing it in float32 would round differently.
|
|
|
|
uint8 resolves to a backend too, but can_quantize always refuses it: torchao
|
|
0.10.0's uint8 path raised inside UintxTensor.from_uint8 (packing only supports
|
|
1..7 bits) before the module was touched, so "uint8" layers have always been
|
|
silently left unquantized. Refusing keeps that exact behavior — flip the check
|
|
in can_quantize if real uint8 quantization is ever wanted for new models.
|
|
|
|
Buffers on the module (registered by quantize_):
|
|
uintx_packed quantized codes packed into power-of-2 bit shards (like
|
|
torchao's UintxTensor: e.g. uint3 = a 2-bit + a 1-bit
|
|
shard), concatenated into one flat uint8 buffer. Shards
|
|
unpack with a couple of elementwise shift/mask kernels,
|
|
which is much cheaper per forward than a generic bitstream.
|
|
uintx_scale per-group scale, stored as a uint8 byte view of the weight
|
|
dtype so module.to(dtype=...) can't cast it
|
|
uintx_zero_point per-group zero point, uint8 (values are in [0, qmax])
|
|
"""
|
|
|
|
import torch
|
|
|
|
from toolkit.util.ostris_quant import OstrisLinear, OstrisQuantizer
|
|
|
|
UINTX_QTYPES = {f"uint{bits}": bits for bits in range(2, 9)}
|
|
|
|
_EPS = torch.finfo(torch.float32).eps
|
|
|
|
|
|
def _pack_shard(vals: torch.Tensor, k: int) -> torch.Tensor:
|
|
"""Pack flat uint8 values (< 2**k) into bytes, 8 // k values per byte.
|
|
Values are laid out in 8//k contiguous chunks (chunk j holds bits
|
|
[j*k, j*k+k) of every byte) so pack and unpack touch memory coalesced."""
|
|
vpb = 8 // k
|
|
if vpb == 1:
|
|
return vals.clone()
|
|
pad = (-vals.numel()) % vpb
|
|
if pad:
|
|
vals = torch.cat([vals, vals.new_zeros(pad)])
|
|
chunks = vals.view(vpb, -1)
|
|
out = chunks[0].clone()
|
|
for j in range(1, vpb):
|
|
out |= chunks[j] << (j * k)
|
|
return out
|
|
|
|
|
|
def _unpack_shard(packed: torch.Tensor, k: int, numel: int) -> torch.Tensor:
|
|
vpb = 8 // k
|
|
if vpb == 1:
|
|
return packed[:numel]
|
|
out = torch.empty(vpb, packed.numel(), dtype=torch.uint8, device=packed.device)
|
|
for j in range(vpb):
|
|
torch.bitwise_right_shift(packed, j * k, out=out[j])
|
|
out.bitwise_and_((1 << k) - 1)
|
|
return out.view(-1)[:numel]
|
|
|
|
|
|
def pack_uintx(codes: torch.Tensor, nbits: int) -> torch.Tensor:
|
|
"""Pack integer codes (values < 2**nbits) into concatenated power-of-2 bit
|
|
shards; bits [offset, offset+k) of each code land in the k-bit shard."""
|
|
flat = codes.flatten().to(torch.uint8)
|
|
shards = []
|
|
offset = 0
|
|
for k in (8, 4, 2, 1):
|
|
if nbits & k:
|
|
if k == nbits: # single shard, values are already < 2**k
|
|
shards.append(_pack_shard(flat, k))
|
|
else:
|
|
shards.append(_pack_shard((flat >> offset) & ((1 << k) - 1), k))
|
|
offset += k
|
|
return torch.cat(shards) if len(shards) > 1 else shards[0]
|
|
|
|
|
|
def unpack_uintx(packed: torch.Tensor, nbits: int, numel: int) -> torch.Tensor:
|
|
"""Inverse of pack_uintx. Returns a flat uint8 tensor of length numel."""
|
|
out = None
|
|
offset = 0
|
|
pos = 0
|
|
for k in (8, 4, 2, 1):
|
|
if nbits & k:
|
|
vpb = 8 // k
|
|
nbytes = -(-numel // vpb)
|
|
vals = _unpack_shard(packed[pos : pos + nbytes], k, numel)
|
|
if out is None:
|
|
# a multi-shard first value is always a fresh tensor from
|
|
# _unpack_shard, safe to mutate; a lone 8-bit shard is a view
|
|
# of the buffer but is never combined, so read-only is fine
|
|
out = vals
|
|
else:
|
|
out |= vals.bitwise_left_shift_(offset)
|
|
offset += k
|
|
pos += nbytes
|
|
return out
|
|
|
|
|
|
class UIntXQuantizer(OstrisQuantizer):
|
|
# quantization runs in the weight's own dtype (that is what torchao did);
|
|
# skipping the float32 copy keeps peak memory at torchao levels
|
|
wants_fp32_weight = False
|
|
|
|
def __init__(self, nbits: int, group_size: int = 64):
|
|
self.nbits = nbits
|
|
self.group_size = group_size
|
|
self.qmin = 0
|
|
self.qmax = (1 << nbits) - 1
|
|
|
|
def can_quantize(self, module: torch.nn.Linear) -> bool:
|
|
if self.nbits == 8:
|
|
# see module docstring: uint8 has always been a silent no-op
|
|
return False
|
|
weight = getattr(module, "weight", None)
|
|
if weight is None or weight.dim() != 2:
|
|
return False
|
|
# torchao asserted on non-divisible in_features, which the quantize loop
|
|
# caught, leaving the layer unquantized; refuse for the same end state
|
|
return weight.shape[1] % self.group_size == 0
|
|
|
|
@torch.no_grad()
|
|
def quantize_(self, module: torch.nn.Linear, weight: torch.Tensor) -> None:
|
|
# weight arrives in its original dtype (wants_fp32_weight is False)
|
|
self._quantize_impl(module, weight)
|
|
|
|
@torch.no_grad()
|
|
def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None:
|
|
# the torchao path cast the merged weight to the original dtype and
|
|
# re-quantized in that dtype
|
|
self._quantize_impl(module, fp_weight.to(module.ostris_orig_dtype))
|
|
|
|
def _quantize_impl(self, module: torch.nn.Module, w: torch.Tensor) -> None:
|
|
out_f, in_f = w.shape
|
|
groups = in_f // self.group_size
|
|
wv = w.contiguous().view(out_f, groups, self.group_size)
|
|
|
|
min_val = torch.amin(wv, dim=2)
|
|
max_val = torch.amax(wv, dim=2)
|
|
# preserve_zero: the qparams must be able to represent 0.0 exactly
|
|
min_val_neg = torch.min(min_val, torch.zeros_like(min_val))
|
|
max_val_pos = torch.max(max_val, torch.zeros_like(max_val))
|
|
|
|
scale = (max_val_pos - min_val_neg) / float(self.qmax - self.qmin)
|
|
scale = torch.clamp(scale, min=_EPS)
|
|
zero_point = self.qmin - torch.round(min_val_neg / scale)
|
|
zero_point = torch.clamp(zero_point, self.qmin, self.qmax).to(torch.int32)
|
|
|
|
q = torch.clamp(
|
|
torch.round(wv * (1.0 / scale.view(out_f, groups, 1)))
|
|
+ zero_point.view(out_f, groups, 1),
|
|
self.qmin,
|
|
self.qmax,
|
|
)
|
|
q = q.view(out_f, in_f).to(torch.uint8)
|
|
|
|
module.register_buffer(
|
|
"uintx_packed", pack_uintx(q, self.nbits), persistent=False
|
|
)
|
|
module.register_buffer(
|
|
"uintx_scale", scale.contiguous().view(torch.uint8), persistent=False
|
|
)
|
|
module.register_buffer(
|
|
"uintx_zero_point", zero_point.to(torch.uint8), persistent=False
|
|
)
|
|
|
|
def _dequantize_native(self, module: "OstrisLinear") -> torch.Tensor:
|
|
"""Reconstruct the weight in its original dtype. torchao subtracted the
|
|
zero point in int32 then cast; doing it directly in the weight dtype is
|
|
bit-identical (codes and zero points are integers <= 255, and every
|
|
integer in [-255, 255] is exact in bf16/fp16/fp32) with half the
|
|
intermediate memory traffic."""
|
|
out_f, in_f = module.out_features, module.in_features
|
|
groups = in_f // self.group_size
|
|
scale = module.uintx_scale.view(module.ostris_orig_dtype)
|
|
q = unpack_uintx(module.uintx_packed, self.nbits, out_f * in_f)
|
|
dq = q.view(out_f, groups, self.group_size).to(scale.dtype)
|
|
dq -= module.uintx_zero_point.to(scale.dtype).view(out_f, groups, 1)
|
|
dq *= scale.view(out_f, groups, 1)
|
|
return dq.view(out_f, in_f)
|
|
|
|
def dequantize(self, module: "OstrisLinear") -> torch.Tensor:
|
|
return self._dequantize_native(module).to(torch.float32)
|
|
|
|
def forward(self, module: "OstrisLinear", x: torch.Tensor) -> torch.Tensor:
|
|
# skip the float32 round-trip of the base implementation; the weight is
|
|
# frozen, so build it outside autograd
|
|
with torch.no_grad():
|
|
w = self._dequantize_native(module)
|
|
if w.dtype != x.dtype:
|
|
w = w.to(x.dtype)
|
|
return torch.nn.functional.linear(x, w, module.bias)
|