[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>

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.