[Boogu-Image] Fuse QKV/FFN projections + switch to diffusion RMSNorm (#6649)

Signed-off-by: dengyunyang <584797741@qq.com>

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -440,29 +440,86 @@ def test_transformer_rejects_prompt_tuning():
         BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
 
 
-def _native_to_checkpoint_name(name: str) -> str:
-    """Inverse of ``load_weights`` remapping: native param name -> diffusers name.
+_QKV_FANOUT = {
+    "to_qkv": ("to_q", "to_k", "to_v"),
+    "img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"),
+    "instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"),
+}
+
+
+def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
+    """Inverse of ``load_weights`` remapping: native param -> diffusers weights.
+
+    The diffusers checkpoint stores fused projections as separate matrices, so a
+    single merged native param (QKV or FFN gate/up) fans out into several
+    checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values,
+    when fed through ``load_weights``, reassemble into the original ``param``.
 
     - ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
     - promoted joint-attention projections move back under ``.processor.``.
     """
+    q_size = NUM_HEADS * HEAD_DIM
+    kv_size = NUM_KV_HEADS * HEAD_DIM
+
+    # Fused QKV projections split back into per-matrix q/k/v weights.
+    for token, (q_name, k_name, v_name) in _QKV_FANOUT.items():
+        marker = f".{token}."
+        if marker in name:
+            q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :]
+            results = []
+            for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)):
+                ckpt_name = name.replace(marker, f".{sub_name}.")
+                if token != "to_qkv":
+                    # Joint-attention projections live under `.processor.` upstream.
+                    ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.")
+                results.append((ckpt_name, value))
+            return results
+
+    # Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input).
+    if ".gate_up_proj." in name:
+        inner = param.shape[0] // 2
+        return [
+            (name.replace(".gate_up_proj.", ".linear_1."), param[:inner]),
+            (name.replace(".gate_up_proj.", ".linear_3."), param[inner:]),
+        ]
+
+    # Default: 1:1 with the existing diffusers name promotions.
     if ".to_out." in name:
         name = name.replace(".to_out.", ".to_out.0.")
-    for proj in (
-        "img_to_q",
-        "img_to_k",
-        "img_to_v",
-        "instruct_to_q",
-        "instruct_to_k",
-        "instruct_to_v",
-        "instruct_out",
-        "img_out",
-    ):
+    for proj in ("instruct_out", "img_out"):
         token = f".img_instruct_attn.{proj}."
         if token in name:
             name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
             break
-    return name
+    return [(name, param)]
+
+
+def test_transformer_exposes_stacked_params_mapping():
+    """The packed -> sub-layer mapping must be discoverable from the module tree.
+
+    ``diffusion/lora/loader.py`` and the quantized weight loaders read
+    ``stacked_params_mapping`` off the model, so it has to exist before (and
+    independently of) ``load_weights``.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _BOOGU_STACKED_PARAMS_MAPPING,
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+
+    assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING
+    # A per-instance copy, so consumers cannot mutate the module constant.
+    assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING
+
+    # Every entry is a (param, shard, shard_id) triple.
+    for param_name, shard_name, shard_id in model.stacked_params_mapping:
+        assert param_name.startswith(".") and shard_name.startswith(".")
+        assert shard_id in {"q", "k", "v", 0, 1}
+
+    # The mapping targets exactly the fused projections this port creates.
+    mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping}
+    assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"}
 
 
 def test_transformer_load_weights_round_trip():
@@ -473,24 +530,59 @@ def test_transformer_load_weights_round_trip():
     model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
     native_params = dict(model.named_parameters())
 
-    # Build synthetic diffusers-named weights (one per native parameter).
-    checkpoint_weights = {}
+    # Build synthetic diffusers-named weights. Fused native projections fan out
+    # into the separate matrices the diffusers checkpoint stores.
+    expected: dict[str, torch.Tensor] = {}
+    checkpoint_weights: dict[str, torch.Tensor] = {}
     for native_name, param in native_params.items():
-        checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
+        full = torch.randn_like(param)
+        expected[native_name] = full
+        for ckpt_name, value in _native_to_checkpoint_weights(native_name, full):
+            checkpoint_weights[ckpt_name] = value
 
-    # The remapping must be a bijection over the parameter set.
-    assert len(checkpoint_weights) == len(native_params)
+    # Every checkpoint name is unique; merged params fan out to >=1 weight.
+    assert len(checkpoint_weights) >= len(native_params)
 
     loaded = model.load_weights(list(checkpoint_weights.items()))
 
     # No missing / unexpected parameters.
     assert loaded == set(native_params.keys())
 
-    # Values landed on the right parameters (TP=1: weight_loader copies verbatim).
+    # Values landed on the right parameters (TP=1: weight_loader copies verbatim;
+    # fused shards are placed at their q/k/v / gate/up offsets).
     reloaded = dict(model.named_parameters())
     for native_name in native_params:
-        expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
-        assert torch.allclose(reloaded[native_name], expected)
+        assert torch.allclose(reloaded[native_name], expected[native_name])
+
+
+def test_transformer_load_weights_rejects_partial_fused_params():
+    """A fused parameter must not count as loaded until all its shards arrive.
+
+    The loader compares parameter *names*, so reporting ``to_qkv`` complete
+    after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start
+    up with the ``k``/``v`` slices left uninitialized.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    native_params = dict(model.named_parameters())
+
+    complete_name = "noise_refiner.0.attn.to_qkv.weight"
+    partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight"
+    assert complete_name in native_params
+    assert partial_name in native_params
+
+    weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name]))
+    # Only the gate half of the FFN, so its fused parameter stays incomplete.
+    gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0]
+    weights.append((gate_ckpt, gate_value))
+
+    loaded = model.load_weights(weights)
+
+    assert complete_name in loaded
+    assert partial_name not in loaded
 
 
 def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
@@ -513,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
     try:
         model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
         native_params = dict(model.named_parameters())
-        loaded_name = next(iter(native_params))
-        checkpoint_name = _native_to_checkpoint_name(loaded_name)
+        # A non-fused parameter, so one checkpoint entry loads it completely.
+        loaded_name = "x_embedder.weight"
+        assert loaded_name in native_params
+        (checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name])
+        checkpoint_value = torch.randn_like(checkpoint_value)
 
         loaded = model.load_weights(
             [
-                (checkpoint_name, torch.randn_like(native_params[loaded_name])),
+                (checkpoint_name, checkpoint_value),
                 ("unexpected.weight", torch.ones(1)),
             ]
         )

vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py

@@ -15,6 +15,7 @@
 #     inference caches (TeaCache/TaylorSeer) are not ported.
 
 import itertools
+from collections import defaultdict
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
@@ -24,9 +25,9 @@ import torch.nn.functional as F
 from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed
 from einops import rearrange, repeat
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
     ReplicatedLinear,
     RowParallelLinear,
 )
@@ -42,6 +43,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
     fused_qk_norm_rope,
     fused_qk_norm_rope_min_tokens,
 )
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.models.utils import make_attention_mask
 from vllm_omni.platforms import current_omni_platform
 
@@ -260,7 +262,14 @@ class LuminaLayerNormContinuous(nn.Module):
 
 
 class LuminaFeedForward(nn.Module):
-    """SwiGLU feed-forward with tensor-parallel projections."""
+    """SwiGLU feed-forward with tensor-parallel projections.
+
+    The ``gate`` and ``input`` projections are fused into a single
+    ``MergedColumnParallelLinear`` (one GEMM instead of two), matching the
+    ``gate_up_proj`` convention used by the Hunyuan Image 3 port. The SwiGLU
+    activation keeps its float32 computation for numerical parity with
+    upstream.
+    """
 
     def __init__(
         self,
@@ -276,20 +285,13 @@ class LuminaFeedForward(nn.Module):
             inner_dim = int(ffn_dim_multiplier * inner_dim)
         inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
 
-        self.linear_1 = ColumnParallelLinear(
-            dim,
-            inner_dim,
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=dim,
+            output_sizes=[inner_dim, inner_dim],
             bias=False,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "linear_1"),
-        )  # gate
-        self.linear_3 = ColumnParallelLinear(
-            dim,
-            inner_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "linear_3"),
-        )  # input
+            prefix=_join_prefix(prefix, "gate_up_proj"),
+        )
         self.linear_2 = RowParallelLinear(
             inner_dim,
             dim,
@@ -299,8 +301,8 @@ class LuminaFeedForward(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        h1, _ = self.linear_1(x)
-        h2, _ = self.linear_3(x)
+        gate_up, _ = self.gate_up_proj(x)
+        h1, h2 = gate_up.chunk(2, dim=-1)
         out, _ = self.linear_2(swiglu(h1, h2))
         return out
 
@@ -593,28 +595,15 @@ class BooguImageSelfAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.head_dim = dim // num_attention_heads
-        kv_dim = self.head_dim * num_kv_heads
 
-        self.to_q = ColumnParallelLinear(
-            dim,
-            dim,
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_attention_heads,
+            total_num_kv_heads=num_kv_heads,
             bias=False,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_q"),
-        )
-        self.to_k = ColumnParallelLinear(
-            dim,
-            kv_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_k"),
-        )
-        self.to_v = ColumnParallelLinear(
-            dim,
-            kv_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_v"),
+            prefix=_join_prefix(prefix, "to_qkv"),
         )
         self.norm_q = RMSNorm(self.head_dim, eps=1e-5)
         self.norm_k = RMSNorm(self.head_dim, eps=1e-5)
@@ -626,8 +615,10 @@ class BooguImageSelfAttention(nn.Module):
             prefix=_join_prefix(prefix, "to_out"),
         )
 
-        self.num_local_heads = self.to_q.output_size_per_partition // self.head_dim
-        self.num_local_kv_heads = self.to_k.output_size_per_partition // self.head_dim
+        self.num_local_heads = self.to_qkv.num_heads
+        self.num_local_kv_heads = self.to_qkv.num_kv_heads
+        self.q_size = self.num_local_heads * self.head_dim
+        self.kv_size = self.num_local_kv_heads * self.head_dim
 
         self.attn = Attention(
             num_heads=self.num_local_heads,
@@ -645,9 +636,8 @@ class BooguImageSelfAttention(nn.Module):
     ) -> torch.Tensor:
         dtype = hidden_states.dtype
 
-        query, _ = self.to_q(hidden_states)
-        key, _ = self.to_k(hidden_states)
-        value, _ = self.to_v(hidden_states)
+        qkv, _ = self.to_qkv(hidden_states)
+        query, key, value = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
         query = query.unflatten(-1, (self.num_local_heads, self.head_dim))
         key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim))
@@ -683,50 +673,24 @@ class BooguImageJointAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.head_dim = dim // num_attention_heads
-        kv_dim = self.head_dim * num_kv_heads
 
-        self.img_to_q = ColumnParallelLinear(
-            dim,
-            dim,
+        self.img_to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_attention_heads,
+            total_num_kv_heads=num_kv_heads,
             bias=False,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "img_to_q"),
+            prefix=_join_prefix(prefix, "img_to_qkv"),
         )
-        self.img_to_k = ColumnParallelLinear(
-            dim,
-            kv_dim,
+        self.instruct_to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_attention_heads,
+            total_num_kv_heads=num_kv_heads,
             bias=False,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "img_to_k"),
-        )
-        self.img_to_v = ColumnParallelLinear(
-            dim,
-            kv_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "img_to_v"),
-        )
-
-        self.instruct_to_q = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "instruct_to_q"),
-        )
-        self.instruct_to_k = ColumnParallelLinear(
-            dim,
-            kv_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "instruct_to_k"),
-        )
-        self.instruct_to_v = ColumnParallelLinear(
-            dim,
-            kv_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=_join_prefix(prefix, "instruct_to_v"),
+            prefix=_join_prefix(prefix, "instruct_to_qkv"),
         )
 
         self.norm_q = RMSNorm(self.head_dim, eps=1e-5)
@@ -757,8 +721,10 @@ class BooguImageJointAttention(nn.Module):
             prefix=_join_prefix(prefix, "to_out"),
         )
 
-        self.num_local_heads = self.img_to_q.output_size_per_partition // self.head_dim
-        self.num_local_kv_heads = self.img_to_k.output_size_per_partition // self.head_dim
+        self.num_local_heads = self.img_to_qkv.num_heads
+        self.num_local_kv_heads = self.img_to_qkv.num_kv_heads
+        self.q_size = self.num_local_heads * self.head_dim
+        self.kv_size = self.num_local_kv_heads * self.head_dim
 
         self.attn = Attention(
             num_heads=self.num_local_heads,
@@ -780,13 +746,13 @@ class BooguImageJointAttention(nn.Module):
         dtype = img_hidden_states.dtype
         batch_size = img_hidden_states.shape[0]
 
-        img_query, _ = self.img_to_q(img_hidden_states)
-        img_key, _ = self.img_to_k(img_hidden_states)
-        img_value, _ = self.img_to_v(img_hidden_states)
+        img_qkv, _ = self.img_to_qkv(img_hidden_states)
+        img_query, img_key, img_value = img_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
-        instruct_query, _ = self.instruct_to_q(instruct_hidden_states)
-        instruct_key, _ = self.instruct_to_k(instruct_hidden_states)
-        instruct_value, _ = self.instruct_to_v(instruct_hidden_states)
+        instruct_qkv, _ = self.instruct_to_qkv(instruct_hidden_states)
+        instruct_query, instruct_key, instruct_value = instruct_qkv.split(
+            [self.q_size, self.kv_size, self.kv_size], dim=-1
+        )
 
         query, key, value = _concat_instruction_image_features(
             [img_query, img_key, img_value],
@@ -1135,6 +1101,44 @@ def _cal_preprocessed_instruction_feat_dim(instruction_feature_configs: dict) ->
         raise ValueError(f"Invalid reduce_type: {reduce_type}")
 
 
+# Packed (fused) projections vs. the logical sub-projections the diffusers
+# checkpoint stores, in the ``(param_name, shard_name, shard_id)`` form used
+# across vllm-omni. ``load_weights`` consumes it directly, and loader consumers
+# that discover it via ``stacked_params_mapping`` (LoRA, quantized loaders) get
+# the packed -> sub-layer relationship for free.
+#
+# The QKV entries stay leaf-scoped: each fused source (``to_q`` on the
+# self-attention module, ``img_to_q`` / ``instruct_to_q`` on the joint
+# attention) has a distinct leaf, and ``.to_q.`` cannot collide with
+# ``.img_to_q.`` because the preceding character is ``_`` rather than ``.``.
+# The FFN entries are path-qualified instead: ``linear_1`` also names the
+# timestep embedder and the ``norm_out`` projections, which must load 1:1 and
+# are not fused.
+_BOOGU_STACKED_PARAMS_MAPPING = (
+    # self-attention: noise / reference-image / context refiners, the
+    # single-stream blocks, and the double-stream image self-attention.
+    (".to_qkv.", ".to_q.", "q"),
+    (".to_qkv.", ".to_k.", "k"),
+    (".to_qkv.", ".to_v.", "v"),
+    # joint (instruction + image) attention of the double-stream blocks.
+    (".img_to_qkv.", ".img_to_q.", "q"),
+    (".img_to_qkv.", ".img_to_k.", "k"),
+    (".img_to_qkv.", ".img_to_v.", "v"),
+    (".instruct_to_qkv.", ".instruct_to_q.", "q"),
+    (".instruct_to_qkv.", ".instruct_to_k.", "k"),
+    (".instruct_to_qkv.", ".instruct_to_v.", "v"),
+    # feed-forward gate/up: ``linear_1`` is the gate (shard 0) and
+    # ``linear_3`` the input (shard 1), matching the packed order the SwiGLU
+    # activation consumes.
+    (".feed_forward.gate_up_proj.", ".feed_forward.linear_1.", 0),
+    (".feed_forward.gate_up_proj.", ".feed_forward.linear_3.", 1),
+    (".img_feed_forward.gate_up_proj.", ".img_feed_forward.linear_1.", 0),
+    (".img_feed_forward.gate_up_proj.", ".img_feed_forward.linear_3.", 1),
+    (".instruct_feed_forward.gate_up_proj.", ".instruct_feed_forward.linear_1.", 0),
+    (".instruct_feed_forward.gate_up_proj.", ".instruct_feed_forward.linear_3.", 1),
+)
+
+
 class BooguImageTransformer2DModel(nn.Module):
     """Boogu-Image transformer with mixed stream topology.
 
@@ -1160,6 +1164,10 @@ class BooguImageTransformer2DModel(nn.Module):
         prefix: str = "",
     ) -> None:
         super().__init__()
+        # Published here rather than in ``load_weights`` so loader consumers
+        # that only inspect the module tree (LoRA discovery, quantized weight
+        # loaders) always see the packed -> sub-layer relationship.
+        self.stacked_params_mapping = list(_BOOGU_STACKED_PARAMS_MAPPING)
         self.od_config = od_config
         cfg = od_config.tf_model_config
 
@@ -1710,7 +1718,7 @@ class BooguImageTransformer2DModel(nn.Module):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load diffusers-named checkpoint weights into the native module.
 
-        Two name promotions relative to upstream (see step 8/10 findings):
+        Name promotions relative to upstream (see step 8/10 findings):
 
         - ``*.img_instruct_attn.processor.{img,instruct}_{to_q,to_k,to_v}`` /
           ``{instruct,img}_out`` -> drop ``.processor`` (upstream keeps the
@@ -1719,9 +1727,29 @@ class BooguImageTransformer2DModel(nn.Module):
         - ``*.to_out.0.weight`` -> ``*.to_out.weight`` (diffusers wraps the
           output projection in a ``ModuleList``; the native module uses a plain
           linear).
+
+        Packed Q/K/V and FFN gate/input matrices are folded onto the fused
+        ``QKVParallelLinear`` / ``MergedColumnParallelLinear`` parameters using
+        :attr:`stacked_params_mapping` — the same mapping loader consumers such
+        as LoRA discovery and quantized weight loaders read.
+
+        A fused parameter is only reported as loaded once *every* one of its
+        source matrices has arrived. The caller compares parameter names, so
+        reporting e.g. ``to_qkv`` complete after the first shard would let a
+        checkpoint carrying only ``to_q`` start up with the ``k``/``v`` slices
+        left uninitialized.
         """
+        stacked_params_mapping = self.stacked_params_mapping
+        # The shard ids each fused parameter is assembled from, keyed by the
+        # mapping entry that produces it.
+        expected_shards: dict[str, set[str | int]] = defaultdict(set)
+        for param_name, _weight_name, packed_shard_id in stacked_params_mapping:
+            expected_shards[param_name].add(packed_shard_id)
+
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        # Fused parameter -> (mapping entry that produced it, shards received).
+        fused_params: dict[str, tuple[str, set[str | int]]] = {}
 
         for name, loaded_weight in weights:
             original_name = name
@@ -1730,15 +1758,36 @@ class BooguImageTransformer2DModel(nn.Module):
             if ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
 
+            shard_id: str | int | None = None
+            matched_entry: str | None = None
+            for param_name, weight_name, packed_shard_id in stacked_params_mapping:
+                if weight_name in name:
+                    name = name.replace(weight_name, param_name)
+                    shard_id = packed_shard_id
+                    matched_entry = param_name
+                    break
+
             if name not in params_dict:
                 logger.warning("Skipping unexpected checkpoint weight %s", original_name)
                 continue
 
             param = params_dict[name]
             weight_loader = getattr(param, "weight_loader", default_weight_loader)
-            weight_loader(param, loaded_weight)
+            if matched_entry is None:
+                weight_loader(param, loaded_weight)
+            else:
+                weight_loader(param, loaded_weight, shard_id)
+                _, received = fused_params.setdefault(name, (matched_entry, set()))
+                received.add(shard_id)
             loaded_params.add(name)
 
+        # Withdraw any fused parameter that only received part of its source
+        # matrices, so the caller's "not initialized from checkpoint" check
+        # reports it as missing instead of accepting a half-filled parameter.
+        for name, (matched_entry, received) in fused_params.items():
+            if received != expected_shards[matched_entry]:
+                loaded_params.discard(name)
+
         unloaded_params = sorted(params_dict.keys() - loaded_params)
         if unloaded_params:
             logger.warning("Model parameters not loaded from checkpoint: %s", unloaded_params)