| |
| """ |
| dequant-loader.py — Pre-dequantize quantized PersonaPlex weights to bf16 cache. |
| |
| For NF4 (INT4) or TurboQuant 2-bit weights, dequantizes to a temporary |
| bf16 safetensors file that moshi.server can load natively. |
| |
| Usage: |
| python dequant-loader.py --input model-nf4.safetensors --output /tmp/model-bf16.safetensors |
| python dequant-loader.py --input model-turbo2bit.safetensors --output /tmp/model-bf16.safetensors |
| |
| The output file can then be passed to moshi.server via --moshi-weight. |
| """ |
|
|
| import os, sys, math, time |
| import torch |
| from safetensors.torch import load_file, save_file |
|
|
| NF2_CENTROIDS = torch.tensor([-1.5104, -0.4528, 0.4528, 1.5104]) |
|
|
|
|
| def fast_wht(x): |
| """Vectorized Walsh-Hadamard Transform.""" |
| n = x.shape[-1] |
| h = 1 |
| while h < n: |
| x_view = x.view(*x.shape[:-1], -1, 2, h) |
| a = x_view[..., 0, :].clone() |
| b = x_view[..., 1, :].clone() |
| x_view[..., 0, :] = a + b |
| x_view[..., 1, :] = a - b |
| x = x_view.reshape(*x.shape) |
| h *= 2 |
| return x / math.sqrt(n) |
|
|
|
|
| def detect_format(state): |
| """Detect if weights are NF4 (INT4), TurboQuant 2-bit, or plain.""" |
| has_scales = any(k.endswith(".__scales__") for k in state) |
| has_packed = any(k.endswith(".packed") for k in state) |
| if has_packed: |
| return "turbo2bit" |
| if has_scales: |
| return "nf4" |
| return "plain" |
|
|
|
|
| def dequant_nf4(state): |
| """Dequantize INT4 NF4 weights.""" |
| result = {} |
| processed = set() |
|
|
| for name in list(state.keys()): |
| if name.endswith(".__scales__") or name.endswith(".__shape__") or name.endswith(".__numel__"): |
| continue |
| if name in processed: |
| continue |
|
|
| scales_key = f"{name}.__scales__" |
| if scales_key in state: |
| packed = state[name] |
| scales = state[scales_key].float() |
| shape = state[f"{name}.__shape__"].tolist() |
| numel = state[f"{name}.__numel__"].item() |
| group_size = 64 |
|
|
| lo = (packed & 0x0F).to(torch.int8) - 8 |
| hi = ((packed >> 4) & 0x0F).to(torch.int8) - 8 |
| unpacked = torch.zeros(packed.numel() * 2, dtype=torch.float32) |
| unpacked[0::2] = lo.float() |
| unpacked[1::2] = hi.float() |
|
|
| n_groups = scales.numel() |
| groups = unpacked[:n_groups * group_size].reshape(n_groups, group_size) |
| deq = (groups * scales.unsqueeze(1)).reshape(-1)[:numel] |
|
|
| orig_shape = [s for s in shape if s > 0] |
| result[name] = deq.reshape(orig_shape).to(torch.bfloat16) |
| processed.add(name) |
| else: |
| result[name] = state[name].to(torch.bfloat16) |
| processed.add(name) |
|
|
| return result |
|
|
|
|
| def dequant_turbo2bit(state): |
| """Dequantize TurboQuant 2-bit (NF2 + WHT) weights.""" |
| result = {} |
| processed = set() |
|
|
| for name in list(state.keys()): |
| if any(name.endswith(f".{s}") for s in ["packed", "scales", "shape", "numel", "gs", "np2"]): |
| continue |
| if name in processed: |
| continue |
|
|
| packed_key = f"{name}.packed" |
| if packed_key in state: |
| gs = state[f"{name}.gs"].item() |
| gs_pow2 = state[f"{name}.np2"].item() |
| numel = state[f"{name}.numel"].item() |
| shape = [s for s in state[f"{name}.shape"].tolist() if s > 0] |
| scales = state[f"{name}.scales"].float() |
| packed = state[packed_key] |
| n_groups = scales.numel() |
|
|
| |
| p = packed.reshape(n_groups, gs // 4) |
| codes = torch.zeros(n_groups, gs, dtype=torch.long) |
| for i in range(4): |
| codes[:, i::4] = (p >> (2 * i)) & 0x03 |
|
|
| dequant = NF2_CENTROIDS[codes] |
|
|
| |
| if gs_pow2 > gs: |
| dequant = torch.cat([dequant, torch.zeros(n_groups, gs_pow2 - gs)], dim=1) |
| dequant = fast_wht(dequant) |
| dequant = dequant[:, :gs] |
|
|
| dequant = dequant * scales.unsqueeze(1) |
| result[name] = dequant.reshape(-1)[:numel].reshape(shape).to(torch.bfloat16) |
| processed.add(name) |
| else: |
| result[name] = state[name].to(torch.bfloat16) |
| processed.add(name) |
|
|
| return result |
|
|
|
|
| def main(): |
| import argparse |
| parser = argparse.ArgumentParser(description="Dequantize PersonaPlex weights to bf16") |
| parser.add_argument("--input", "-i", required=True, help="Quantized safetensors file") |
| parser.add_argument("--output", "-o", required=True, help="Output bf16 safetensors file") |
| parser.add_argument("--device", "-d", default="cpu", help="Device for dequantization") |
| args = parser.parse_args() |
|
|
| if not os.path.exists(args.input): |
| print(f"Error: {args.input} not found") |
| sys.exit(1) |
|
|
| |
| if os.path.exists(args.output) and os.path.getmtime(args.output) > os.path.getmtime(args.input): |
| print(f"Cached: {args.output} is up to date") |
| sys.exit(0) |
|
|
| print(f"Loading {args.input}...") |
| t0 = time.time() |
| state = load_file(args.input, device=args.device) |
|
|
| fmt = detect_format(state) |
| print(f"Format: {fmt}") |
|
|
| if fmt == "nf4": |
| result = dequant_nf4(state) |
| elif fmt == "turbo2bit": |
| result = dequant_turbo2bit(state) |
| else: |
| print("Already plain bf16/fp16 — copying") |
| result = {k: v.to(torch.bfloat16) for k, v in state.items()} |
|
|
| t1 = time.time() |
| print(f"Dequantized {len(result)} tensors in {t1-t0:.1f}s") |
|
|
| print(f"Saving to {args.output}...") |
| save_file(result, args.output) |
| size_gb = os.path.getsize(args.output) / 1024**3 |
| print(f"Done: {size_gb:.2f} GB") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
