[Model] Fail MiniMax H3 encoder load when a weight or fused shard is missing (#5824)

Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

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:

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())