From d1af3c89233c475fc1bf11939d86787e6cddd58c Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 23 Aug 2026 22:13:59 +0800 Subject: [PATCH] [diffusion] feat: support loading native diffusers miniMax h3 components (#36067) --- .../configs/models/dits/minimax_h3.py | 65 +++++++++++++------ .../component_loaders/transformer_loader.py | 13 ++-- .../runtime/models/dits/minimax_h3.py | 51 +++++++++++++++ .../test/unit/test_minimax_h3_dit_contract.py | 26 ++++++++ 4 files changed, 130 insertions(+), 25 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py index 4751ff961..376fd87b2 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py @@ -16,51 +16,59 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): r"^(.*\.lora_[AB])\.[^.]+$": r"\1", r"^base_model\.model\.(.*\.lora_[AB])$": r"\1", r"^transformer\.(.*\.lora_[AB])$": r"\1", - r"^proj_in\.(lora_[AB])$": r"video_patch_proj.\1", - r"^audio_proj_in\.(lora_[AB])$": r"audio_patch_proj.\1", - r"^context_embedder\.(lora_[AB])$": r"condition_proj.\1", - r"^time_embedder\.linear_1\.(lora_[AB])$": r"time_embedder.proj_in.\1", - r"^time_embedder\.linear_2\.(lora_[AB])$": r"time_embedder.proj_out.\1", - r"^norm_out\.linear\.(lora_[AB])$": r"final_layer.adaln_proj.linear.\1", - r"^proj_out\.(lora_[AB])$": r"final_layer.video_out.\1", - r"^audio_proj_out\.(lora_[AB])$": r"final_layer.audio_out.\1", - r"^transformer_blocks\.(\d+)\.adaln_proj\.linear\.(lora_[AB])$": r"blocks.\1.adaln_proj.linear.\2", - r"^transformer_blocks\.(\d+)\.attn\.to_q\.(lora_[AB])$": ( + r"^proj_in\.(.*)$": r"video_patch_proj.\1", + r"^audio_proj_in\.(.*)$": r"audio_patch_proj.\1", + r"^context_embedder\.(.*)$": r"condition_proj.\1", + r"^time_embedder\.linear_1\.(.*)$": r"time_embedder.proj_in.\1", + r"^time_embedder\.linear_2\.(.*)$": r"time_embedder.proj_out.\1", + r"^norm_out\.norm\.(.*)$": r"final_layer.norm.\1", + r"^norm_out\.linear\.(.*)$": r"final_layer.adaln_proj.linear.\1", + r"^proj_out\.(.*)$": r"final_layer.video_out.\1", + r"^audio_proj_out\.(.*)$": r"final_layer.audio_out.\1", + r"^transformer_blocks\.(\d+)\.adaln_proj\.linear\.(.*)$": r"blocks.\1.adaln_proj.linear.\2", + r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( r"blocks.\1.attn.qkv_proj.\2", 0, 3, ), - r"^transformer_blocks\.(\d+)\.attn\.to_k\.(lora_[AB])$": ( + r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( r"blocks.\1.attn.qkv_proj.\2", 1, 3, ), - r"^transformer_blocks\.(\d+)\.attn\.to_v\.(lora_[AB])$": ( + r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( r"blocks.\1.attn.qkv_proj.\2", 2, 3, ), - r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(lora_[AB])$": r"blocks.\1.attn.out_proj.\2", - r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(lora_[AB])$": r"blocks.\1.mlp.fc1.\2", - r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(lora_[AB])$": r"blocks.\1.mlp.fc2.\2", - r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(lora_[AB])$": ( + r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"blocks.\1.attn.out_proj.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"blocks.\1.attn.q_norm.\2", + r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"blocks.\1.attn.k_norm.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"blocks.\1.mlp.fc1.\2", + r"^transformer_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"blocks.\1.mlp.fc2.\2", + r"^transformer_blocks\.(\d+)\.norm([12])\.(.*)$": r"blocks.\1.norm\2.\3", + r"^token_refiner\.final_norm\.(.*)$": r"token_refiner.final_norm.\1", + r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$": ( r"token_refiner.blocks.\1.attn.qkv_proj.\2", 0, 3, ), - r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(lora_[AB])$": ( + r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$": ( r"token_refiner.blocks.\1.attn.qkv_proj.\2", 1, 3, ), - r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(lora_[AB])$": ( + r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$": ( r"token_refiner.blocks.\1.attn.qkv_proj.\2", 2, 3, ), - r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(lora_[AB])$": r"token_refiner.blocks.\1.attn.out_proj.\2", - r"^token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0\.proj\.(lora_[AB])$": r"token_refiner.blocks.\1.mlp.fc1.\2", - r"^token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2\.(lora_[AB])$": r"token_refiner.blocks.\1.mlp.fc2.\2", + r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"token_refiner.blocks.\1.attn.out_proj.\2", + r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"token_refiner.blocks.\1.attn.q_norm.\2", + r"^token_refiner\.refiner_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"token_refiner.blocks.\1.attn.k_norm.\2", + r"^token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"token_refiner.blocks.\1.mlp.fc1.\2", + r"^token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2\.(.*)$": r"token_refiner.blocks.\1.mlp.fc2.\2", + r"^token_refiner\.refiner_blocks\.(\d+)\.norm([12])\.(.*)$": r"token_refiner.blocks.\1.norm\2.\3", } ) @@ -85,6 +93,7 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): norm_eps: float = 1e-5 qk_norm_eps: float = 1e-5 final_norm_eps: float = 1e-5 + checkpoint_uses_diffusers_layout: bool = False def __post_init__(self) -> None: super().__post_init__() @@ -99,6 +108,20 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): class MiniMaxH3DiTConfig(DiTConfig): arch_config: MiniMaxH3DiTArchConfig = field(default_factory=MiniMaxH3DiTArchConfig) + def update_model_arch(self, source_model_dict: dict) -> None: + aliases = { + "num_refiner_layers": "token_refiner_num_layers", + "ffn_dim": "ffn_hidden_size", + "in_channels": "latents_dim", + "audio_in_channels": "audio_latents_dim", + "freq_dim": "timestep_input_dim", + "time_embed_hidden_dim": "time_embed_hidden_size", + "rope_freq_dim": "rope_inv_freq_len", + } + super().update_model_arch( + {aliases.get(key, key): value for key, value in source_model_dict.items()} + ) + __all__ = [ "MINIMAX_H3_ADALN_MODALITY_NUM", diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 6422f2da3..55d4ffee9 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -219,9 +219,14 @@ class TransformerLoader(ComponentLoader): cls_name = config.pop("_class_name") model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) + is_minimax_h3 = model_cls.__name__ == "MiniMaxH3DiTModel" + if is_minimax_h3: + dit_config.arch_config.checkpoint_uses_diffusers_layout = ( + cls_name == "MiniMaxH3Transformer3DModel" + ) checkpoint_quant_config = None - if cls_name == "MiniMaxH3DiTModel": + if is_minimax_h3: selected_variant = str(component_server_args.model_variant or "fl2va") if gguf_file is not None: validate_minimax_h3_checkpoint_variant([gguf_file], selected_variant) @@ -260,7 +265,7 @@ class TransformerLoader(ComponentLoader): gguf_file=gguf_file, checkpoint_quant_config=checkpoint_quant_config, ) - if quant_spec.gguf_file is not None and cls_name == "MiniMaxH3DiTModel": + if quant_spec.gguf_file is not None and is_minimax_h3: assert quant_spec.quant_config is not None curve = quant_spec.quant_config.tensor_meta.get("adaln_t_table") if curve is not None: @@ -312,7 +317,7 @@ class TransformerLoader(ComponentLoader): ) adaln_cache_path = component_server_args.minimax_h3_adaln_cache_path if adaln_cache_path is not None: - if cls_name != "MiniMaxH3DiTModel": + if not is_minimax_h3: raise ValueError( "--minimax-h3-adaln-cache-path is only supported by MiniMax H3" ) @@ -326,7 +331,7 @@ class TransformerLoader(ComponentLoader): ) checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter if component_server_args.minimax_h3_adaln_online: - if cls_name != "MiniMaxH3DiTModel": + if not is_minimax_h3: raise ValueError( "--minimax-h3-adaln-online is only supported by MiniMax H3" ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 808600ac1..8f190f54d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -10,6 +10,8 @@ from __future__ import annotations import math import os import struct +from collections import defaultdict +from collections.abc import Iterable, Iterator from contextlib import ExitStack from typing import Any, Callable @@ -58,6 +60,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.usp import _ring_attention_varlen +from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, is_layerwise_offloaded_module, @@ -75,6 +78,48 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import logger = init_logger(__name__) _ARCH_DEFAULTS = MiniMaxH3DiTArchConfig() + + +def _diffusers_h3_checkpoint( + iterator: Iterable[tuple[str, torch.Tensor]], +) -> Iterator[tuple[str, torch.Tensor]]: + """Map Diffusers H3 names/layout to the fused native checkpoint layout.""" + mapping = get_param_names_mapping(_ARCH_DEFAULTS.param_names_mapping) + pending: dict[str, dict[int, torch.Tensor]] = defaultdict(dict) + + for source_name, tensor in iterator: + target_name, merge_index, merge_count = mapping(source_name) + + # Diffusers SwiGLU stores [value, gate]; the native fused MLP consumes + # [gate, value]. Packed GPTQ tensors carry output channels on dim 1. + if ".ff.net.0.proj." in source_name: + output_dim = ( + 1 if source_name.endswith((".qweight", ".qzeros", ".scales")) else 0 + ) + value, gate = tensor.chunk(2, dim=output_dim) + tensor = torch.cat((gate, value), dim=output_dim) + + if merge_index is None: + yield target_name, tensor + continue + + assert merge_count is not None + pending[target_name][merge_index] = tensor + if len(pending[target_name]) != merge_count: + continue + + merge_dim = 1 if target_name.endswith((".qweight", ".qzeros", ".scales")) else 0 + yield target_name, torch.cat( + [pending[target_name][index] for index in range(merge_count)], + dim=merge_dim, + ) + del pending[target_name] + + if pending: + incomplete = ", ".join(sorted(pending)) + raise ValueError(f"Incomplete Diffusers H3 fused parameters: {incomplete}") + + _BF16_DTYPE = torch.bfloat16 _FP32_DTYPE = torch.float32 _MPS_MLP_TOKEN_CHUNK_SIZE = 128 @@ -618,6 +663,9 @@ class MiniMaxH3Attention(nn.Module): quant_config.get_name() == "gguf" or quant_config.checkpoint_uses_native_qkv_layout ) + checkpoint_qkv_is_native = ( + checkpoint_qkv_is_native or arch.checkpoint_uses_diffusers_layout + ) if not checkpoint_qkv_is_native: self._install_qkv_weight_loader(arch) self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps) @@ -1605,6 +1653,7 @@ class MiniMaxH3FinalLayer(nn.Module): class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): + _aliases = ["MiniMaxH3Transformer3DModel"] _fsdp_shard_conditions = [is_block] # refine_prompt_embeds drives a forward pass outside __call__. _fsdp_forward_methods = ("refine_prompt_embeds",) @@ -1735,6 +1784,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): adaln_cache_path is not None or adaln_weight_files is not None ) self.arch = arch + if arch.checkpoint_uses_diffusers_layout: + self.preprocess_loaded_state_dict = _diffusers_h3_checkpoint self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.latents_dim diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py index c0be860ae..de3d8f3d6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py @@ -28,6 +28,7 @@ from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( MiniMaxH3DiTBlock, MiniMaxH3DiTModel, _copy_grouped_qkv_tp_shard, + _diffusers_h3_checkpoint, _modulate_gate, _reorder_grouped_qkv_to_qkv, ) @@ -84,6 +85,31 @@ def test_native_weight_names_and_grouped_qkv_reorder(): None, ) + assert mapping("transformer_blocks.7.attn.to_k.qweight") == ( + "blocks.7.attn.qkv_proj.qweight", + 1, + 3, + ) + + diffusers_weights = [ + ("transformer_blocks.0.attn.to_q.qweight", torch.full((2, 3), 1)), + ("transformer_blocks.0.attn.to_k.qweight", torch.full((2, 3), 2)), + ("transformer_blocks.0.attn.to_v.qweight", torch.full((2, 3), 3)), + ( + "transformer_blocks.0.ff.net.0.proj.weight", + torch.arange(8).reshape(4, 2), + ), + ] + converted = dict(_diffusers_h3_checkpoint(diffusers_weights)) + assert torch.equal( + converted["blocks.0.attn.qkv_proj.qweight"], + torch.cat([tensor for _, tensor in diffusers_weights[:3]], dim=1), + ) + assert torch.equal( + converted["blocks.0.mlp.fc1.weight"], + torch.tensor([[4, 5], [6, 7], [0, 1], [2, 3]]), + ) + weight = torch.arange(12, dtype=torch.float32).reshape(12, 1) actual = _reorder_grouped_qkv_to_qkv( weight,