Upload convert_nvfp4.py with huggingface_hub

convert_nvfp4.py

@@ -0,0 +1,108 @@
+#!/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()