From 026f6ab2689397e64c1f53a5ac57d37c4dc46b67 Mon Sep 17 00:00:00 2001 From: gitbisector Date: Sat, 17 Oct 2026 02:37:56 -0700 Subject: [PATCH] [Model] DeepSeek-V4/V4.1 VL: stream language-model weights instead of sorting the checkpoint (#58857) Signed-off-by: git bisector Signed-off-by: Yifan Qiao Co-authored-by: Claude Co-authored-by: Yifan Qiao --- .buildkite/test_areas/models_basic.yaml | 3 +- .../test_deepseek_v4_vl_stream_weights.py | 78 +++++++++++++++++++ vllm/models/deepseek_v4/common/vl_model.py | 20 ++++- vllm/models/deepseek_v41/nvidia/vl_model.py | 10 +-- 4 files changed, 103 insertions(+), 8 deletions(-) create mode 100644 tests/models/test_deepseek_v4_vl_stream_weights.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 029ccbbffa..e76a99d059 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.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)" diff --git a/tests/models/test_deepseek_v4_vl_stream_weights.py b/tests/models/test_deepseek_v4_vl_stream_weights.py new file mode 100644 index 0000000000..66fc10be43 --- /dev/null +++ b/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 diff --git a/vllm/models/deepseek_v4/common/vl_model.py b/vllm/models/deepseek_v4/common/vl_model.py index 708a57d737..7868c59f33 100644 --- a/vllm/models/deepseek_v4/common/vl_model.py +++ b/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 diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index d67753405a..348c31d722 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/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.