Claude
b108d4f4f5
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr
64 lines
2.9 KiB
Python
64 lines
2.9 KiB
Python
"""MXFP8 weights for MiniMax H3, made from an ordinary (bf16/fp16) file at load.
|
|
|
|
Blackwell's tensor cores run FP8 at twice the bf16 rate, and the MXFP8 format keeps one power-of-two scale
|
|
per 32 weights, so no lone outlier spoils a whole row (the failure that gave plain fp8 NaN audio). The
|
|
weights are stored in exactly the layout a pre-quantized file has, so ComfyUI's own loader and kernels
|
|
(comfy_kitchen scaled_mm_mxfp8) take it from there.
|
|
|
|
The recipe follows NVIDIA's Sol-H3 (arXiv 2609.35110): the attention and MLP projections of the middle
|
|
blocks only. The first two blocks and the last three stay as they are, the timestep modulation (adaLN),
|
|
embeddings, the output layer and the text refiner too. The result is an approximation of the file.
|
|
"""
|
|
|
|
import json
|
|
import re
|
|
|
|
import torch
|
|
|
|
_BIG = re.compile(r"(?:^|\.)(?<!refiner\.)blocks\.(\d+)\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))\.weight$")
|
|
KEEP_FIRST = 2
|
|
KEEP_LAST = 3
|
|
GROUP = 32
|
|
_FLOATS = (torch.bfloat16, torch.float16, torch.float32)
|
|
_TAG = torch.tensor(list(json.dumps({"format": "mxfp8"}).encode()), dtype=torch.uint8)
|
|
|
|
|
|
def available(device=None):
|
|
"""True when this machine's ComfyUI can run MXFP8 matmuls (Blackwell, torch 2.10+, comfy_kitchen)."""
|
|
try:
|
|
import comfy.model_management as mm
|
|
from comfy.quant_ops import QUANT_ALGOS
|
|
return "mxfp8" in QUANT_ALGOS and bool(mm.supports_mxfp8_compute(device))
|
|
except Exception: # noqa: BLE001
|
|
return False
|
|
|
|
|
|
def quantize_state_dict(sd, device=None):
|
|
"""Rewrite `sd` in place. -> (layers converted, note). Already-quantized layers and weights whose width is
|
|
not a multiple of 32 are left as they are, and the note says so."""
|
|
from comfy.quant_ops import QuantizedTensor
|
|
keys = [(k, int(m.group(1))) for k in sd if (m := _BIG.search(k))]
|
|
if not keys:
|
|
return 0, "mxfp8: no H3 attention/MLP weights found to convert (not an H3 file, or already quantized) -- loaded as is"
|
|
count = max(i for _, i in keys) + 1
|
|
done, odd = 0, 0
|
|
for key, index in keys:
|
|
if index < KEEP_FIRST or index >= count - KEEP_LAST:
|
|
continue
|
|
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, "TensorCoreMXFP8Layout")
|
|
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, "mxfp8: nothing to convert (already quantized, or too few blocks) -- loaded as is"
|
|
return done, (f"mxfp8: {done} layers in blocks {KEEP_FIRST}-{count - KEEP_LAST - 1} stored as MXFP8"
|
|
+ (f"; {odd} left as they were (width not a multiple of {GROUP})" if odd else ""))
|