File size: 4,439 Bytes
72ea840 | 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 | #!/usr/bin/env python3
"""Convert the MiniMax H3 bf16 DiT to NVFP4 in ComfyUI's native quant layout.
Runs ON the instance (needs /workspace/ComfyUI + comfy_kitchen + a GPU).
Mirrors the quantization policy of the released int8_convrot DiT exactly:
only the 50 main blocks' attn.qkv_proj / attn.out_proj / mlp.fc1 / mlp.fc2
are quantized (200 layers); norms, adaln, patch/condition projections,
token_refiner and final layers stay bf16.
Per quantized layer (identical tensor layout to the released nvfp4_awq TE):
.weight U8 [out, in/2] packed FP4 E2M1 pairs
.weight_scale F8_E4M3 [out, in/16] per-16-block scales
.weight_scale_2 F32 [] global scale (amax / (448*6))
.comfy_quant U8 [n] JSON layer config
Usage:
python3 convert_nvfp4.py --src models/diffusion_models/minimax_h3_ref2va_bf16.safetensors \
--dst models/diffusion_models/minimax_h3_ref2va_nvfp4.safetensors [--fpmm]
--fpmm adds {"full_precision_matrix_mult": true} (dequant->bf16 GEMM, the
quality-safe path the official TE uses). Without it, comfy_kitchen's
native FP4 tensor-core GEMM path is used (faster, more quality risk).
"""
import argparse
import json
import sys
import time
sys.path.insert(0, "/workspace/ComfyUI")
import torch # noqa: E402
from safetensors import safe_open # noqa: E402
from safetensors.torch import save_file # noqa: E402
from comfy.quant_ops import ( # noqa: E402
TensorCoreConvRotW4A4Layout,
TensorCoreNVFP4Layout,
)
TARGET_SUFFIXES = (".attn.qkv_proj.weight", ".attn.out_proj.weight",
".mlp.fc1.weight", ".mlp.fc2.weight")
def should_quantize(key, shape):
return (key.startswith("blocks.")
and key.endswith(TARGET_SUFFIXES)
and len(shape) == 2)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--dst", required=True)
ap.add_argument("--algo", choices=["nvfp4", "convrot_w4a4"], default="nvfp4")
ap.add_argument("--fpmm", action="store_true",
help="nvfp4 only: full_precision_matrix_mult=true "
"(dequant->bf16 GEMM)")
args = ap.parse_args()
if args.algo == "convrot_w4a4":
# int4 weights + int4 activations, rotation-assisted (groupsize 256)
cfg = {"format": "convrot_w4a4", "convrot_groupsize": 256,
"linear_dtype": "int4"}
else:
cfg = {"format": "nvfp4"}
if args.fpmm:
cfg["full_precision_matrix_mult"] = True
cfg_tensor = torch.tensor(list(json.dumps(cfg).encode("utf-8")),
dtype=torch.uint8)
out, n_q, n_keep, t0 = {}, 0, 0, time.time()
with safe_open(args.src, framework="pt", device="cpu") as f:
keys = list(f.keys())
for i, k in enumerate(keys):
t = f.get_tensor(k)
if should_quantize(k, t.shape):
layer = k[:-len(".weight")]
w = t.cuda()
if args.algo == "convrot_w4a4":
qdata, params = TensorCoreConvRotW4A4Layout.quantize(
w, convrot_groupsize=256, linear_dtype="int4")
out[layer + ".weight"] = qdata.contiguous().cpu()
out[layer + ".weight_scale"] = params.scale.contiguous().cpu()
else:
qdata, params = TensorCoreNVFP4Layout.quantize(w)
out[layer + ".weight"] = qdata.contiguous().cpu()
out[layer + ".weight_scale"] = params.block_scale.contiguous().cpu()
out[layer + ".weight_scale_2"] = params.scale.to(torch.float32).cpu()
out[layer + ".comfy_quant"] = cfg_tensor.clone()
del w, qdata, params
n_q += 1
if n_q % 20 == 0:
torch.cuda.empty_cache()
print(f"[{time.time()-t0:6.0f}s] quantized {n_q} layers "
f"({i+1}/{len(keys)} tensors)", flush=True)
else:
out[k] = t
n_keep += 1
print(f"quantized {n_q} layers, kept {n_keep} tensors; saving {args.dst}")
save_file(out, args.dst)
size = sum(v.numel() * v.element_size() for v in out.values())
print(f"done in {time.time()-t0:.0f}s — ~{size/1e9:.1f} GB "
f"(config: {json.dumps(cfg)})")
if __name__ == "__main__":
main()
|