feat: int8_convrot option on the diffusion model loader (H3 bf16 file -> int8 + rotation at load)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

core/tests/test_shipped_nodes.py

@@ -73,7 +73,7 @@ FROZEN = {
     "FunPackDiffusionModelLoader": [
         "model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation",
         "sla", "sla_sparsity", "sla_block_size", "sla_protect_audio", "sla_min_seq_len",
-        "sla_dense_last_steps",
+        "sla_dense_last_steps", "int8_convrot",
     ],
     "FunPackLoraLoader": [
         "model", "lora_name", "strength_model", "clip", "strength_clip",

modules/loaders/diffusion_model/nodes.py

@@ -5,6 +5,7 @@ model IS, which is why they are node widgets and not settings: nobody downstream
 can supply them, and a saved workflow has to record them.
 """
 
+import comfy.model_management
 import comfy.sd
 import comfy.utils
 import folder_paths
@@ -12,6 +13,7 @@ from comfy_api.latest import io
 
 from ..._core import log
 from .. import gguf_support, sla_attention
+from .. import int8_convrot as int8_convrot_module
 from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices,
                       attention_override, dtype_of, set_fp16_accumulation,
                       weight_model_options)
@@ -98,6 +100,14 @@ class FunPackDiffusionModelLoader(io.ComfyNode):
                             tooltip="Run the last N sampling steps at full attention. 0 "
                                     "matches lightx2v; 1 was tested and did not help, for "
                                     "+20% time."),
+                io.Boolean.Input("int8_convrot", default=False, optional=True,
+                                 tooltip="MiniMax H3 from an ordinary (bf16/fp16) file: store the "
+                                         "attention and MLP weights as int8 (rotated first, so no lone "
+                                         "outlier spoils the rounding), the way Comfy-Org's own int8_convrot "
+                                         "files are. About half the memory and int8 matmuls on GPUs that "
+                                         "have them; slightly different pictures. Does nothing to a file "
+                                         "that is already quantized or is not H3 (the status says so), and "
+                                         "is skipped when a weight dtype is forced above."),
             ],
             outputs=[
                 io.Model.Output(display_name="model"),
@@ -110,7 +120,7 @@ class FunPackDiffusionModelLoader(io.ComfyNode):
                 attention: str, fp16_accumulation: bool = False, sla: bool = False,
                 sla_sparsity: float = None, sla_block_size: str = None,
                 sla_protect_audio: bool = None, sla_min_seq_len: int = None,
-                sla_dense_last_steps: int = None) -> io.NodeOutput:
+                sla_dense_last_steps: int = None, int8_convrot: bool = False) -> io.NodeOutput:
         notes = [f"FunPack Diffusion Model Loader | {model_name}"]
 
         accumulation = set_fp16_accumulation(fp16_accumulation)
@@ -147,6 +157,12 @@ class FunPackDiffusionModelLoader(io.ComfyNode):
             notes.append(gguf_note)
         else:
             state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
+        if int8_convrot:
+            if weight_dtype != "default":
+                notes.append(f"int8_convrot: skipped, weight dtype is forced to {weight_dtype}")
+            else:
+                _, note = int8_convrot_module.quantize_state_dict(state_dict, comfy.model_management.get_torch_device())
+                notes.append(note)
         model = comfy.sd.load_diffusion_model_state_dict(
             state_dict, model_options=model_options, metadata=metadata)
         if model is None:

modules/loaders/int8_convrot.py

@@ -0,0 +1,47 @@
+"""Int8 + rotation weights for MiniMax H3, made from an ordinary (bf16/fp16) file at load.
+
+Comfy-Org ships H3 already quantized this way (half the memory, int8 matmuls on the tensor cores). This
+does the same to a file that is not: each big Linear's weight is rotated (Hadamard, in groups) so no
+weight has a lone huge value, then stored as int8 with one scale per output row -- in exactly the layout
+a pre-quantized file has, so ComfyUI's own loader and kernels take it from there. Rotation is what
+spares the outliers that broke plain fp8 (NaN audio), but the result is still an approximation of the
+file: unvalidated on a GPU.
+
+Only the transformer's attention and MLP projections are touched. The timestep modulation (adaLN, ~40% of
+the bytes but tiny rows), embeddings, the output layer and the text refiner stay as they are.
+"""
+
+import json
+import re
+
+import torch
+
+# blocks.N.attn.{qkv_proj,out_proj} / blocks.N.mlp.{fc1,fc2}, not token_refiner.blocks.*
+_BIG = re.compile(r"(?:^|\.)(?<!refiner\.)blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))\.weight$")
+GROUP = 256
+_FLOATS = (torch.bfloat16, torch.float16, torch.float32)
+_TAG = torch.tensor(list(json.dumps({"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": GROUP}).encode()), dtype=torch.uint8)
+
+
+def quantize_state_dict(sd, device=None):
+    """Rewrite `sd` in place. -> (layers converted, note). Already-quantized layers and rows that cannot be
+    rotated in groups (width not a multiple of 256) are left as they are, and the note says so."""
+    from comfy.quant_ops import QuantizedTensor
+    done, odd = 0, 0
+    for key in [k for k in sd if _BIG.search(k)]:
+        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, "TensorWiseINT8Layout",
+                                       per_channel=True, convrot=True, convrot_groupsize=GROUP, scale="recalculate")
+        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, "int8_convrot: no H3 attention/MLP weights found to convert (not an H3 file, or already quantized) -- loaded as is"
+    return done, f"int8_convrot: {done} layers stored as int8" + (f"; {odd} left as they were (width not a multiple of {GROUP})" if odd else "")

modules/loaders/tests/test_int8_convrot.py

@@ -0,0 +1,51 @@
+import json
+
+import torch
+
+from .. import int8_convrot as m
+
+
+def _sd():
+    g = torch.Generator().manual_seed(0)
+    r = lambda *s: torch.randn(*s, generator=g, dtype=torch.bfloat16)
+    return {
+        "blocks.0.attn.qkv_proj.weight": r(96, 512), "blocks.0.attn.out_proj.weight": r(32, 512),
+        "blocks.0.mlp.fc1.weight": r(64, 512), "blocks.0.mlp.fc2.weight": r(32, 100),     # 100 is not a multiple of 256
+        "blocks.0.adaln_proj.linear.weight": r(64, 128), "blocks.0.adaln_proj.linear.bias": r(64),
+        "token_refiner.blocks.0.attn.qkv_proj.weight": r(96, 512), "video_patch_proj.weight": r(32, 512),
+    }
+
+
+def test_only_the_big_projections_are_converted():
+    sd = _sd()
+    keep = {k: v.clone() for k, v in sd.items()}
+    n, note = m.quantize_state_dict(sd)
+    assert n == 3 and "1 left" in note
+    for k in ("blocks.0.attn.qkv_proj", "blocks.0.attn.out_proj", "blocks.0.mlp.fc1"):
+        assert sd[k + ".weight"].dtype == torch.int8 and sd[k + ".weight_scale"].shape == (sd[k + ".weight"].shape[0], 1)
+        conf = json.loads(sd[k + ".comfy_quant"].numpy().tobytes())
+        assert conf == {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}
+    for k in ("blocks.0.mlp.fc2.weight", "blocks.0.adaln_proj.linear.weight", "token_refiner.blocks.0.attn.qkv_proj.weight", "video_patch_proj.weight"):
+        assert torch.equal(sd[k], keep[k])
+
+
+def test_second_pass_and_foreign_files_are_left_alone():
+    sd = _sd()
+    m.quantize_state_dict(sd)
+    again = {k: v.clone() for k, v in sd.items()}
+    assert m.quantize_state_dict(sd)[0] == 0 and all(torch.equal(sd[k], again[k]) for k in sd)
+    n, note = m.quantize_state_dict({"unet.weight": torch.zeros(4, 4)})
+    assert n == 0 and "loaded as is" in note
+
+
+def test_comfy_loads_it_and_the_output_matches():
+    import comfy.ops
+    sd = {"blocks.0.mlp.fc1.weight": torch.randn(64, 512, dtype=torch.bfloat16)}
+    ref = sd["blocks.0.mlp.fc1.weight"].clone()
+    m.quantize_state_dict(sd)
+    lin = comfy.ops.mixed_precision_ops({"mixed_ops": True}, torch.bfloat16).Linear(512, 64, bias=False)
+    lin.load_state_dict({k.split("fc1.")[1]: v for k, v in sd.items()}, strict=False)
+    x = torch.randn(3, 512, dtype=torch.bfloat16)
+    want = torch.nn.functional.linear(x, ref)
+    got = lin(x)
+    assert (got.float() - want.float()).abs().mean() / want.float().abs().mean() < 0.03