[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