File size: 2,570 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 | #!/usr/bin/env python3
"""All-int8 (per-output-channel symmetric) quantization of the DiT weights."""
from __future__ import annotations
import json
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 / "int8"
OUT.mkdir(exist_ok=True)
MIN_NUMEL = 8_000_000
def map_key(k: str) -> str:
if k.startswith("double_blocks."):
p = k.split(".")
if len(p) > 2 and p[2] == "attn":
return ".".join(p[:2] + p[3:])
return k
def q8(w: torch.Tensor):
wf = w.float()
scale = wf.abs().amax(dim=1, keepdim=True).clamp_min(1e-8) / 127.0
q = torch.round(wf / scale).clamp_(-127, 127).to(torch.int8).contiguous()
return q, scale.to(torch.float16).contiguous()
def main():
t0 = time.time()
idx = json.load(open(SRC / "transformer_diffusion_pytorch_model.safetensors.index.json"))
per_shard = {}
for k, s in idx["weight_map"].items():
per_shard.setdefault(s, []).append(k)
out_t, meta = {}, {}
n_q = 0
errs = []
for shard in sorted(per_shard):
with safe_open(str(SRC / "transformer" / shard), framework="pt") 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:
q, sc = q8(t)
out_t[nk + ".q8weight"] = q
out_t[nk + ".q8scales"] = sc
meta[nk] = {"dtype": "int8", "shape": list(t.shape)}
n_q += 1
if n_q % 40 == 0:
deq = q.float() * sc.float()
errs.append(((deq - t.float()).pow(2).mean() / t.float().pow(2).mean()).item())
else:
out_t[nk] = t.contiguous()
meta[nk] = {"dtype": "bf16", "shape": list(t.shape)}
keys = sorted(out_t)
third = len(keys) // 3
parts = [keys[:third], keys[third:2 * third], keys[2 * third:]]
for pi, ks in enumerate(parts, 1):
save_file({k: out_t[k] for k in ks}, str(OUT / f"int8-{pi:05d}-of-0000{len(parts)}.safetensors"))
json.dump({"group_size": 0, "tensors": meta}, open(OUT / "meta.json", "w"))
tot = sum(p.numel() * p.element_size() for p in out_t.values())
print(f"quantized {n_q} tensors, size {tot/2**30:.2f} GiB, "
f"mean rel_mse {sum(errs)/len(errs):.5f} ({time.time()-t0:.0f}s)")
if __name__ == "__main__":
main()
|