File size: 5,150 Bytes
b9fc164 | 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 | #!/usr/bin/env python3
"""Group-wise symmetric int4 quantization of the JoyAI-Video-Edit DiT.
Reads the official diffusers shards, renames `double_blocks.N.attn.X` ->
`double_blocks.N.X` (the runtime's layout), quantizes every large 2D Linear
weight to int4 (group_size=64) and writes packed weights + bf16 group scales.
out/int4-0000X-of-00002.safetensors : packed uint8 [out, in/2] + scales
out/meta.json : group size, key list, dtype map
"""
from __future__ import annotations
import json
import os
import time
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import save_file
SRC = Path("/home/ubuntu/jve-diff")
OUT = SRC / "int4"
OUT.mkdir(exist_ok=True)
GROUP = 64
MIN_NUMEL = 8_000_000 # quantize the big Linears only
def map_key(k: str) -> str:
if k.startswith("double_blocks."):
parts = k.split(".")
if parts[2] == "attn":
return ".".join(parts[:2] + parts[3:])
return k
def quantize(w: torch.Tensor):
"""w: [out, in] -> (packed uint8 [out, in/2], scales bf16 [out, in/64])"""
out_f, in_f = w.shape
assert in_f % GROUP == 0 and in_f % 2 == 0, w.shape
wr = w.float().reshape(out_f, in_f // GROUP, GROUP)
scale = wr.abs().amax(dim=-1).clamp_min(1e-8) / 7.0
q = torch.round(wr / scale.unsqueeze(-1)).clamp_(-8, 7).to(torch.int8)
q = q.reshape(out_f, in_f) + 8 # 0..15
q8 = q.to(torch.uint8)
packed = (q8[:, 0::2] | (q8[:, 1::2] << 4)).contiguous()
return packed, scale.to(torch.bfloat16).contiguous()
def main() -> None:
t0 = time.time()
idx = json.load(open(SRC / "transformer_diffusion_pytorch_model.safetensors.index.json"))
weight_map = idx["weight_map"]
shards = sorted(set(weight_map.values()))
print(f"{len(weight_map)} tensors in {len(shards)} shards", flush=True)
# group keys by shard so we stream one shard at a time (RAM stays ~5 GB)
per_shard: dict[str, list[str]] = {}
for k, s in weight_map.items():
per_shard.setdefault(s, []).append(k)
packed_t: dict[str, torch.Tensor] = {}
meta: dict[str, dict] = {}
stats = []
for si, shard in enumerate(shards, 1):
path = SRC / "transformer" / shard
with safe_open(str(path), framework="pt", device="cpu") as f:
for k in per_shard[shard]:
t = f.get_tensor(k).to(torch.bfloat16)
nk = map_key(k)
if t.dim() == 2 and t.numel() >= MIN_NUMEL:
packed, scales = quantize(t)
packed_t[nk + ".qweight"] = packed
packed_t[nk + ".scales"] = scales
# reconstruction error
deq = (packed & 0x0F).to(torch.float32) - 8
deq2 = (packed >> 4).to(torch.float32) - 8
w_deq = torch.empty(t.shape, dtype=torch.float32)
w_deq[:, 0::2] = deq
w_deq[:, 1::2] = deq2
w_deq = (w_deq.reshape(t.shape[0], -1, GROUP) *
scales.float().unsqueeze(-1)).reshape(t.shape)
tf = t.float()
mse = (tf - w_deq).pow(2).mean().item()
var = tf.pow(2).mean().item()
cos = torch.nn.functional.cosine_similarity(
tf.flatten(), w_deq.flatten(), dim=0).item()
stats.append((nk, mse / max(var, 1e-12), cos, t.numel()))
meta[nk] = {"group": GROUP, "shape": list(t.shape),
"dtype": "int4", "rel_mse": mse / max(var, 1e-12)}
else:
packed_t[nk] = t.contiguous()
meta[nk] = {"dtype": "bf16", "shape": list(t.shape)}
print(f" shard {si}/{len(shards)} done tensors={len(packed_t)} "
f"{time.time() - t0:.0f}s", flush=True)
# split the output into two files to keep each under ~5 GB
keys = sorted(packed_t)
half = len(keys) // 2
parts = [keys[:half], keys[half:]]
for pi, ks in enumerate(parts, 1):
save_file({k: packed_t[k] for k in ks},
str(OUT / f"int4-{pi:05d}-of-{len(parts):05d}.safetensors"))
print(f"wrote part {pi} ({len(ks)} tensors)", flush=True)
json.dump({"group_size": GROUP, "tensors": meta},
open(OUT / "meta.json", "w"), indent=1)
n_q = len(stats)
params_q = sum(s[3] for s in stats)
rel = sorted(s[1] for s in stats)
cosl = sorted(s[2] for s in stats)
print(f"\nquantized {n_q} tensors, {params_q/1e9:.2f} B params")
print(f"rel_mse mean={sum(rel)/len(rel):.5f} p50={rel[len(rel)//2]:.5f} max={rel[-1]:.5f}")
print(f"cosine mean={sum(cosl)/len(cosl):.6f} min={cosl[0]:.6f}")
worst = sorted(stats, key=lambda s: -s[1])[:3]
for nk, m, c, n in worst:
print(f" worst {nk:<48} rel_mse={m:.5f} cos={c:.5f}")
tot = sum(p.numel() * p.element_size() for p in packed_t.values())
print(f"output size {tot/2**30:.2f} GiB")
print(f"done in {time.time() - t0:.0f}s")
if __name__ == "__main__":
main()
|