diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f1322d..525bcb1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,20 @@ # Changelog -
2.0.0 +
2.1.0 + +### 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. + +
+ +
2.0.0 ### Added diff --git a/int8_model_adapter.py b/int8_model_adapter.py index 5744416..7df1919 100644 --- a/int8_model_adapter.py +++ b/int8_model_adapter.py @@ -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"), diff --git a/int8_quant.py b/int8_quant.py index dc5fb03..c335d92 100644 --- a/int8_quant.py +++ b/int8_quant.py @@ -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 ''}", + ) 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 ''}", + ) 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 ''}: 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 ''}", + ) 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 ''}", + ) + if not isinstance(hadamard, torch.Tensor) or tuple(hadamard.shape) != (convrot_group_size, convrot_group_size): + raise ValueError( + f"Quantized layer {prefix.rstrip('.') or ''}: 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) diff --git a/int8_unet_loader.py b/int8_unet_loader.py index cd65a6b..67d41b2 100644 --- a/int8_unet_loader.py +++ b/int8_unet_loader.py @@ -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 [] diff --git a/pyproject.toml b/pyproject.toml index 0afcbec..d82aa07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"] diff --git a/tests/test_quantization_modes.py b/tests/test_quantization_modes.py index 703d21b..59e0084 100644 --- a/tests/test_quantization_modes.py +++ b/tests/test_quantization_modes.py @@ -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 = [