From 939c00a7e306a2265602f6febca946e2c84e065a Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 23 Aug 2026 21:25:11 +0800 Subject: [PATCH] [diffusion] feat: support compact qwen3-vl conditioning for minimax h3 (#36076) --- .../models/encoders/minimax_h3_qwen3vl.py | 3 + .../component_loaders/text_encoder_loader.py | 5 + .../runtime/models/encoders/base.py | 8 + .../models/encoders/minimax_h3_qwen3vl.py | 229 ++++++++++++++++-- .../test/unit/test_server_args.py | 16 ++ .../test/unit/test_text_encoder_loader.py | 77 +++++- 6 files changed, 322 insertions(+), 16 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py index 62d9d3bdc..6148731c8 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/minimax_h3_qwen3vl.py @@ -26,6 +26,8 @@ class MiniMaxH3Qwen3VLArchConfig(Qwen3VLArchConfig): head_dim: int = 128 text_len: int = 262144 hidden_state_skip_layer: int = 0 + checkpoint_num_hidden_layers: int = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER + conditioning_projection_path: str | None = None @dataclass @@ -38,6 +40,7 @@ class MiniMaxH3Qwen3VLConfig(Qwen3VLConfig): """Select the in-tree extractor after loading the HF architecture.""" arch = self.arch_config + arch.checkpoint_num_hidden_layers = int(arch.text_config.num_hidden_layers) arch.architectures = ["MiniMaxH3Qwen3VLEncoder"] arch.hidden_size = int(arch.text_config.hidden_size) arch.intermediate_size = int(arch.text_config.intermediate_size) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 0103f8df2..c690e5ea7 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -600,6 +600,11 @@ class TextEncoderLoader(ComponentLoader): component_weights_path, component_name, ) + if issubclass(model_cls, EncoderTensorParallelMixin): + model_cls.configure_component_paths( + encoder_config, + server_args.component_paths, + ) encoder_dp_group = get_encoder_data_parallel_group() prefer_dp = ( server_args.batching_max_size > 1 diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py index f58ee332a..1f5147d02 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -162,6 +162,14 @@ class EncoderTensorParallelMixin: # states or sharding contract cannot use the generic loader lifecycle. manages_checkpoint_quantization = False + @classmethod + def configure_component_paths( + cls, + config: EncoderConfig, + component_paths: dict[str, str], + ) -> None: + """Apply optional runtime components before parallel layout is resolved.""" + def bind_encoder_tp_group(self, tp_group: GroupCoordinator) -> None: self._encoder_tp_group = tp_group diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py index 7c1a31f87..bec9ffb83 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Native, TP-foldable Qwen3-VL layer-50 encoder for MiniMax H3.""" +"""Native, TP-foldable Qwen3-VL conditioning encoder for MiniMax H3.""" from __future__ import annotations @@ -9,6 +9,9 @@ from typing import Any import torch import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open +from safetensors.torch import load_file from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import ( @@ -20,6 +23,10 @@ from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel +from sglang.multimodal_gen.runtime.weights.source import ( + materialize_weight, + resolve_weight, +) MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 _LAYER_WEIGHT_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.") @@ -35,17 +42,155 @@ def _map_checkpoint_name(name: str) -> str: return _MAP_CHECKPOINT_NAME(name)[0] -def _is_unconsumed_checkpoint_weight(name: str) -> bool: - """Weights intentionally absent from the layer-50 feature extractor.""" +def _is_unconsumed_checkpoint_weight(name: str, selected_layer: int) -> bool: + """Weights intentionally absent from the selected feature extractor.""" if name == "lm_head.weight" or name.startswith("model.language_model.norm."): return True match = _LAYER_WEIGHT_RE.match(name) - return bool(match and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER) + return bool(match and int(match.group(1)) >= selected_layer) + + +class _FrozenLinear(nn.Module): + def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None) -> None: + super().__init__() + self.register_buffer("weight", weight) + self.register_buffer("bias", bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.linear(hidden_states, self.weight, self.bias) + + +class MiniMaxH3ConditioningProjection(nn.Module): + """Apply the safe ClipProj conditioning format without a custom runtime.""" + + _REQUIRED_TENSORS = ("mean_in", "std_in", "mean_out", "std_out") + + @staticmethod + def inspect(path: str) -> tuple[int, int, int]: + if not path.endswith(".safetensors"): + raise ValueError("H3 conditioning projections must use safetensors") + with safe_open(path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + try: + tap = int(metadata["tap"]) + input_dim = int(handle.get_slice("mean_in").get_shape()[0]) + output_dim = int(handle.get_slice("mean_out").get_shape()[0]) + except (IndexError, KeyError, TypeError, ValueError) as error: + raise ValueError( + f"Invalid H3 conditioning projection metadata in {path!r}" + ) from error + return tap, input_dim, output_dim + + def __init__(self, path: str) -> None: + super().__init__() + self.tap, self.input_dim, self.output_dim = self.inspect(path) + tensors = load_file(path, device="cpu") + missing = set(self._REQUIRED_TENSORS) - set(tensors) + if missing: + raise ValueError( + f"H3 conditioning projection is missing tensors: {sorted(missing)}" + ) + + for name in self._REQUIRED_TENSORS: + self.register_buffer(name, tensors.pop(name).float()) + expected_shapes = { + "mean_in": (self.input_dim,), + "std_in": (self.input_dim,), + "mean_out": (self.output_dim,), + "std_out": (self.output_dim,), + } + for name, expected_shape in expected_shapes.items(): + actual_shape = tuple(self.get_buffer(name).shape) + if actual_shape != expected_shape: + raise ValueError( + f"H3 conditioning projection {name} has shape " + f"{actual_shape}, expected {expected_shape}" + ) + self.register_buffer( + "weight", + tensors.pop("W").float() if "W" in tensors else None, + ) + self.register_buffer( + "sink_out", + tensors.pop("sink_out").float() if "sink_out" in tensors else None, + ) + + layer_indices = sorted( + { + int(name.split(".")[1]) + for name in tensors + if re.fullmatch(r"mlp\.\d+\.weight", name) + } + ) + layers: list[nn.Module] = [] + layer_input_dim = self.input_dim + for layer_index in layer_indices: + weight_name = f"mlp.{layer_index}.weight" + bias_name = f"mlp.{layer_index}.bias" + weight = tensors.pop(weight_name) + bias = tensors.pop(bias_name, None) + if weight.ndim != 2 or int(weight.shape[1]) != layer_input_dim: + raise ValueError( + f"H3 conditioning projection {weight_name} cannot follow " + f"width {layer_input_dim}: got {tuple(weight.shape)}" + ) + if bias is not None and tuple(bias.shape) != (int(weight.shape[0]),): + raise ValueError( + f"H3 conditioning projection {bias_name} has shape " + f"{tuple(bias.shape)}, expected ({int(weight.shape[0])},)" + ) + layers.append(_FrozenLinear(weight, bias)) + layer_input_dim = int(weight.shape[0]) + if tensors: + raise ValueError( + "H3 conditioning projection contains unsupported tensors: " + f"{sorted(tensors)}" + ) + if self.weight is None and not layers: + raise ValueError("H3 conditioning projection has neither W nor an MLP") + if layers and layer_input_dim != self.output_dim: + raise ValueError( + f"H3 conditioning projection MLP outputs width {layer_input_dim}, " + f"expected {self.output_dim}" + ) + if self.weight is not None and tuple(self.weight.shape) != ( + self.input_dim, + self.output_dim, + ): + raise ValueError( + "H3 conditioning projection W has shape " + f"{tuple(self.weight.shape)}, expected " + f"({self.input_dim}, {self.output_dim})" + ) + self.layers = nn.ModuleList(layers) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if int(hidden_states.shape[-1]) != self.input_dim: + raise ValueError( + f"H3 conditioning projection expects width {self.input_dim}, " + f"got {int(hidden_states.shape[-1])}" + ) + normalized = (hidden_states.float() - self.mean_in) / self.std_in + projected = normalized @ self.weight if self.weight is not None else None + if self.layers: + residual = normalized.to(self.layers[0].weight.dtype) + for index, layer in enumerate(self.layers): + residual = layer(residual) + if index + 1 < len(self.layers): + residual = F.gelu(residual) + residual = residual.float() + projected = residual if projected is None else projected + residual + if projected is None: + raise RuntimeError("H3 conditioning projection produced no output") + output = projected * self.std_out + self.mean_out + if self.sink_out is not None and int(output.shape[-2]) > 0: + output[..., 0, :] = self.sink_out + return output class MiniMaxH3Qwen3VLEncoder(TextEncoder): - """Qwen3-VL-32B multimodal backbone ending at hidden_states[50]. + """Qwen3-VL multimodal backbone producing MiniMax H3 conditioning. The component loader builds and loads this module under the encoder-folding TP group. A TP=1/SP=8 DiT deployment therefore shards the encoder over all @@ -54,27 +199,73 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): # The inherited text-layer list covers Qwen's language stack; reference # modes also execute the embedded visual tower. - layer_names = [*TextEncoder.layer_names, "model.visual.blocks"] + layer_names = [ + *TextEncoder.layer_names, + "model.visual.blocks", + "conditioning_projection.layers", + ] supports_dp_encode = True param_names_mapping = _PARAM_NAMES_MAPPING - @staticmethod - def should_materialize_checkpoint_weight(name: str) -> bool: + @classmethod + def configure_component_paths( + cls, + config: MiniMaxH3Qwen3VLConfig, + component_paths: dict[str, str], + ) -> None: + arch = config.arch_config + source = component_paths.get("conditioning_projection") + if source is None: + if ( + int(arch.hidden_size) != MINIMAX_H3_QWEN3VL_HIDDEN_DIM + or int(arch.checkpoint_num_hidden_layers) + < MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER + ): + raise ValueError( + "MiniMax H3 Qwen3-VL encoders smaller than 32B require " + "--component-paths.conditioning_projection" + ) + return + + projection_path = materialize_weight(resolve_weight(source)) + tap, input_dim, output_dim = MiniMaxH3ConditioningProjection.inspect( + projection_path + ) + if input_dim != int(arch.hidden_size): + raise ValueError( + f"H3 conditioning projection expects encoder width {input_dim}, " + f"but the selected text encoder has width {int(arch.hidden_size)}" + ) + if output_dim != MINIMAX_H3_QWEN3VL_HIDDEN_DIM: + raise ValueError( + f"H3 conditioning projection must output width " + f"{MINIMAX_H3_QWEN3VL_HIDDEN_DIM}, got {output_dim}" + ) + if tap <= 0 or tap > int(arch.checkpoint_num_hidden_layers): + raise ValueError( + f"H3 conditioning projection tap {tap} is outside the selected " + f"encoder's {int(arch.checkpoint_num_hidden_layers)} layers" + ) + arch.conditioning_projection_path = projection_path + arch.num_hidden_layers = tap + arch.text_config.num_hidden_layers = tap + + def should_materialize_checkpoint_weight(self, name: str) -> bool: name = _map_checkpoint_name(name) return ( "rotary_emb.inv_freq" not in name - and not _is_unconsumed_checkpoint_weight(name) + and not _is_unconsumed_checkpoint_weight(name, self.selected_lm_layer) ) def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__(config) arch = config.arch_config - selected_layer = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER - if int(arch.text_config.num_hidden_layers) != selected_layer: + selected_layer = int(arch.text_config.num_hidden_layers) + if selected_layer <= 0 or int(arch.num_hidden_layers) != selected_layer: raise ValueError( - "MiniMax H3 Qwen3-VL config must be trimmed to " - f"{selected_layer} language layers before construction" + "MiniMax H3 Qwen3-VL language-layer configuration is " + f"inconsistent: {selected_layer} vs {int(arch.num_hidden_layers)}" ) self.model = Qwen3VLModel( arch, @@ -82,12 +273,17 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): use_tensor_parallel=True, prefix="model", ) - # H3 consumes the unnormalized output immediately after layer 49. + # H3 and ClipProj consume an unnormalized intermediate residual stream. self.model.language_model.norm = nn.Identity() self.image_token_id = int(arch.image_token_id) self.video_token_id = int(arch.video_token_id) self.selected_lm_layer = selected_layer self.hidden_dim = MINIMAX_H3_QWEN3VL_HIDDEN_DIM + self.conditioning_projection = ( + MiniMaxH3ConditioningProjection(arch.conditioning_projection_path) + if arch.conditioning_projection_path is not None + else None + ) @property def device(self) -> torch.device: @@ -179,7 +375,10 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): ) call_kwargs["video_grid_thw"] = host_video_grid_thw - hidden = self(**call_kwargs).last_hidden_state[0].to(torch.bfloat16) + hidden = self(**call_kwargs).last_hidden_state[0] + if self.conditioning_projection is not None: + hidden = self.conditioning_projection(hidden) + hidden = hidden.to(torch.bfloat16) expected_shape = [int(ids.shape[1]), self.hidden_dim] if list(hidden.shape) != expected_shape: raise ValueError( diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 7dd737f1b..6a73359c9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -180,6 +180,22 @@ class TestServerArgsPathExpansion(unittest.TestCase): {"text_encoder": "owner/repo/text_encoder/model.safetensors"}, ) + def test_supplemental_weight_file_remains_a_component_path(self): + args = self._from_dict_without_model_resolution( + { + "model_path": "/data/my-model", + "component_paths": { + "conditioning_projection": "owner/repo/projection.safetensors" + }, + } + ) + + self.assertEqual( + args.component_paths, + {"conditioning_projection": "owner/repo/projection.safetensors"}, + ) + self.assertEqual(args.component_weights_paths, {}) + def test_component_attention_backends_are_normalized(self): args = self._from_dict_without_model_resolution( { diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index 5bef8c2c3..60aa00c48 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -27,6 +27,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader ) from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import ( + MiniMaxH3ConditioningProjection, MiniMaxH3Qwen3VLEncoder, ) @@ -132,7 +133,9 @@ class TestTextEncoderClassResolution(unittest.TestCase): class TestMiniMaxH3CheckpointFilter(unittest.TestCase): def test_only_known_unconsumed_weights_are_filtered(self): - should_load = MiniMaxH3Qwen3VLEncoder.should_materialize_checkpoint_weight + encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder) + encoder.selected_lm_layer = 50 + should_load = encoder.should_materialize_checkpoint_weight expected = { "model.language_model.layers.49.self_attn.q_proj.weight": True, "model.language_model.layers.50.self_attn.q_proj.weight": False, @@ -151,10 +154,18 @@ class TestMiniMaxH3CheckpointFilter(unittest.TestCase): {name: should_load(name) for name in expected}, expected, ) + encoder.selected_lm_layer = 24 + self.assertTrue( + should_load("model.language_model.layers.23.mlp.down_proj.weight") + ) + self.assertFalse( + should_load("model.language_model.layers.24.mlp.down_proj.weight") + ) def test_vision_qkv_checkpoint_name_maps_to_native_projection(self): encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder) torch.nn.Module.__init__(encoder) + encoder.selected_lm_layer = 50 encoder.model = torch.nn.Module() encoder.model.visual = torch.nn.Module() block = torch.nn.Module() @@ -175,6 +186,7 @@ class TestMiniMaxH3CheckpointFilter(unittest.TestCase): def test_comfy_language_checkpoint_name_maps_to_native_namespace(self): encoder = MiniMaxH3Qwen3VLEncoder.__new__(MiniMaxH3Qwen3VLEncoder) torch.nn.Module.__init__(encoder) + encoder.selected_lm_layer = 50 encoder.model = torch.nn.Module() encoder.model.language_model = torch.nn.Module() layer = torch.nn.Module() @@ -192,6 +204,69 @@ class TestMiniMaxH3CheckpointFilter(unittest.TestCase): torch.testing.assert_close(layer.self_attn.q_proj.weight, source) +class TestMiniMaxH3ConditioningProjection(unittest.TestCase): + @staticmethod + def _save_projection(path, tensors): + save_file(tensors, path, metadata={"tap": "2"}) + + def test_linear_projection_and_attention_sink(self): + with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint: + tensors = { + "W": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "mean_in": torch.tensor([1.0, -1.0]), + "std_in": torch.tensor([2.0, 4.0]), + "mean_out": torch.tensor([0.5, -0.5]), + "std_out": torch.tensor([2.0, 3.0]), + "sink_out": torch.tensor([9.0, 8.0]), + } + self._save_projection(checkpoint.name, tensors) + projection = MiniMaxH3ConditioningProjection(checkpoint.name) + hidden = torch.tensor([[5.0, 3.0], [3.0, -1.0]]) + expected = (hidden - tensors["mean_in"]) / tensors["std_in"] + expected = expected @ tensors["W"] + expected = expected * tensors["std_out"] + tensors["mean_out"] + expected[0] = tensors["sink_out"] + + torch.testing.assert_close(projection(hidden), expected) + self.assertEqual(projection.tap, 2) + + def test_mlp_only_projection(self): + with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint: + tensors = { + "mean_in": torch.zeros(2), + "std_in": torch.ones(2), + "mean_out": torch.tensor([0.5]), + "std_out": torch.tensor([2.0]), + "mlp.0.weight": torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + "mlp.0.bias": torch.tensor([0.25, -0.25]), + "mlp.2.weight": torch.tensor([[2.0, -1.0]]), + "mlp.2.bias": torch.tensor([0.75]), + } + self._save_projection(checkpoint.name, tensors) + projection = MiniMaxH3ConditioningProjection(checkpoint.name) + hidden = torch.tensor([[1.0, 2.0]]) + residual = torch.nn.functional.linear( + hidden, tensors["mlp.0.weight"], tensors["mlp.0.bias"] + ) + residual = torch.nn.functional.gelu(residual) + residual = torch.nn.functional.linear( + residual, tensors["mlp.2.weight"], tensors["mlp.2.bias"] + ) + expected = residual * tensors["std_out"] + tensors["mean_out"] + + torch.testing.assert_close(projection(hidden), expected) + + def test_small_encoder_requires_matching_projection(self): + config = SimpleNamespace( + arch_config=SimpleNamespace( + hidden_size=2560, + checkpoint_num_hidden_layers=36, + ) + ) + with self.assertRaisesRegex(ValueError, "conditioning_projection"): + MiniMaxH3Qwen3VLEncoder.configure_component_paths(config, {}) + + class TestTextEncoderQuantization(unittest.TestCase): def setUp(self): serialized = Fp8Config(