diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index d2cbc4551..8cb4a5311 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -9,6 +9,7 @@ from unittest.mock import Mock import numpy as np import pytest import torch +import torch.nn as nn pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -1465,3 +1466,149 @@ def test_g4_reference_video_metadata_validation(field, value, message, tmp_path) metadata[field] = value with pytest.raises(ValueError, match=message): _validate_reference_video_metadata(metadata, index=0, source=str(tmp_path / "reference.mp4")) + + +_ENCODER_HIDDEN = 4 +_ENCODER_HEAD_DIM = 2 +_ENCODER_NUM_HEADS = 2 +_ENCODER_NUM_KV_HEADS = 1 +_ENCODER_INTERMEDIATE = 4 + +# One distinct value per checkpoint tensor, none of them a parameter initializer +# (fused weights start as `torch.empty`, norms as ones). A uniform fill would +# prove only that every row was written, not that each source landed in its own +# row range, which is the invariant the fused loaders have to get right. +_SOURCE_FILL = { + "self_attn.q_proj.weight": 3.0, + "self_attn.k_proj.weight": 4.0, + "self_attn.v_proj.weight": 5.0, + "mlp.gate_proj.weight": 6.0, + "mlp.up_proj.weight": 7.0, + "input_layernorm.weight": 8.0, +} + +_QKV_TARGET = "text_model.layers.0.self_attn.qkv_proj.weight" +_GATE_UP_TARGET = "text_model.layers.0.mlp.gate_up_proj.weight" +_NORM_TARGET = "text_model.layers.0.input_layernorm.weight" + + +def _one_layer_text_encoder(): + """Encoder stub holding both fused weights plus one plain parameter.""" + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + MiniMaxH3Qwen3VLMergedColumnParallelLinear, + MiniMaxH3Qwen3VLQKVParallelLinear, + MiniMaxH3Qwen3VLRMSNorm, + ) + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _SingleRankEncoderGroup + + group = _SingleRankEncoderGroup(0) + layer = nn.Module() + layer.self_attn = nn.Module() + layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear( + group, + hidden_size=_ENCODER_HIDDEN, + num_heads=_ENCODER_NUM_HEADS, + num_kv_heads=_ENCODER_NUM_KV_HEADS, + head_dim=_ENCODER_HEAD_DIM, + dtype=torch.float32, + ) + layer.mlp = nn.Module() + layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear( + group, + input_size=_ENCODER_HIDDEN, + intermediate_size=_ENCODER_INTERMEDIATE, + dtype=torch.float32, + ) + layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN) + + encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) + nn.Module.__init__(encoder) + encoder.text_model = nn.Module() + encoder.text_model.layers = nn.ModuleList([layer]) + return encoder + + +def _write_text_encoder_checkpoint(path, *, omit=()): + """Write a one-layer Qwen3-VL-named checkpoint, one fill value per source.""" + import json + + import safetensors.torch + + prefix = "model.language_model.layers.0" + q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM + kv_rows = _ENCODER_NUM_KV_HEADS * _ENCODER_HEAD_DIM + shapes = { + "self_attn.q_proj.weight": (q_rows, _ENCODER_HIDDEN), + "self_attn.k_proj.weight": (kv_rows, _ENCODER_HIDDEN), + "self_attn.v_proj.weight": (kv_rows, _ENCODER_HIDDEN), + "mlp.gate_proj.weight": (_ENCODER_INTERMEDIATE, _ENCODER_HIDDEN), + "mlp.up_proj.weight": (_ENCODER_INTERMEDIATE, _ENCODER_HIDDEN), + "input_layernorm.weight": (_ENCODER_HIDDEN,), + } + tensors = {f"{prefix}.{source}": torch.full(shape, _SOURCE_FILL[source]) for source, shape in shapes.items()} + for source in omit: + del tensors[f"{prefix}.{source}"] + safetensors.torch.save_file(tensors, str(path / "model.safetensors")) + index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")} + (path / "model.safetensors.index.json").write_text(json.dumps(index)) + + +@pytest.mark.parametrize( + ("omitted_sources", "expected_detail", "unreported_target"), + [ + (("self_attn.q_proj.weight",), f"{_QKV_TARGET}: ['q']", _GATE_UP_TARGET), + (("self_attn.k_proj.weight",), f"{_QKV_TARGET}: ['k']", _GATE_UP_TARGET), + (("self_attn.v_proj.weight",), f"{_QKV_TARGET}: ['v']", _GATE_UP_TARGET), + (("mlp.gate_proj.weight",), f"{_GATE_UP_TARGET}: ['gate']", _QKV_TARGET), + (("mlp.up_proj.weight",), f"{_GATE_UP_TARGET}: ['up']", _QKV_TARGET), + ( + ("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"), + f"{_QKV_TARGET}: ['q', 'k', 'v']", + _GATE_UP_TARGET, + ), + ], +) +def test_encoder_load_reports_missing_fused_source_shards( + tmp_path, omitted_sources, expected_detail, unreported_target +): + encoder = _one_layer_text_encoder() + _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources) + + with pytest.raises(RuntimeError) as excinfo: + encoder._load_weights(str(tmp_path)) + + message = str(excinfo.value) + assert expected_detail in message + assert unreported_target not in message + + +def test_encoder_load_rejects_unloaded_plain_param(tmp_path): + encoder = _one_layer_text_encoder() + _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",)) + + with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"): + encoder._load_weights(str(tmp_path)) + + +def test_encoder_load_places_every_source_in_its_own_rows(tmp_path): + encoder = _one_layer_text_encoder() + _write_text_encoder_checkpoint(tmp_path) + + encoder._load_weights(str(tmp_path)) + + params = dict(encoder.named_parameters()) + q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM + kv_rows = _ENCODER_NUM_KV_HEADS * _ENCODER_HEAD_DIM + qkv = params[_QKV_TARGET] + gate_up = params[_GATE_UP_TARGET] + slices = { + "self_attn.q_proj.weight": qkv[:q_rows], + "self_attn.k_proj.weight": qkv[q_rows : q_rows + kv_rows], + "self_attn.v_proj.weight": qkv[q_rows + kv_rows :], + "mlp.gate_proj.weight": gate_up[:_ENCODER_INTERMEDIATE], + "mlp.up_proj.weight": gate_up[_ENCODER_INTERMEDIATE:], + "input_layernorm.weight": params[_NORM_TARGET], + } + for source, rows in slices.items(): + assert torch.equal(rows, torch.full_like(rows, _SOURCE_FILL[source])), source diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index f811eeecb..15098b098 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -206,6 +206,38 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): raise ValueError(f"unexpected QKV shard id {loaded_shard_id!r}") +# A fused weight is filled from several checkpoint tensors, so a partial load leaves +# uninitialized `torch.empty` rows that no later check would catch. Keyed by the class +# owning the multi-shard `weight_loader`; each entry is the source name reported to the +# user paired with the shard id that loader expects, in projection order. +_FUSED_SOURCE_SHARDS: dict[type[nn.Module], tuple[tuple[str, str | int], ...]] = { + MiniMaxH3Qwen3VLQKVParallelLinear: (("q", "q"), ("k", "k"), ("v", "v")), + MiniMaxH3Qwen3VLMergedColumnParallelLinear: (("gate", 0), ("up", 1)), +} + + +def _fused_source_shards(module: nn.Module) -> tuple[tuple[str, str | int], ...] | None: + for cls, sources in _FUSED_SOURCE_SHARDS.items(): + if isinstance(module, cls): + return sources + return None + + +def _verify_fused_shards_complete( + expected: dict[str, tuple[tuple[str, str | int], ...]], + loaded: dict[str, set[str | int]], +) -> None: + missing = { + name: [source for source, shard_id in sources if shard_id not in loaded[name]] + for name, sources in expected.items() + } + missing = {name: sources for name, sources in missing.items() if sources} + if not missing: + return + details = "; ".join(f"{name}: {sources}" for name, sources in sorted(missing.items())) + raise RuntimeError(f"MiniMax H3 text encoder fused weights are missing source shards: {details}") + + class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): def __init__( self, @@ -952,7 +984,7 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): def tp_size(self) -> int: return self._tp_size - def _map_weight_name(self, name: str) -> tuple[str, str | None] | None: + def _map_weight_name(self, name: str) -> tuple[str, str | int | None] | None: if name == "lm_head.weight" or name == "model.language_model.norm.weight": return None if name.startswith("model.visual."): @@ -1004,6 +1036,15 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): files = sorted(set(weight_map.values())) params = dict(self.named_parameters()) loaded: set[str] = set() + # Seeded from the modules that own a multi-shard loader, so a fused weight + # that receives no source shard at all is caught too, not only a partially + # filled one. + expected_fused_shards: dict[str, tuple[tuple[str, str | int], ...]] = { + f"{module_name}.weight": sources + for module_name, module in self.named_modules() + if (sources := _fused_source_shards(module)) is not None + } + loaded_fused_shards: dict[str, set[str | int]] = {name: set() for name in expected_fused_shards} for shard in files: tensors = safetensors.torch.load_file(os.path.join(model_path, shard), device="cpu") for name, tensor in tensors.items(): @@ -1018,13 +1059,13 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): weight_loader = getattr(param, "weight_loader", _default_weight_loader) weight_loader(param, tensor, shard_id) loaded.add(param_name) + if shard_id is not None and param_name in loaded_fused_shards: + loaded_fused_shards[param_name].add(shard_id) + _verify_fused_shards_complete(expected_fused_shards, loaded_fused_shards) missing = sorted(set(params) - loaded) if missing: - logger.warning( - "MiniMax H3 text encoder weights not loaded (retained-layer subset): %d params, e.g. %s", - len(missing), - missing[:5], - ) + # Listed in full: this aborts startup, so the message is the only diagnostic. + raise RuntimeError(f"MiniMax H3 text encoder weights not loaded: {len(missing)} params: {missing}") def load_to_device(self) -> None: if not self.is_loaded: diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 610e1925f..cac987d60 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -724,7 +724,10 @@ class MiniMaxH3Pipeline( loaded_with_prefix.update(prefix + name for name in loaded) # The text encoder and both VAEs load eagerly in ``__init__`` rather # than through ``weights_sources``. Record them for the runner's strict - # missing-parameter check. + # missing-parameter check. For the text encoder that report is exact + # because its loader raises on any unloaded parameter or partially filled + # fused parameter; weaken that and gaps here become invisible. The two + # VAEs carry no equivalent guarantee. for component_name in ("text_encoder", "video_vae", "audio_vae"): component = getattr(self, component_name) loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())