File size: 3,364 Bytes
1564656 | 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 | """Requantize MiniMax-H3 pruned_int8_convrot -> NVFP4.
The StarNodes converter cannot do this: it passes non-floating-point tensors through
untouched, so an already-int8 model comes out byte-identical. This walks the state dict
and does a real dequantize -> requantize on the 200 attn/mlp layers, leaving AdaLN, the
norms, embedders and token_refiner exactly as they are.
Per-layer input: .weight (I8), .weight_scale (F32), .comfy_quant (U8 json bytes)
Per-layer output: .weight + NVFP4 layout suffixes, .comfy_quant = {"format": "nvfp4"}
"""
import json, os, sys, time, collections
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from comfy_kitchen import tensor as ckt
SRC = "/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors"
DST = "/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
DEV = "cuda"
FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
f = safe_open(SRC, "pt")
keys = list(f.keys())
# layers that carry a comfy_quant blob are the quantized ones
quant_bases = sorted({k[: -len(".comfy_quant")] for k in keys if k.endswith(".comfy_quant")})
print(f"quantized layers found: {len(quant_bases)}", flush=True)
out, stats = {}, collections.Counter()
consumed = set()
t0 = time.time()
for i, base in enumerate(quant_bases):
wk, sk, ck_ = f"{base}.weight", f"{base}.weight_scale", f"{base}.comfy_quant"
cfg = json.loads(bytes(f.get_tensor(ck_).tolist()).decode())
if cfg.get("format") != "int8_tensorwise":
print(f" SKIP {base}: unexpected format {cfg}", flush=True)
stats["skipped_unknown_format"] += 1
continue
qdata = f.get_tensor(wk).to(DEV)
scale = f.get_tensor(sk).to(DEV)
params = ckt.TensorWiseINT8Layout.Params(
scale=scale,
orig_dtype=torch.bfloat16,
orig_shape=tuple(qdata.shape),
is_weight=True,
convrot=cfg.get("convrot", False),
convrot_groupsize=cfg.get("convrot_groupsize", 256),
)
deq = ckt.TensorWiseINT8Layout.dequantize(qdata, params) # -> bf16
nq, nparams = ckt.TensorCoreNVFP4Layout.quantize(deq.float().contiguous())
tensors = ckt.TensorCoreNVFP4Layout.state_dict_tensors(nq, nparams)
for suffix, t in tensors.items():
name = f"{base}.weight{suffix}"
if t.dtype == torch.float8_e8m0fnu:
out[name] = t.view(torch.uint8).cpu()
elif t.dtype in FP8_DTYPES:
out[name] = t.view(torch.uint8).cpu().view(t.dtype)
else:
out[name] = t.cpu()
blob = json.dumps({"format": "nvfp4"}).encode()
out[ck_] = torch.tensor(list(blob), dtype=torch.uint8)
consumed.update({wk, sk, ck_})
stats["nvfp4"] += 1
del qdata, scale, deq, nq
if (i + 1) % 25 == 0:
torch.cuda.empty_cache()
print(f" {i+1}/{len(quant_bases)} {time.time()-t0:.0f}s", flush=True)
# copy everything untouched
for k in keys:
if k in consumed or k in out:
continue
out[k] = f.get_tensor(k)
stats["copied"] += 1
print(f"stats: {dict(stats)}", flush=True)
print(f"writing {DST} ...", flush=True)
save_file(out, DST)
gb = os.path.getsize(DST) / 1e9
print(f"DONE {gb:.2f} GB in {time.time()-t0:.0f}s", flush=True)
c = collections.Counter(str(v.dtype) for v in out.values())
print("dtype census:", dict(c), flush=True)
|