| """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()) |
|
|
| |
| 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) |
|
|
| 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) |
|
|
| |
| 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) |
|
|
