Merge pull request #3202 from bghira/fix/minimax-h3-mixed-convrot-groups

Allow mixed ConvRot groups in MiniMax H3 checkpoints

Author: bghira

simpletuner/helpers/models/minimaxh3/transformer.py

@@ -69,6 +69,16 @@
 _H3_MASKED_CONTEXT_PARALLEL_BACKENDS = frozenset({AttentionBackendName.NATIVE, AttentionBackendName._NATIVE_CUDNN})
 
 
+def _linear_compute_dtype(linear: nn.Module) -> torch.dtype:
+    compute_dtype = getattr(linear, "compute_dtype", None)
+    if isinstance(compute_dtype, torch.dtype):
+        return compute_dtype
+    weight = linear.weight
+    dequantizer = getattr(weight, "sdnq_dequantizer", None)
+    result_dtype = getattr(dequantizer, "result_dtype", None)
+    return result_dtype if isinstance(result_dtype, torch.dtype) else weight.dtype
+
+
 class _MiniMaxH3AllGather(torch.autograd.Function):
     """Gather sequence shards without PyTorch's unsupported NCCL coalesced path."""
 
@@ -440,7 +450,13 @@ def _infer_minimax_h3_config_from_checkpoint(checkpoint) -> dict[str, Any]:
         audio_weight = _get_checkpoint_tensor(checkpoint, "audio_proj_in.weight")
         context_weight = _get_checkpoint_tensor(checkpoint, "context_embedder.weight")
         q_norm_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.norm_q.weight")
-        q_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.to_q.weight")
+        if "transformer_blocks.0.attn.to_q.weight" in raw_keys:
+            q_output_dim = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.to_q.weight").shape[0]
+        else:
+            qkv_weight = _get_checkpoint_tensor(checkpoint, "blocks.0.attn.qkv_proj.weight")
+            if qkv_weight.shape[0] % 3 != 0:
+                raise RuntimeError("MiniMax-H3 fused QKV tensor blocks.0.attn.qkv_proj.weight cannot be split into q/k/v")
+            q_output_dim = qkv_weight.shape[0] // 3
         ffn_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.ff.net.0.proj.weight")
         has_adaln_curve = "adaln_t_table" in raw_keys
         adaln_curve_table = _get_checkpoint_tensor(checkpoint, "adaln_t_table") if has_adaln_curve else None
@@ -456,7 +472,7 @@ def _infer_minimax_h3_config_from_checkpoint(checkpoint) -> dict[str, Any]:
             "audio_in_channels": audio_weight.shape[1],
             "text_dim": context_weight.shape[1],
             "attention_head_dim": q_norm_weight.shape[0],
-            "num_attention_heads": q_weight.shape[0] // q_norm_weight.shape[0],
+            "num_attention_heads": q_output_dim // q_norm_weight.shape[0],
             "freq_dim": time_in.shape[1] if time_in is not None else 256,
             "time_embed_hidden_dim": time_in.shape[0] if time_in is not None else 5376,
             "time_embed_dim": adaln_curve_table.shape[1] if has_adaln_curve else time_out.shape[0],
@@ -631,7 +647,7 @@ def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]:
         # The activation runs at `temb`'s own precision and only the projection input is aligned to the projection
         # weight. Every block reads the same `temb`, so early rounding biases every block's modulation coherently.
         temb = nn.functional.silu(temb) if self.apply_silu else temb
-        temb = self.linear(temb.to(self.linear.weight.dtype))
+        temb = self.linear(temb.to(_linear_compute_dtype(self.linear)))
         temb = temb.view(-1, 6 * self.hidden_size)
         return temb.chunk(6, dim=-1)
 
@@ -661,7 +677,7 @@ def forward(
     ) -> torch.Tensor:
         # As in `MiniMaxH3AdaLayerNormModulation`: activate at `temb`'s precision, cast to the projection's dtype after.
         temb = nn.functional.silu(temb) if self.apply_silu else temb
-        shift, scale = self.linear(temb.to(self.linear.weight.dtype)).chunk(2, dim=-1)
+        shift, scale = self.linear(temb.to(_linear_compute_dtype(self.linear))).chunk(2, dim=-1)
         activation_dtype = hidden_states.dtype
         hidden_states = self.norm(hidden_states)
         shift = _select_modulation(shift, timestep_indices).to(dtype=activation_dtype)
@@ -1522,7 +1538,7 @@ def _time_embedding(
                 temb = blend_flowmap_embeddings(temb, delta_temb, self.flowmap_delta_emb_gate)
             return temb
 
-        dtype = self.time_embedder.linear_1.weight.dtype
+        dtype = _linear_compute_dtype(self.time_embedder.linear_1)
         temb = flowmap_timestep_embedding(
             time_proj=self.time_proj,
             timestep_embedder=self.time_embedder,
@@ -1857,21 +1873,20 @@ def from_single_file(
                     result_dtype=torch_dtype or torch.bfloat16,
                     hadamard_group_size=hadamard_group_size,
                 )
-            if len(hadamard_group_sizes) != 1:
-                raise RuntimeError(
-                    f"MiniMax-H3 ConvRot checkpoint uses multiple Hadamard group sizes: {sorted(hadamard_group_sizes)}"
-                )
-            group_size = hadamard_group_sizes.pop()
             model.quantization_method = "minimax_h3_comfy_convrot_sdnq"
             model.quantization_config = {
                 "quant_method": "sdnq_training",
                 "weights_dtype": "int8",
                 "quantized_matmul_dtype": "int8",
                 "use_hadamard": True,
-                "hadamard_group_size": group_size,
                 "group_size": -1,
                 "source_format": "comfy_minimax_h3_convrot",
             }
+            sorted_group_sizes = sorted(hadamard_group_sizes)
+            if len(sorted_group_sizes) == 1:
+                model.quantization_config["hadamard_group_size"] = sorted_group_sizes[0]
+            else:
+                model.quantization_config["hadamard_group_sizes"] = sorted_group_sizes
         elif fp8_state_dict:
             model.quantization_method = "minimax_h3_comfy_fp8"
             model.quantization_config = {
@@ -2022,9 +2037,9 @@ def forward(
         # mixed-precision (the two patch projections are float32 while `context_embedder` and the block stack are
         # bfloat16 — see `_keep_in_fp32_modules`), so every input is aligned with its projection's parameter dtype,
         # mirroring the reference's explicit casts. The text stream sets the dtype of the packed sequence.
-        video_embeds = self.proj_in(hidden_states.to(self.proj_in.weight.dtype))
-        audio_embeds = self.audio_proj_in(audio_hidden_states.to(self.audio_proj_in.weight.dtype))
-        text_embeds = self.context_embedder(encoder_hidden_states.to(self.context_embedder.weight.dtype))
+        video_embeds = self.proj_in(hidden_states.to(_linear_compute_dtype(self.proj_in)))
+        audio_embeds = self.audio_proj_in(audio_hidden_states.to(_linear_compute_dtype(self.audio_proj_in)))
+        text_embeds = self.context_embedder(encoder_hidden_states.to(_linear_compute_dtype(self.context_embedder)))
         self.token_refiner.gradient_checkpointing = self.gradient_checkpointing
         text_attention_mask = None
         if packed_valid_mask is not None:
@@ -2375,7 +2390,7 @@ def run_checkpointed_block(
         # 5. Both heads run over every row, then the rows of each modality are selected. The heads are listed in
         # `_keep_in_fp32_modules`, so they stay float32 while the block stack runs in the requested `torch_dtype`;
         # align the activation with their parameter dtype.
-        hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(self.proj_out.weight.dtype)
+        hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(_linear_compute_dtype(self.proj_out))
         video_output = _gather_h3_context_parallel_output(self.proj_out(hidden_states), cp_config, dim=1).index_select(
             1, video_indices.to(hidden_states.device)
         )

simpletuner/helpers/models/z_image/quantized_loading.py

@@ -200,7 +200,9 @@ def _wrap_convrot_linear(
         True,
         -1,
     )
-    _set_module(model, module_name, get_sdnq_wrapper_class(module, forward))
+    wrapped_module = get_sdnq_wrapper_class(module, forward)
+    wrapped_module.compute_dtype = result_dtype
+    _set_module(model, module_name, wrapped_module)
 
 
 def _validate_quant_metadata(checkpoint, key: str) -> int:

tests/test_minimaxh3.py

@@ -64,6 +64,7 @@
     _convert_minimax_h3_native_swiglu_scale_to_diffusers,
     _convert_minimax_h3_native_swiglu_to_diffusers,
     _gather_h3_context_parallel_output,
+    _linear_compute_dtype,
     _pad_h3_context_parallel_layout,
     resolve_h3_reference_mode,
 )
@@ -768,6 +769,17 @@ def sample(self, generator=None):
 
 
 class MiniMaxH3Tests(unittest.TestCase):
+    def test_linear_compute_dtype_uses_quantized_result_dtype(self):
+        weight = SimpleNamespace(
+            dtype=torch.int8,
+            sdnq_dequantizer=SimpleNamespace(result_dtype=torch.bfloat16),
+        )
+        self.assertEqual(_linear_compute_dtype(SimpleNamespace(weight=weight)), torch.bfloat16)
+
+    def test_linear_compute_dtype_prefers_module_contract(self):
+        linear = SimpleNamespace(weight=SimpleNamespace(dtype=torch.int8), compute_dtype=torch.float16)
+        self.assertEqual(_linear_compute_dtype(linear), torch.float16)
+
     def test_registry_metadata_resolves(self):
         model_cls = ModelRegistry.get("minimaxh3")
         self.assertEqual(model_cls.NAME, "MiniMax H3")
@@ -3655,6 +3667,47 @@ def test_single_file_loader_accepts_abiray_convrot_metadata(self):
         self.assertEqual(wrap_convrot.call_args.args[1], "transformer_blocks.0.attn.to_out.0")
         self.assertEqual(wrap_convrot.call_args.kwargs["hadamard_group_size"], 256)
 
+    def test_single_file_loader_accepts_mixed_convrot_group_sizes(self):
+        model = tiny_h3_transformer(num_layers=1)
+        state_dict = dict(model.state_dict())
+
+        qkv_weights = [state_dict.pop(f"transformer_blocks.0.attn.to_{branch}.weight") for branch in ("q", "k", "v")]
+        qkv_source = "blocks.0.attn.qkv_proj"
+        qkv_weight = torch.cat(qkv_weights, dim=0)
+        state_dict[f"{qkv_source}.weight"] = torch.zeros(qkv_weight.shape, dtype=torch.int8)
+        state_dict[f"{qkv_source}.weight_scale"] = torch.ones(qkv_weight.shape[0], 1, dtype=torch.float32)
+        state_dict[f"{qkv_source}.comfy_quant"] = comfy_quant_metadata_tensor(
+            {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 64}
+        )
+
+        out_target = "transformer_blocks.0.attn.to_out.0.weight"
+        out_weight = state_dict.pop(out_target)
+        out_source = "blocks.0.attn.out_proj"
+        state_dict[f"{out_source}.weight"] = torch.zeros(out_weight.shape, dtype=torch.int8)
+        state_dict[f"{out_source}.weight_scale"] = torch.ones(out_weight.shape[0], 1, dtype=torch.float32)
+        state_dict[f"{out_source}.comfy_quant"] = comfy_quant_metadata_tensor(
+            {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256}
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = f"{tmpdir}/tiny-h3-mixed-convrot.safetensors"
+            save_file(state_dict, path)
+            with patch("simpletuner.helpers.models.z_image.quantized_loading._wrap_convrot_linear") as wrap_convrot:
+                loaded = MiniMaxH3Transformer3DModel.from_single_file(path, torch_dtype=torch.float32)
+
+        group_sizes_by_module = {call.args[1]: call.kwargs["hadamard_group_size"] for call in wrap_convrot.call_args_list}
+        self.assertEqual(
+            group_sizes_by_module,
+            {
+                "transformer_blocks.0.attn.to_q": 64,
+                "transformer_blocks.0.attn.to_k": 64,
+                "transformer_blocks.0.attn.to_v": 64,
+                "transformer_blocks.0.attn.to_out.0": 256,
+            },
+        )
+        self.assertEqual(loaded.quantization_config["hadamard_group_sizes"], [64, 256])
+        self.assertNotIn("hadamard_group_size", loaded.quantization_config)
+
     def test_single_file_loader_accepts_comfy_fp8_scale_metadata(self):
         model = tiny_h3_transformer(num_layers=1)
         state_dict = dict(model.state_dict())