File size: 8,900 Bytes
7ccb33d | 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """True int8/int4/int3 packing for canonical RTN fake-quant weights."""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn.functional as F
from .rtn_quant import _validate_rtn_args
_DTYPE_NAMES = {
torch.float16: "float16",
torch.bfloat16: "bfloat16",
torch.float32: "float32",
torch.float64: "float64",
}
_NAME_DTYPES = {name: dtype for dtype, name in _DTYPE_NAMES.items()}
@dataclass(frozen=True)
class PackedRTN:
"""Integer payload, fp16 scales, and the metadata needed for decoding."""
data: torch.Tensor
scales: torch.Tensor
shape: tuple[int, int]
bits: int
granularity: str
group_size: int
original_dtype: str
@property
def num_weights(self) -> int:
return self.shape[0] * self.shape[1]
@property
def effective_bits(self) -> float:
"""Actual bits/weight including fp16 scales and payload byte padding."""
return (self.data.numel() * 8 + self.scales.numel() * 16) / self.num_weights
def to_bytes(self) -> bytes:
if self.data.dtype == torch.int8:
return self.data.cpu().numpy().tobytes()
return bytes(self.data.cpu().tolist())
def _canonical_codes_and_scales(
weight: torch.Tensor,
bits: int,
granularity: str,
group_size: int,
) -> tuple[torch.Tensor, torch.Tensor]:
_validate_rtn_args(weight, bits, granularity, group_size)
if weight.dtype not in _DTYPE_NAMES:
raise ValueError(f"unsupported fake-quant dtype: {weight.dtype}")
cpu_weight = weight.detach().cpu().contiguous()
rows, columns = cpu_weight.shape
actual_group_size = columns if granularity == "per_channel" else group_size
number_of_groups = (columns + actual_group_size - 1) // actual_group_size
padded_columns = number_of_groups * actual_group_size
padded = F.pad(cpu_weight, (0, padded_columns - columns))
groups = padded.reshape(rows, number_of_groups, actual_group_size)
qmax = 2 ** (bits - 1) - 1
scales = (groups.float().abs().amax(dim=-1) / qmax).to(torch.float16)
safe_scales = torch.where(scales == 0, torch.ones_like(scales), scales).float()
codes = torch.round(groups.float() / safe_scales[..., None]).clamp(-qmax, qmax)
reconstructed = (codes * scales.float()[..., None]).reshape(rows, -1)[:, :columns]
reconstructed = reconstructed.to(cpu_weight.dtype)
if not torch.equal(reconstructed, cpu_weight):
raise ValueError(
"w_fakequant is not exactly representable by RTN integer codes and fp16 "
"scales; fake-quantize an fp16 tensor before packing"
)
return codes.to(torch.int8).reshape(-1), scales.contiguous()
def _pack_w3(codes: torch.Tensor) -> torch.Tensor:
"""Pack signed W3 codes little-endian, eight codes per three bytes."""
logical_codes = codes.numel()
if logical_codes % 8:
codes = F.pad(codes, (0, 8 - logical_codes % 8))
values = (codes.to(torch.int16) & 0x7).reshape(-1, 8)
payload = torch.empty((values.shape[0], 3), dtype=torch.int16)
payload[:, 0] = values[:, 0] | (values[:, 1] << 3) | (values[:, 2] << 6)
payload[:, 1] = (
(values[:, 2] >> 2)
| (values[:, 3] << 1)
| (values[:, 4] << 4)
| (values[:, 5] << 7)
)
payload[:, 2] = (values[:, 5] >> 1) | (values[:, 6] << 2) | (values[:, 7] << 5)
payload_bytes = (logical_codes * 3 + 7) // 8
return payload.reshape(-1)[:payload_bytes].to(torch.uint8).contiguous()
def _unpack_w3(payload: torch.Tensor, logical_codes: int) -> torch.Tensor:
"""Unpack the little-endian W3 stream to signed two's-complement codes."""
padding = (-payload.numel()) % 3
if padding:
payload = F.pad(payload, (0, padding))
packed = payload.to(torch.int16).reshape(-1, 3)
values = torch.empty((packed.shape[0], 8), dtype=torch.int16)
values[:, 0] = packed[:, 0] & 0x7
values[:, 1] = (packed[:, 0] >> 3) & 0x7
values[:, 2] = ((packed[:, 0] >> 6) | (packed[:, 1] << 2)) & 0x7
values[:, 3] = (packed[:, 1] >> 1) & 0x7
values[:, 4] = (packed[:, 1] >> 4) & 0x7
values[:, 5] = ((packed[:, 1] >> 7) | (packed[:, 2] << 1)) & 0x7
values[:, 6] = (packed[:, 2] >> 2) & 0x7
values[:, 7] = (packed[:, 2] >> 5) & 0x7
values = values.reshape(-1)[:logical_codes]
return torch.where(values >= 4, values - 8, values)
def pack_rtn(
w_fakequant: torch.Tensor,
bits: int = 8,
granularity: str = "per_channel",
group_size: int = 128,
) -> PackedRTN:
"""Pack canonical RTN values into int8, int4, or dense int3 payloads."""
codes, scales = _canonical_codes_and_scales(
w_fakequant, bits, granularity, group_size
)
rows, columns = w_fakequant.shape
actual_group_size = columns if granularity == "per_channel" else group_size
number_of_groups = (columns + actual_group_size - 1) // actual_group_size
padded_columns = number_of_groups * actual_group_size
logical_codes = codes.reshape(rows, padded_columns)[:, :columns].reshape(-1)
if bits == 8:
data = logical_codes.contiguous()
elif bits == 4:
if logical_codes.numel() % 2:
logical_codes = F.pad(logical_codes, (0, 1))
nibbles = logical_codes.to(torch.int16) & 0xF
data = (nibbles[0::2] | (nibbles[1::2] << 4)).to(torch.uint8)
else:
data = _pack_w3(logical_codes)
return PackedRTN(
data=data,
scales=scales,
shape=(w_fakequant.shape[0], w_fakequant.shape[1]),
bits=bits,
granularity=granularity,
group_size=group_size,
original_dtype=_DTYPE_NAMES[w_fakequant.dtype],
)
def unpack_rtn(packed: PackedRTN) -> torch.Tensor:
"""Decode a packed RTN tensor exactly to its canonical fake-quant values."""
rows, columns = packed.shape
if packed.bits not in {3, 4, 8}:
raise ValueError("packed bits must be 3, 4, or 8")
if packed.granularity not in {"per_channel", "per_group"}:
raise ValueError("invalid packed granularity")
if packed.group_size <= 0 or rows <= 0 or columns <= 0:
raise ValueError("invalid packed shape or group size")
if packed.scales.dtype != torch.float16 or packed.scales.ndim != 2:
raise ValueError("packed scales must be a two-dimensional float16 tensor")
if packed.original_dtype not in _NAME_DTYPES:
raise ValueError(f"unsupported original dtype metadata: {packed.original_dtype}")
actual_group_size = columns if packed.granularity == "per_channel" else packed.group_size
number_of_groups = (columns + actual_group_size - 1) // actual_group_size
expected_scale_shape = (rows, number_of_groups)
if tuple(packed.scales.shape) != expected_scale_shape:
raise ValueError("scale shape does not match packed metadata")
logical_weights = rows * columns
if packed.bits == 8:
if packed.data.dtype != torch.int8 or packed.data.ndim != 1:
raise ValueError("W8 packed data must be a one-dimensional int8 tensor")
if packed.data.numel() != logical_weights:
raise ValueError("W8 payload size does not match packed shape")
logical_codes = packed.data.cpu()
elif packed.bits == 4:
if packed.data.dtype != torch.uint8 or packed.data.ndim != 1:
raise ValueError("W4 packed data must be a one-dimensional uint8 tensor")
if packed.data.numel() != (logical_weights + 1) // 2:
raise ValueError("W4 payload size does not match packed shape")
payload = packed.data.cpu()
nibbles = torch.empty(payload.numel() * 2, dtype=torch.int16)
nibbles[0::2] = payload.to(torch.int16) & 0xF
nibbles[1::2] = payload.to(torch.int16) >> 4
logical_codes = torch.where(nibbles >= 8, nibbles - 16, nibbles)[:logical_weights]
else:
if packed.data.dtype != torch.uint8 or packed.data.ndim != 1:
raise ValueError("W3 packed data must be a one-dimensional uint8 tensor")
if packed.data.numel() != (logical_weights * 3 + 7) // 8:
raise ValueError("W3 payload size does not match packed shape")
logical_codes = _unpack_w3(packed.data.cpu(), logical_weights)
# Packing omits row-tail padding codes; restore it before applying group scales.
padded_columns = number_of_groups * actual_group_size
codes_by_row = logical_codes.reshape(rows, columns)
codes_by_row = F.pad(codes_by_row, (0, padded_columns - columns))
groups = codes_by_row.reshape(rows, number_of_groups, actual_group_size).float()
decoded = groups * packed.scales.cpu().float()[..., None]
return decoded.reshape(rows, padded_columns)[:, :columns].to(
_NAME_DTYPES[packed.original_dtype]
)
# Module-local conventional names are convenient without shadowing STQ1 exports
# from quantlib's package root.
pack = pack_rtn
unpack = unpack_rtn
|