Spaces:
Running on Zero
Running on Zero
| """Native post-training int8 runtime for the packed LSCodec / FocalCodec model repos. | |
| Packed format (what `export_lscodec_int8.py` / `export_quantized.py` write): every weight as | |
| int8 levels under its original state-dict key plus an fp32 scale under `_scale/<key>`. | |
| `apply_int8(model, packed)` swaps each Linear / Conv1d / ConvTranspose1d / Embedding for a twin | |
| that keeps the weight int8-resident and computes int8 x int8 -> int32 GEMMs (`torch._int_mm`) | |
| with per-tensor dynamic activation quantization (W8A8 dynamic). Biases, norms and other | |
| activation-side constants are materialized as fp32 at load (KB-scale); no weight matrix is ever | |
| dequantized to run it. fp4/fp8/fp16 have no int8 GEMM path — fake-quantize those instead. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| SCALE_PREFIX = "_scale/" | |
| _ROUND = 32 # cublasLt int8 GEMM (via torch._int_mm) wants %32 dims on CUDA; %16 is not enough (K=16 fails) | |
| def _up(n: int) -> int: | |
| return (n + _ROUND - 1) // _ROUND * _ROUND | |
| def _act_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Per-row (last-dim) symmetric absmax activation quantization -> (int8 levels, scales). | |
| One scale per GEMM row (token / output position): a single per-tensor absmax lets a few | |
| outlier channels set the step for everything, which LSCodec's GroupNorms amplify badly. | |
| """ | |
| scale = (x.abs().amax(dim=-1, keepdim=True) / 127.0).clamp_min(1e-12) | |
| return (x / scale).round().clamp(-128, 127).to(torch.int8), scale | |
| def _int_mm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: | |
| """(M, K) int8 x (K, N) int8 -> (M, N) int32. M is padded: CUDA rejects M <= 16.""" | |
| m = a.shape[0] | |
| mp = max(_ROUND, _up(m)) | |
| if mp != m: | |
| padded = torch.zeros((mp, a.shape[1]), dtype=a.dtype, device=a.device) | |
| padded[:m] = a | |
| a = padded | |
| try: | |
| return torch._int_mm(a, b)[:m] | |
| except (RuntimeError, NotImplementedError): # backend refuses the shape | |
| if a.is_cuda: | |
| # exact in float32 (|int8 * int8 * K| << 2^24); int32 matmul has no CUDA kernel | |
| return (a.to(torch.float32) @ b.to(torch.float32))[:m] | |
| return (a.to(torch.int32) @ b.to(torch.int32))[:m] | |
| def _k_pad(a: torch.Tensor, k: int) -> torch.Tensor: | |
| """Zero-pad the K dim; the weight's K is padded the same way, so the product is exact.""" | |
| if a.shape[-1] == k: | |
| return a | |
| padded = torch.zeros((*a.shape[:-1], k), dtype=a.dtype, device=a.device) | |
| padded[..., : a.shape[-1]] = a | |
| return padded | |
| def _quant_channel(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| """fp32 -> int8 + per-output-channel absmax scale (same grid as core.int_levels).""" | |
| scale = w.abs().amax(dim=tuple(range(1, w.dim())), keepdim=True) if w.dim() > 1 else w.abs().amax() | |
| scale = (scale / 128.0).clamp_min(1e-12) | |
| return (w / scale).round().clamp(-128, 127).to(torch.int8), scale.float() | |
| def _quant_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| """fp32 -> int8 + one per-tensor scale (used where per-channel scales don't fit the GEMM).""" | |
| scale = (w.abs().amax() / 128.0).clamp_min(1e-12) | |
| return (w / scale).round().clamp(-128, 127).to(torch.int8), scale.float() | |
| def _dequant(packed: dict, key: str, device: torch.device) -> torch.Tensor: | |
| value = packed[key].to(device) | |
| scale = packed.get(SCALE_PREFIX + key) | |
| return value if scale is None else value.float() * scale.to(device) | |
| def _int1(v) -> int: | |
| if isinstance(v, int): | |
| return v | |
| assert len(v) == 1, f"unsupported tuple attribute {v}" | |
| return v[0] | |
| def _pair(v) -> tuple[int, int]: | |
| if isinstance(v, int): | |
| return (v, v) | |
| if len(v) == 1: | |
| return (v[0], v[0]) | |
| assert len(v) == 2, f"unsupported padding {v}" | |
| return (v[0], v[1]) | |
| class QuantLinear(nn.Module): | |
| def __init__(self, mod: nn.Linear, w: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None): | |
| super().__init__() | |
| n, k = w.shape | |
| wt = torch.zeros((_up(n), _up(k)), dtype=torch.int8, device=w.device) | |
| wt[:n, :k] = w | |
| self.register_buffer("wT", wt.t().contiguous()) # (Kp, Np), the _int_mm operand | |
| sc = torch.zeros((_up(n),), dtype=torch.float32, device=w.device) | |
| sc[:n] = scale.reshape(-1) | |
| self.register_buffer("w_scale", sc) | |
| self.register_buffer("bias", bias) | |
| self.k, self.n = k, n | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| assert x.shape[-1] == self.k, f"expected in_features {self.k}, got {tuple(x.shape)}" | |
| q, sx = _act_quant(x) | |
| a = _k_pad(q.reshape(-1, self.k), self.wT.shape[0]) | |
| y = _int_mm(a, self.wT)[:, : self.n].float() * (sx.reshape(-1, 1) * self.w_scale[: self.n]) | |
| y = y.view(*x.shape[:-1], self.n) | |
| return y if self.bias is None else y + self.bias | |
| class QuantConv1d(nn.Module): | |
| def __init__(self, mod: nn.Conv1d, w: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None): | |
| super().__init__() | |
| self.register_buffer("bias", bias) # before the early returns below | |
| self.k, self.s, self.d = _int1(mod.kernel_size), _int1(mod.stride), _int1(mod.dilation) | |
| if isinstance(mod.padding, str): # FocalCodec's "same" convs: stride 1, left-heavy pad | |
| assert mod.padding in {"same", "valid"} and self.s == 1, f"unsupported padding {mod.padding!r}" | |
| left = 0 if mod.padding == "valid" else (self.k - 1) // 2 | |
| self.pl, self.pr = (left, (self.k - 1) - left) if mod.padding == "same" else (0, 0) | |
| else: | |
| self.pl, self.pr = _pair(mod.padding) | |
| assert mod.groups == 1 or mod.groups == mod.in_channels or mod.in_channels % mod.groups == 0 | |
| self.kind = "dense" | |
| if mod.groups > 1: | |
| self.kind = "depthwise" if mod.groups == mod.in_channels and mod.out_channels == mod.groups else "grouped" | |
| self.groups = mod.groups | |
| self.ci = mod.in_channels | |
| w = w.reshape(w.shape[0], -1) # (co, (ci / groups) * k); accepts a 3-D conv weight | |
| n, ck = w.shape | |
| self.n, self.ck = n, ck | |
| sc = scale.reshape(-1) | |
| if self.kind == "depthwise": | |
| self.register_buffer("w_d", w.reshape(n, -1)) # (C, k), elementwise int MAC below | |
| self.register_buffer("w_scale_d", sc) | |
| return | |
| if self.kind == "grouped": | |
| co_g, ck_g = n // mod.groups, ck # ck is already (ci / groups) * k, group-major | |
| wts, scales = [], [] | |
| for gi in range(mod.groups): | |
| wt = torch.zeros((_up(co_g), _up(ck_g)), dtype=torch.int8, device=w.device) | |
| wt[:co_g, :ck_g] = w[gi * co_g : (gi + 1) * co_g] | |
| wts.append(wt.t().contiguous()) | |
| s = torch.zeros((_up(co_g),), dtype=torch.float32, device=w.device) | |
| s[:co_g] = sc[gi * co_g : (gi + 1) * co_g] | |
| scales.append(s) | |
| self.wT, self.w_scale_g = wts, scales | |
| self.co_g, self.ck_g = co_g, ck_g | |
| return | |
| wt = torch.zeros((_up(n), _up(ck)), dtype=torch.int8, device=w.device) | |
| wt[:n, :ck] = w | |
| self.register_buffer("wT", wt.t().contiguous()) | |
| s = torch.zeros((_up(n),), dtype=torch.float32, device=w.device) | |
| s[:n] = sc | |
| self.register_buffer("w_scale", s) | |
| def _cols(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: | |
| """fp32 input -> quantized im2col windows (B, Lout, ck) + per-position scales.""" | |
| x = F.pad(x, (self.pl, self.pr)) if (self.pl or self.pr) else x | |
| cols = F.unfold(x.unsqueeze(2), (1, self.k), dilation=(1, self.d), stride=(1, self.s)) | |
| q, sx = _act_quant(cols.permute(0, 2, 1)) # rows = output positions | |
| return q, sx, cols.shape[-1] | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| squeeze = x.dim() == 2 | |
| if squeeze: | |
| x = x.unsqueeze(0) | |
| assert x.shape[1] == self.ci, f"expected {self.ci} input channels, got {x.shape[1]}" | |
| b = x.shape[0] | |
| q, sx, lout = self._cols(x) | |
| sx_rows = sx.reshape(-1, 1) | |
| if self.kind == "depthwise": | |
| c = self.w_d.shape[0] | |
| acc = q.permute(0, 2, 1).view(b, c, self.k, lout).to(torch.int32) * self.w_d.to(torch.int32).view(c, self.k, 1) | |
| y = acc.sum(2).float() * sx.transpose(1, 2) * self.w_scale_d.view(1, -1, 1) # (B, C, Lout) | |
| y = y.permute(0, 2, 1) | |
| elif self.kind == "grouped": | |
| qg = q.view(b, lout, self.groups, self.ck_g) | |
| outs = [] | |
| for gi in range(self.groups): | |
| a = _k_pad(qg[:, :, gi].reshape(-1, self.ck_g), self.wT[gi].shape[0]) | |
| yi = _int_mm(a, self.wT[gi])[:, : self.co_g].float() * (sx_rows * self.w_scale_g[gi][: self.co_g]) | |
| outs.append(yi.view(b, lout, self.co_g)) | |
| y = torch.cat(outs, dim=-1) | |
| else: | |
| a = _k_pad(q.reshape(-1, self.ck), self.wT.shape[0]) | |
| y = _int_mm(a, self.wT)[:, : self.n].float() * (sx_rows * self.w_scale[: self.n]) | |
| y = y.view(b, lout, self.n) | |
| if self.bias is not None: | |
| y = y + self.bias | |
| y = y.transpose(1, 2) | |
| return y.squeeze(0) if squeeze else y | |
| class QuantConvTranspose1d(nn.Module): | |
| def __init__(self, mod: nn.ConvTranspose1d, w_eff: torch.Tensor, bias: torch.Tensor | None): | |
| super().__init__() | |
| assert _int1(mod.groups) == 1, "grouped ConvTranspose1d not supported" | |
| self.k, self.s, self.d = _int1(mod.kernel_size), _int1(mod.stride), _int1(mod.dilation) | |
| self.pl, self.pr = _pair(mod.padding) | |
| self.op = _int1(mod.output_padding) | |
| ci, co = w_eff.shape[0], w_eff.shape[1] | |
| self.ci, self.n = ci, co | |
| w, sc = _quant_tensor(w_eff) # per-K scales in the pack don't fit a GEMM: per-tensor here | |
| wts = [] | |
| for ki in range(self.k): | |
| wt = torch.zeros((_up(ci), _up(co)), dtype=torch.int8, device=w.device) | |
| wt[:ci, :co] = w[:, :, ki] | |
| wts.append(wt) | |
| self.wT, self.scale = wts, sc | |
| self.register_buffer("bias", bias) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| assert x.dim() == 3 and x.shape[1] == self.ci, f"expected (B, {self.ci}, L), got {tuple(x.shape)}" | |
| b, _, l = x.shape | |
| lout = (l - 1) * self.s - (self.pl + self.pr) + self.d * (self.k - 1) + 1 + self.op | |
| q, sx = _act_quant(x.permute(0, 2, 1)) # rows = (batch, time) positions | |
| a = _k_pad(q.reshape(-1, self.ci), self.wT[0].shape[0]) | |
| acc = torch.zeros((b, lout, self.n), dtype=torch.float32, device=x.device) | |
| rows = torch.arange(l, device=x.device) * self.s | |
| sx_rows = sx.reshape(-1, 1) | |
| for ki, wt in enumerate(self.wT): | |
| y = _int_mm(a, wt)[:, : self.n].float() * (sx_rows * self.scale) | |
| t = rows + ki * self.d - self.pl | |
| keep = (t >= 0) & (t < lout) | |
| if keep.any(): | |
| acc[:, t[keep]] += y.view(b, l, self.n)[:, keep] | |
| y = acc if self.bias is None else acc + self.bias | |
| return y.transpose(1, 2) | |
| class QuantEmbedding(nn.Module): | |
| def __init__(self, w: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None = None): | |
| super().__init__() | |
| # named `weight` so legacy hooks that touch `.weight.device` keep working; stays int8 | |
| self.register_buffer("weight", w) | |
| self.register_buffer("w_scale", scale.reshape(-1)) | |
| def forward(self, idx: torch.Tensor) -> torch.Tensor: | |
| return self.weight[idx].float() * self.w_scale[idx].unsqueeze(-1) | |
| _TARGETS = (nn.Linear, nn.Conv1d, nn.ConvTranspose1d, nn.Embedding) | |
| def _effective_weight(mod, name: str, packed: dict, device: torch.device, consumed: set) -> torch.Tensor: | |
| """Materialize a weight-normed layer's effective weight from the packed g/v via its own pre-hook.""" | |
| g_key, v_key = f"{name}.weight_g", f"{name}.weight_v" | |
| g = _dequant(packed, g_key, device) | |
| v = _dequant(packed, v_key, device) | |
| consumed.update({g_key, v_key}) | |
| hooks = tuple(mod._forward_pre_hooks.values()) | |
| assert hooks, f"{name}: weight-normed but no forward pre-hook to recompute `weight`" | |
| with torch.no_grad(): | |
| mod.weight_g.copy_(g) | |
| mod.weight_v.copy_(v) | |
| for hook in hooks: | |
| hook(mod, ()) | |
| return mod.weight.detach().clone() | |
| def _build_twin(mod, name: str, packed: dict, device: torch.device, consumed: set) -> nn.Module: | |
| bias = None | |
| bias_key = f"{name}.bias" | |
| if bias_key in packed: | |
| consumed.add(bias_key) | |
| bias = _dequant(packed, bias_key, device) | |
| if f"{name}.weight_g" in packed: | |
| w, scale, effective = _effective_weight(mod, name, packed, device, consumed), None, True | |
| elif f"{name}.weight" in packed: | |
| consumed.add(f"{name}.weight") | |
| raw = packed[f"{name}.weight"].to(device) | |
| scale = packed.get(SCALE_PREFIX + f"{name}.weight") | |
| scale = scale.to(device) if scale is not None else None | |
| w, effective = raw, raw.is_floating_point() # int8 packed -> direct; fp32 -> requantize effective | |
| else: | |
| raise KeyError(f"{name}: no weight in packed state") | |
| if isinstance(mod, nn.ConvTranspose1d): | |
| w_eff = w if effective else w.float() * scale | |
| return QuantConvTranspose1d(mod, w_eff, bias) | |
| if effective: # weight-normed (or caller passed fp32): quantize the effective weight now | |
| w, scale = _quant_channel(w) | |
| if isinstance(mod, nn.Linear): | |
| return QuantLinear(mod, w, scale, bias) | |
| if isinstance(mod, nn.Embedding): | |
| return QuantEmbedding(w, scale) | |
| return QuantConv1d(mod, w, scale, bias) | |
| def apply_int8(model: nn.Module, packed: dict) -> nn.Module: | |
| """Swap every Linear/Conv1d/ConvTranspose1d/Embedding for its int8 twin, in place. | |
| `packed` is the export format: int8 levels + `_scale/` fp32 scales (fp32 tensors are | |
| accepted too). Everything else in `packed` (biases of unswapped modules, norms, buffers) | |
| is materialized as fp32 and loaded; weights of swapped modules stay int8. | |
| """ | |
| device = next(model.parameters()).device | |
| consumed: set = set() | |
| twins: dict[str, nn.Module] = {} | |
| for name, mod in model.named_modules(): | |
| if name and isinstance(mod, _TARGETS): | |
| twin = _build_twin(mod, name, packed, device, consumed) | |
| twin.train(mod.training) | |
| twins[name] = twin | |
| for name, twin in twins.items(): | |
| parent_name, _, attr = name.rpartition(".") | |
| parent = model.get_submodule(parent_name) if parent_name else model | |
| parent._modules[attr] = twin | |
| fp_state = {} | |
| for key, value in packed.items(): | |
| if key.startswith(SCALE_PREFIX) or key in consumed: | |
| continue | |
| scale = packed.get(SCALE_PREFIX + key) | |
| # packed int8/fp8 -> dequantize (int8 without a scale would be a bug, not a buffer); | |
| # unscaled non-float buffers (int64 masks etc.) pass through with their own dtype | |
| fp_state[key] = value.float() * scale if scale is not None else value | |
| _, unexpected = model.load_state_dict(fp_state, strict=False) | |
| assert not unexpected, f"packed keys with no home in the model: {unexpected}" | |
| return model | |
| def load_packed(repo_id: str, filename: str) -> dict: | |
| from huggingface_hub import hf_hub_download | |
| from safetensors.torch import load_file | |
| return load_file(hf_hub_download(repo_id, filename, repo_type="model")) | |