digital-garbageandClaude Sonnet 5.5
b9717ed1cc
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
48 lines
2.5 KiB
Python
48 lines
2.5 KiB
Python
"""Int8 + rotation weights for MiniMax H3, made from an ordinary (bf16/fp16) file at load.
|
|
|
|
Comfy-Org ships H3 already quantized this way (half the memory, int8 matmuls on the tensor cores). This
|
|
does the same to a file that is not: each big Linear's weight is rotated (Hadamard, in groups) so no
|
|
weight has a lone huge value, then stored as int8 with one scale per output row -- in exactly the layout
|
|
a pre-quantized file has, so ComfyUI's own loader and kernels take it from there. Rotation is what
|
|
spares the outliers that broke plain fp8 (NaN audio), but the result is still an approximation of the
|
|
file: unvalidated on a GPU.
|
|
|
|
Only the transformer's attention and MLP projections are touched. The timestep modulation (adaLN, ~40% of
|
|
the bytes but tiny rows), embeddings, the output layer and the text refiner stay as they are.
|
|
"""
|
|
|
|
import json
|
|
import re
|
|
|
|
import torch
|
|
|
|
# blocks.N.attn.{qkv_proj,out_proj} / blocks.N.mlp.{fc1,fc2}, not token_refiner.blocks.*
|
|
_BIG = re.compile(r"(?:^|\.)(?<!refiner\.)blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))\.weight$")
|
|
GROUP = 256
|
|
_FLOATS = (torch.bfloat16, torch.float16, torch.float32)
|
|
_TAG = torch.tensor(list(json.dumps({"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": GROUP}).encode()), dtype=torch.uint8)
|
|
|
|
|
|
def quantize_state_dict(sd, device=None):
|
|
"""Rewrite `sd` in place. -> (layers converted, note). Already-quantized layers and rows that cannot be
|
|
rotated in groups (width not a multiple of 256) are left as they are, and the note says so."""
|
|
from comfy.quant_ops import QuantizedTensor
|
|
done, odd = 0, 0
|
|
for key in [k for k in sd if _BIG.search(k)]:
|
|
w = sd[key]
|
|
layer = key[:-len(".weight")]
|
|
if w.dtype not in _FLOATS or w.ndim != 2 or f"{layer}.comfy_quant" in sd:
|
|
continue
|
|
if w.shape[1] % GROUP:
|
|
odd += 1
|
|
continue
|
|
q = QuantizedTensor.from_float(w.to(device) if device is not None else w, "TensorWiseINT8Layout",
|
|
per_channel=True, convrot=True, convrot_groupsize=GROUP, scale="recalculate")
|
|
for name, tensor in q.state_dict(key).items():
|
|
sd[name] = tensor.cpu()
|
|
sd[f"{layer}.comfy_quant"] = _TAG.clone()
|
|
done += 1
|
|
if not done:
|
|
return 0, "int8_convrot: no H3 attention/MLP weights found to convert (not an H3 file, or already quantized) -- loaded as is"
|
|
return done, f"int8_convrot: {done} layers stored as int8" + (f"; {odd} left as they were (width not a multiple of {GROUP})" if odd else "")
|