[Model] DeepSeek-V4/V4.1 VL: stream language-model weights instead of sorting the checkpoint (#58857)

Signed-off-by: git bisector <gitbisector@gmail.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

.buildkite/test_areas/models_basic.yaml

@@ -72,8 +72,9 @@ steps:
   - tests/models/test_deepseek_v41_pcp.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
+  - tests/models/test_deepseek_v4_vl_stream_weights.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v4_vl_stream_weights.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"

tests/models/test_deepseek_v4_vl_stream_weights.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The DeepSeek-V4/V4.1 VL wrappers stream the language model's weights to it
+as one contiguous group."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="CUDA implementations"
+)
+
+
+class _FakeFinalizingLanguageModel(nn.Module):
+    finalizes_weights_during_load = True
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.tensor_a = nn.Parameter(torch.zeros(1))
+        self.tensor_c = nn.Parameter(torch.zeros(1))
+        self.load_calls = 0
+        self.finalized_values: list[tuple[float, float]] = []
+
+    def load_weights(self, weights) -> set[str]:
+        self.load_calls += 1
+        loaded = set()
+        for name, value in weights:
+            getattr(self, name).data.copy_(value)
+            loaded.add(name)
+        self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item()))
+        return loaded
+
+
+def _v4_wrapper():
+    from vllm.models.deepseek_v4.common.vl_model import (
+        DeepseekV4ForConditionalGeneration,
+    )
+
+    return DeepseekV4ForConditionalGeneration
+
+
+def _v41_wrapper():
+    from vllm.models.deepseek_v41.nvidia.vl_model import DeepseekV41ForCausalLM
+
+    return DeepseekV41ForCausalLM
+
+
+@pytest.mark.parametrize("wrapper", [_v4_wrapper, _v41_wrapper], ids=["v4", "v41"])
+def test_load_weights_streams_language_model_group(wrapper) -> None:
+    model = object.__new__(wrapper())
+    nn.Module.__init__(model)
+    model.language_model = _FakeFinalizingLanguageModel()
+    model.vision = nn.Module()
+    model.vision.tensor_b = nn.Parameter(torch.zeros(1))
+    model.hf_to_vllm_mapper = WeightsMapper()
+
+    def interleaved_weights():
+        yield "language_model.tensor_a", torch.tensor([1.0])
+        # Loaded before the stream is drained, not after.
+        assert model.language_model.tensor_a.item() == 1.0
+        # Splits the language model's weights unless the wrapper regroups them.
+        yield "vision.tensor_b", torch.tensor([2.0])
+        yield "language_model.tensor_c", torch.tensor([3.0])
+
+    loaded = model.load_weights(interleaved_weights())
+
+    assert loaded == {
+        "language_model.tensor_a",
+        "vision.tensor_b",
+        "language_model.tensor_c",
+    }
+    assert model.language_model.load_calls == 1
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+    assert model.vision.tensor_b.item() == 2.0

vllm/models/deepseek_v4/common/vl_model.py

@@ -79,6 +79,24 @@ def _make_deepseek_v4_vl_weights_mapper(
     )
 
 
+def stream_language_model_first(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterable[tuple[str, torch.Tensor]]:
+    """Yield ``language_model.*`` weights as they arrive, then the rest.
+
+    Keeps the language model's weights in one contiguous group without
+    materializing the stream. Deferred weights are cloned so they do not keep
+    the loader's per-shard storage alive until the end of the load.
+    """
+    rest = []
+    for name, tensor in weights:
+        if name.startswith("language_model."):
+            yield name, tensor
+        else:
+            rest.append((name, tensor.clone()))
+    yield from rest
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     DeepseekV4VLMultiModalProcessor,
     info=DeepseekV4VLProcessingInfo,
@@ -334,7 +352,7 @@ class DeepseekV4ForConditionalGeneration(
         if child_finalizes:
             # A child which finalizes inside load_weights must see all of its
             # weights in one contiguous delegation from AutoWeightsLoader.
-            mapped = iter(sorted(mapped, key=lambda x: x[0]))
+            mapped = stream_language_model_first(mapped)
         loader = AutoWeightsLoader(self)
         loaded_params = loader.load_weights(mapped)
         self._weights_finalized = child_finalizes

vllm/models/deepseek_v41/nvidia/vl_model.py

@@ -41,6 +41,7 @@ from vllm.models.deepseek_v4.common.vision import (
     DeepseekV4ViT,
     run_dp_sharded_vision_tower,
 )
+from vllm.models.deepseek_v4.common.vl_model import stream_language_model_first
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
@@ -317,12 +318,9 @@ class DeepseekV41ForCausalLM(
         return self.language_model.get_mtp_target_hidden_states()
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        # Map HF names into this wrapper's namespace up front and sort, so
-        # the "language_model." group reaches the child loader as one
-        # contiguous block (AutoWeightsLoader delegates per contiguous group,
-        # and the child's load_weights finalizes fused expert weights, which
-        # must not run on a partially loaded model).
-        mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0])
+        # The child finalizes in load_weights, so its weights must reach it as
+        # one contiguous group.
+        mapped = stream_language_model_first(self.hf_to_vllm_mapper.apply(weights))
         loader = AutoWeightsLoader(self)
         loaded_params = loader.load_weights(mapped)
         # The child's load_weights already ran its post-load finalization.