File size: 7,172 Bytes
fd369ce | 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 | """ComfyUI single file (export_comfy.py) -> ConvRot W8A8, the Comfy-Org/MOSS-Image ``*_int8_convrot`` layout.
The 168 block linears (``layers.N.proj.{qkv,o,ffn_up_gate,ffn_down}``) are stored as ComfyUI's native
``int8_tensorwise`` with ConvRot (arXiv 2512.03673); no rotation matrix is stored, the loader rebuilds it:
<layer>.weight int8 [out, in] W rotated per 256-wide input block by the regular Hadamard, quantized
<layer>.weight_scale f32 [out, 1] per-output-row scale, absmax / 127 of the rotated row
<layer>.comfy_quant uint8 JSON {"format": "int8_tensorwise", "convrot": true, "convrot_groupsize": 256}
layers.N.comfy_attention.config uint8 JSON {"attention": "comfy_kitchen_int8"}
H is H4 kron H4 kron H4 kron H4 / 16 with H4 = [[1,1,1,-1],[1,1,-1,1],[1,-1,1,1],[-1,1,1,1]]: symmetric and
orthogonal, so x W^T = (x H)(W H^T)^T and ComfyUI rotates activations online. Everything else (text embedder,
time / modulation, norms, head) passes through unchanged; empty file metadata, as in the Comfy-Org files. The math
mirrors comfy-kitchen's int8_tensorwise (sign pattern, Kronecker order, rounding) or the file decodes to garbage.
python quantize_convrot.py kromoss_bf16.safetensors kromoss_int8_convrot.safetensors \
[--compare moss_image_sft_int8_convrot.safetensors]
"""
from __future__ import annotations
import os
import argparse
import json
import re
import time
import torch
from safetensors import safe_open
from safetensors.torch import save_file
BLOCK_LINEAR = re.compile(r"^layers\.(\d+)\.proj\.(qkv|o|ffn_up_gate|ffn_down)\.weight$")
H4 = ((1, 1, 1, -1), (1, 1, -1, 1), (1, -1, 1, 1), (-1, 1, 1, 1))
CLIP_GRID = torch.linspace(0.55, 1.0, 80)
def regular_hadamard(n, device):
h = torch.tensor(H4, dtype=torch.float32, device=device)
while h.shape[0] < n:
h = torch.kron(h, torch.tensor(H4, dtype=torch.float32, device=device))
if h.shape[0] != n:
raise ValueError(f"group size must be a power of 4, got {n}")
return h / n ** 0.5
def rotate(w, gs):
out_f, in_f = w.shape
return (w.reshape(out_f, in_f // gs, gs) @ regular_hadamard(gs, w.device).T).reshape(out_f, in_f)
@torch.no_grad()
def quantize(w, gs, mseclip=False):
"""fp32 rotate, then per row scale = absmax / 127 (the Comfy-Org recipe), or with ``mseclip`` the clip ratio
in CLIP_GRID with the lowest reconstruction MSE."""
wf = rotate(w.float(), gs)
absmax = wf.abs().amax(dim=1, keepdim=True).clamp(min=1e-30)
best_mse, best_scale, best_q = torch.full_like(absmax, float("inf")), absmax / 127.0, None
for a in CLIP_GRID.tolist() if mseclip else [1.0]:
scale = (absmax * a / 127.0).clamp(min=1e-30)
q = (wf / scale).round_().clamp_(-127, 127)
mse = ((q * scale - wf) ** 2).mean(dim=1, keepdim=True)
better = mse < best_mse
best_mse = torch.where(better, mse, best_mse)
best_scale = torch.where(better, scale, best_scale)
best_q = q if best_q is None else torch.where(better, q, best_q)
deq = rotate(best_q * best_scale, gs).double()
cos = (deq.flatten() @ w.double().flatten() / (deq.norm() * w.double().norm())).item()
return best_q.to(torch.int8), best_scale, cos
def blob(obj):
return torch.tensor(list(json.dumps(obj).encode()), dtype=torch.uint8)
def convrot(path, groupsize=256, mseclip=False, int8_attention=True, device=None):
"""bf16 ComfyUI single file -> state dict of the ConvRot W8A8 file (save with metadata={})."""
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
out, start, worst = {}, time.time(), (1.0, "")
with safe_open(str(path), "pt") as f:
keys = list(f.keys())
blocks = sorted({int(m.group(1)) for k in keys if (m := BLOCK_LINEAR.match(k))})
for k in keys:
tensor = f.get_tensor(k)
if not BLOCK_LINEAR.match(k):
out[k] = tensor
continue
q, scale, cos = quantize(tensor.to(device), groupsize, mseclip)
if cos <= 0.99:
raise SystemExit(f"broken quantization of {k}: cos {cos:.5f}")
worst = min(worst, (cos, k))
base = k[:-len(".weight")]
out[k], out[f"{base}.weight_scale"] = q.cpu(), scale.cpu()
out[f"{base}.comfy_quant"] = blob({"format": "int8_tensorwise", "convrot": True,
"convrot_groupsize": groupsize})
if int8_attention:
for n in blocks:
out[f"layers.{n}.comfy_attention.config"] = blob({"attention": "comfy_kitchen_int8"})
print(f"quantized {sum(1 for k in out if k.endswith('.comfy_quant'))} linears in {time.time() - start:.0f}s, "
f"worst cos {worst[0]:.6f} ({worst[1]})", flush=True)
return out
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("input", help="bf16 ComfyUI single file (fused proj.qkv), e.g. from export_comfy.py")
parser.add_argument("output")
parser.add_argument("--groupsize", type=int, default=256)
parser.add_argument("--mseclip", action="store_true", help="per-row MSE-optimal clip instead of absmax "
"(~1.5%% lower weight error; not how the Comfy-Org files were made)")
parser.add_argument("--no-int8-attention", action="store_true",
help="omit the per-block comfy_attention.config (ComfyUI then keeps bf16 attention)")
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
parser.add_argument("--compare", help="reference int8 file: require the same keys/shapes/dtypes and report "
"how many int8 values and scales agree")
args = parser.parse_args()
out = convrot(args.input, args.groupsize, args.mseclip, not args.no_int8_attention, args.device)
if args.compare:
with safe_open(args.compare, "pt") as f:
if set(f.keys()) != set(out):
raise SystemExit(f"key sets differ: {sorted(set(f.keys()) ^ set(out))[:6]}")
q_same = q_total = s_same = s_total = exact = 0
for k in f.keys():
ref = f.get_tensor(k)
if ref.shape != out[k].shape or ref.dtype != out[k].dtype:
raise SystemExit(f"{k}: {tuple(ref.shape)} {ref.dtype} vs {tuple(out[k].shape)} {out[k].dtype}")
exact += torch.equal(ref, out[k])
if ref.dtype == torch.int8:
q_same += (ref == out[k]).sum().item()
q_total += ref.numel()
elif k.endswith(".weight_scale"):
s_same += (ref == out[k]).sum().item()
s_total += ref.numel()
print(f"vs {args.compare}: {exact}/{len(out)} tensors bit-identical; int8 values equal "
f"{100 * q_same / q_total:.4f}%, scales equal {100 * s_same / s_total:.4f}%")
save_file(out, args.output, metadata={})
print(f"{args.output}: {len(out)} tensors, {os.path.getsize(args.output) / 1e9:.2f} GB")
if __name__ == "__main__":
main()
|