nodes/SparknightLLC-ComfyUI-INT8-Fast-Fork
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add MiniMax H3 support and ConvRot validation

Browse Source 
Introduce MODEL_TYPE_MINIMAX_H3, keep-float/int4-sensitive presets, and automatic adapter detection; bump package version to 2.1.0 and update changelog. Add robust ConvRot group-size validation (_validate_convrot_group_size), stricter Hadamard shape checks, and improved error handling when loading/saving native INT8 ConvRot metadata. Add unit tests covering ConvRot metadata round-trips, invalid metadata rejection, save-time consistency checks, and MiniMax H3 detection.
This commit is contained in:
ThereforeGames
2026-08-07 06:51:09 -04:00
parent 722dadcbc5
commit dc1c2cc498
6 changed files with 205 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+15 -1
View File
@@ -1,6 +1,20 @@
# Changelog
<details open><summary>2.0.0</summary>
<details open><summary>2.1.0</summary>
### Added
- Added preliminary MiniMax H3 keep-float and INT4-sensitive quantization tiers, including automatic architecture detection.
- Added native ConvRot metadata regression coverage for heterogeneous group sizes and invalid checkpoint metadata.
### Fixed
- Preserved each native INT8 ConvRot layer's checkpoint-defined group size at runtime, including MiniMax H3 AdaLN projections that use 64-channel groups.
- Validated ConvRot group divisibility and Hadamard state while loading and saving so malformed checkpoints fail with a clear error instead of crashing later or exporting inconsistent metadata.
</details>
<details><summary>2.0.0</summary>
### Added
int8_model_adapter.py
+9
View File
@@ -47,6 +47,7 @@ from .int8_unet_loader import MODEL_TYPE_FLUX2_FAST_UNSAFE
from .int8_unet_loader import MODEL_TYPE_HIDREAM_O1
from .int8_unet_loader import MODEL_TYPE_IDEOGRAM4
from .int8_unet_loader import MODEL_TYPE_KREA2
from .int8_unet_loader import MODEL_TYPE_MINIMAX_H3
from .int8_unet_loader import DEFAULT_QUANTIZATION_MODE
from .int8_unet_loader import get_model_type_quantization_preset
from .quantization_policy import DEFAULT_INT4_MIXED_RATIO
@@ -123,6 +124,13 @@ MODEL_TYPE_FINGERPRINTS = {
"txtfusion",
"txtmlp",
),
MODEL_TYPE_MINIMAX_H3: (
"video_patch_proj",
"audio_patch_proj",
"condition_proj",
"token_refiner",
"adaln_proj",
),
"z-image": (
"cap_embedder",
"context_refiner",
@@ -183,6 +191,7 @@ MODEL_TYPE_REQUIRED_MARKERS = {
MODEL_TYPE_HIDREAM_O1: ("language_model.layers",),
MODEL_TYPE_IDEOGRAM4: ("embed_image_indicator", "t_embedding"),
MODEL_TYPE_KREA2: ("txtfusion", "tmlp", "tproj"),
MODEL_TYPE_MINIMAX_H3: ("video_patch_proj", "audio_patch_proj", "token_refiner"),
"z-image": ("cap_embedder", "context_refiner", "noise_refiner"),
"wan": ("patch_embedding", "time_projection"),
"ltx2": ("audio_adaln_single", "audio_caption_projection", "audio_patchify_proj"),
int8_quant.py
+50 -8
View File
@@ -313,6 +313,22 @@ def _build_outlier_hadamard(method, group_size: int, device=None, dtype=torch.fl
return None
return _quarot_build_hadamard(group_size, device=device, dtype=dtype)
def _validate_convrot_group_size(group_size, in_features: int, context: str) -> int:
try:
group_size = int(group_size)
except (TypeError, ValueError) as e:
raise ValueError(f"{context}: invalid ConvRot group size {group_size!r}") from e
if group_size <= 0:
raise ValueError(f"{context}: ConvRot group size must be positive, got {group_size}")
if group_size < 4 or (group_size & (group_size - 1)) != 0 or ((group_size.bit_length() - 1) % 2) != 0:
raise ValueError(f"{context}: ConvRot group size must be a power of 4, got {group_size}")
if in_features % group_size != 0:
raise ValueError(
f"{context}: in_features {in_features} is not divisible by ConvRot group size {group_size}"
)
return group_size
def _rotate_weight_for_outlier_method(weight: Tensor, h_matrix: Tensor, group_size: int, method) -> Tensor:
normalized = _normalize_outlier_method(method)
if normalized == OUTLIER_METHOD_CONVROT:
@@ -2128,7 +2144,11 @@ if _COMFY_OPS_AVAILABLE:
self.hadanorm_sigma = None
self.weight_packed = None
self.weight_int_mm = None
self._convrot_groupsize = int(native_quant_config.get("convrot_groupsize", _CONVROT_GROUP_SIZE))
self._convrot_groupsize = _validate_convrot_group_size(
native_quant_config.get("convrot_groupsize", _CONVROT_GROUP_SIZE),
self.in_features,
f"Quantized layer {prefix.rstrip('.') or '<root>'}",
)
self._quant_group_size = _INT4_QUANT_GROUP_SIZE
self._linear_dtype = str(native_quant_config.get("linear_dtype", "int4"))
orig_dtype = getattr(self, "weight_comfy_model_dtype", None) or torch.bfloat16
@@ -2186,7 +2206,11 @@ if _COMFY_OPS_AVAILABLE:
and native_quant_config.get("format") == "int8_tensorwise"
and native_quant_config.get("convrot", False)
):
convrot_group_size = int(native_quant_config.get("convrot_groupsize", _CONVROT_GROUP_SIZE))
convrot_group_size = _validate_convrot_group_size(
native_quant_config.get("convrot_groupsize", _CONVROT_GROUP_SIZE),
self.in_features,
f"Quantized layer {prefix.rstrip('.') or '<root>'}",
)
try:
h_matrix = _build_outlier_hadamard(
OUTLIER_METHOD_CONVROT,
@@ -2194,11 +2218,15 @@ if _COMFY_OPS_AVAILABLE:
device="cpu",
dtype=torch.float32,
)
if h_matrix is not None:
self.quarot_hadamard = h_matrix
self._outlier_method = OUTLIER_METHOD_CONVROT
if h_matrix is None:
raise RuntimeError("ConvRot Hadamard helper is unavailable")
self._convrot_groupsize = convrot_group_size
self.quarot_hadamard = h_matrix
self._outlier_method = OUTLIER_METHOD_CONVROT
except Exception as e:
logging.warning(f"Quantization Toolkit: ConvRot metadata ignored for {prefix} ({e}).")
raise ValueError(
f"Quantized layer {prefix.rstrip('.') or '<root>'}: invalid ConvRot metadata ({e})"
) from e
elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32) or _is_float8_dtype(weight_tensor.dtype):
# Load High-Precision
@@ -2379,11 +2407,16 @@ if _COMFY_OPS_AVAILABLE:
if self._quant_format == "convrot_w4a4" and _is_native_int4_tensor(self.weight):
params = self.weight._params
convrot_group_size = _validate_convrot_group_size(
params.convrot_groupsize,
self.in_features,
f"Quantized layer {prefix.rstrip('.') or '<root>'}",
)
output[prefix + "weight"] = self.weight._qdata if keep_vars else self.weight._qdata.detach()
output[prefix + "weight_scale"] = params.scale if keep_vars else params.scale.detach()
quant_config = {
"format": "convrot_w4a4",
"convrot_groupsize": int(params.convrot_groupsize),
"convrot_groupsize": convrot_group_size,
}
if params.linear_dtype != "int4":
quant_config["linear_dtype"] = params.linear_dtype
@@ -2397,7 +2430,16 @@ if _COMFY_OPS_AVAILABLE:
quant_config = {"format": "int8_tensorwise"}
if outlier_method == OUTLIER_METHOD_CONVROT:
hadamard = getattr(self, "quarot_hadamard", None)
convrot_group_size = int(hadamard.shape[-1]) if isinstance(hadamard, torch.Tensor) else _CONVROT_GROUP_SIZE
convrot_group_size = _validate_convrot_group_size(
self._convrot_groupsize,
self.in_features,
f"Quantized layer {prefix.rstrip('.') or '<root>'}",
)
if not isinstance(hadamard, torch.Tensor) or tuple(hadamard.shape) != (convrot_group_size, convrot_group_size):
raise ValueError(
f"Quantized layer {prefix.rstrip('.') or '<root>'}: ConvRot Hadamard shape must be "
f"({convrot_group_size}, {convrot_group_size})"
)
quant_config["convrot"] = True
quant_config["convrot_groupsize"] = convrot_group_size
output.pop(prefix + "quarot_hadamard", None)
int8_unet_loader.py
+11
View File
@@ -29,6 +29,7 @@ MODEL_TYPE_BOOGU = "boogu"
MODEL_TYPE_HIDREAM_O1 = "hidream o1"
MODEL_TYPE_IDEOGRAM4 = "ideogram4"
MODEL_TYPE_KREA2 = "krea2"
MODEL_TYPE_MINIMAX_H3 = "minimax_h3"
MODEL_TYPE_CHOICES = [
"anima",
MODEL_TYPE_BOOGU,
@@ -39,6 +40,7 @@ MODEL_TYPE_CHOICES = [
MODEL_TYPE_HIDREAM_O1,
MODEL_TYPE_IDEOGRAM4,
MODEL_TYPE_KREA2,
MODEL_TYPE_MINIMAX_H3,
"ltx2",
"qwen",
"sdxl",
@@ -100,6 +102,11 @@ def _get_model_type_keep_float(model_type):
return [
"first", "last", "tmlp", "tproj", "txtfusion", "txtmlp",
]
if model_type == MODEL_TYPE_MINIMAX_H3:
return [
"video_patch_proj", "audio_patch_proj", "condition_proj", "time_embedder",
"token_refiner", "adaln_proj", "final_layer",
]
if model_type == "sdxl":
return [
"time_embed", "label_emb", "emb_layers", "proj_in", "proj_out",
@@ -129,6 +136,10 @@ def _get_model_type_int4_sensitive(model_type):
return [
"attn.wo", "mlp.down",
]
if model_type == MODEL_TYPE_MINIMAX_H3:
return [
"attn.out_proj", "mlp.fc2",
]
return []
pyproject.toml
+1 -1
View File
@@ -2,7 +2,7 @@
[project]
name = "comfyui-int8-fast-fork" # Immutable Comfy Registry ID; repository/display names may differ.
description = "Quantization toolkit for ComfyUI with native INT4/INT8 conversion, mixed-precision loading, LoRA handling, and low-VRAM runtime controls."
version = "2.0.0"
version = "2.1.0"
license = { file = "LICENSE" }
dependencies = []
keywords = ["comfyui", "quantization", "int4", "int8", "w4a4", "w8a8"]
tests/test_quantization_modes.py
+119
View File
@@ -87,6 +87,100 @@ class QuantizationModeTests(unittest.TestCase):
).to(base_output.dtype)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.02)
def test_native_int8_convrot_metadata_round_trip_preserves_layer_group_size(self):
state_dict = {
"weight": torch.randint(-8, 9, (16, 2688), dtype=torch.int8),
"weight_scale": torch.ones((16, 1), dtype=torch.float32),
"comfy_quant": quant._encode_comfy_quant_config({
"format": "int8_tensorwise",
"convrot": True,
"convrot_groupsize": 64,
}),
}
module = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
module.load_state_dict(state_dict, strict=False)
self.assertEqual(module._convrot_groupsize, 64)
self.assertEqual(tuple(module.quarot_hadamard.shape), (64, 64))
self.assertEqual(module._outlier_method, quant.OUTLIER_METHOD_CONVROT)
with mock.patch.object(
quant,
"_native_int8_convrot_linear",
return_value=torch.zeros((2, 16), dtype=torch.bfloat16),
) as native_linear:
module(torch.randn(2, 2688, dtype=torch.float32))
native_linear.assert_called_once()
self.assertEqual(native_linear.call_args.args[-1], 64)
saved_state_dict = module.state_dict()
saved_config = quant._decode_comfy_quant_config(saved_state_dict["comfy_quant"])
self.assertEqual(saved_config["convrot_groupsize"], 64)
self.assertNotIn("quarot_hadamard", saved_state_dict)
reloaded = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
reloaded.load_state_dict(saved_state_dict, strict=False)
self.assertEqual(reloaded._convrot_groupsize, 64)
self.assertEqual(tuple(reloaded.quarot_hadamard.shape), (64, 64))
def test_native_int8_convrot_load_rejects_invalid_group_metadata(self):
for group_size, message in (
(0, "must be positive"),
(256, "2688 is not divisible"),
(32, "must be a power of 4"),
("invalid", "invalid ConvRot group size"),
):
with self.subTest(group_size=group_size):
module = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
state_dict = {
"weight": torch.randint(-8, 9, (16, 2688), dtype=torch.int8),
"weight_scale": torch.ones((16, 1), dtype=torch.float32),
"comfy_quant": quant._encode_comfy_quant_config({
"format": "int8_tensorwise",
"convrot": True,
"convrot_groupsize": group_size,
}),
}
with self.assertRaisesRegex(ValueError, message):
module.load_state_dict(state_dict, strict=False)
def test_native_int8_convrot_save_rejects_inconsistent_hadamard_shape(self):
module = quant.Int8TensorwiseOps.Linear(256, 16, bias=False, device="cpu", dtype=torch.float32)
module.weight = nn.Parameter(torch.randint(-8, 9, (16, 256), dtype=torch.int8), requires_grad=False)
module.weight_scale = torch.ones((16, 1), dtype=torch.float32)
module._is_quantized = True
module._quant_format = "int8_tensorwise"
module._outlier_method = quant.OUTLIER_METHOD_CONVROT
module._convrot_groupsize = 256
module.quarot_hadamard = quant._build_outlier_hadamard(
quant.OUTLIER_METHOD_CONVROT,
64,
device="cpu",
dtype=torch.float32,
)
with self.assertRaisesRegex(ValueError, "Hadamard shape must be"):
module.state_dict()
def test_native_int8_convrot_save_rejects_nondivisible_group_size(self):
module = quant.Int8TensorwiseOps.Linear(2688, 16, bias=False, device="cpu", dtype=torch.float32)
module.weight = nn.Parameter(torch.randint(-8, 9, (16, 2688), dtype=torch.int8), requires_grad=False)
module.weight_scale = torch.ones((16, 1), dtype=torch.float32)
module._is_quantized = True
module._quant_format = "int8_tensorwise"
module._outlier_method = quant.OUTLIER_METHOD_CONVROT
module._convrot_groupsize = 256
module.quarot_hadamard = quant._build_outlier_hadamard(
quant.OUTLIER_METHOD_CONVROT,
256,
device="cpu",
dtype=torch.float32,
)
with self.assertRaisesRegex(ValueError, "2688 is not divisible"):
module.state_dict()
def test_dynamic_lora_hook_supports_baked_native_int8_linear(self):
class DiffusionModel(nn.Module):
def __init__(self):
@@ -396,6 +490,7 @@ class QuantizationModeTests(unittest.TestCase):
def test_architecture_presets_protect_residual_writeback_paths(self):
krea_preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_KREA2)
anima_preset = loader.get_model_type_quantization_preset("anima")
minimax_preset = loader.get_model_type_quantization_preset(loader.MODEL_TYPE_MINIMAX_H3)
self.assertIn("first", krea_preset["keep_float"])
self.assertEqual(krea_preset["int4_sensitive"], ("attn.wo", "mlp.down"))
@@ -404,6 +499,30 @@ class QuantizationModeTests(unittest.TestCase):
anima_preset["int4_sensitive"],
("self_attn.output_proj", "cross_attn.output_proj", "mlp.layer2"),
)
self.assertEqual(
minimax_preset["keep_float"],
(
"video_patch_proj", "audio_patch_proj", "condition_proj", "time_embedder",
"token_refiner", "adaln_proj", "final_layer",
),
)
self.assertEqual(minimax_preset["int4_sensitive"], ("attn.out_proj", "mlp.fc2"))
def test_model_adapter_auto_detects_minimax_h3(self):
class MiniMaxH3(nn.Module):
def __init__(self):
super().__init__()
self.video_patch_proj = nn.Linear(4, 4)
self.audio_patch_proj = nn.Linear(4, 4)
self.condition_proj = nn.Linear(4, 4)
self.token_refiner = nn.Sequential(nn.Linear(4, 4))
self.blocks = nn.ModuleList([nn.Module()])
self.blocks[0].adaln_proj = nn.Sequential(nn.Linear(4, 4))
self.assertEqual(
model_adapter._infer_model_type_from_modules(MiniMaxH3()),
loader.MODEL_TYPE_MINIMAX_H3,
)
def test_architecture_ratios_cover_residual_writeback_paths(self):
krea_names = [