From b9717ed1ccbd548fb71f1d638cca4fcc12fc53bb Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 4 Oct 2026 14:09:24 +0300 Subject: [PATCH] feat: int8_convrot option on the diffusion model loader (H3 bf16 file -> int8 + rotation at load) Co-Authored-By: Claude Sonnet 5.5 --- core/tests/test_shipped_nodes.py | 2 +- modules/loaders/diffusion_model/nodes.py | 18 +++++++- modules/loaders/int8_convrot.py | 47 ++++++++++++++++++++ modules/loaders/tests/test_int8_convrot.py | 51 ++++++++++++++++++++++ 4 files changed, 116 insertions(+), 2 deletions(-) create mode 100644 modules/loaders/int8_convrot.py create mode 100644 modules/loaders/tests/test_int8_convrot.py diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 6a65c912..50c46625 100644 --- a/core/tests/test_shipped_nodes.py +++ b/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", diff --git a/modules/loaders/diffusion_model/nodes.py b/modules/loaders/diffusion_model/nodes.py index 3c155174..12c40ae9 100644 --- a/modules/loaders/diffusion_model/nodes.py +++ b/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: diff --git a/modules/loaders/int8_convrot.py b/modules/loaders/int8_convrot.py new file mode 100644 index 00000000..7ca5c998 --- /dev/null +++ b/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"(?:^|\.)(? (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 "") diff --git a/modules/loaders/tests/test_int8_convrot.py b/modules/loaders/tests/test_int8_convrot.py new file mode 100644 index 00000000..69f01a7a --- /dev/null +++ b/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