From 88754c0850d66814471ae6bf0924ccdac388edd9 Mon Sep 17 00:00:00 2001 From: git bisector Date: Sat, 26 Sep 2026 13:21:31 -0700 Subject: [PATCH] [Model] DeepSeek-V4/V4.1 VL: stream language-model weights instead of sorting the checkpoint The DeepSeek-V4 VL wrapper and the NVIDIA DeepSeek-V4.1 VL wrapper sort the whole mapped weight stream so the "language_model." group reaches the child loader contiguously; the NVIDIA children finalize fused expert weights inside load_weights. sorted() materializes the iterator, so every tensor the loader yields is alive before the first copy: the whole checkpoint for loaders that yield materialized tensors (--safetensors-load-strategy eager, --load-format fastsafetensors), and a shard-hopping access order for mmap views. Contiguity only needs a non-interleaved group. Add stream_language_model_first, which yields language_model.* as it arrives and defers the rest (vision tower, aligner), cloned in case the loader reuses its buffers, and use it in both wrappers. #58782 makes the same change to the AMD V4.1 wrapper. Co-authored-by: Claude Signed-off-by: git bisector --- .buildkite/test_areas/models_basic.yaml | 3 +- .../test_deepseek_v4_vl_stream_weights.py | 77 +++++++++++++++++++ vllm/models/deepseek_v4/common/vl_model.py | 20 ++++- vllm/models/deepseek_v41/nvidia/vl_model.py | 10 +-- 4 files changed, 102 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 54130d0d020b..16b591625252 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -68,8 +68,9 @@ steps: - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.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/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.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/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_vl_stream_weights.py mirror: amd: label: ":amd: (MI355 DPX) 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 000000000000..32899115b42e --- /dev/null +++ b/tests/models/test_deepseek_v4_vl_stream_weights.py @@ -0,0 +1,77 @@ +# 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 "vision.tensor_b", torch.tensor([2.0]) + 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 + 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 97701cc2d507..1a8002ede2d3 100644 --- a/vllm/models/deepseek_v4/common/vl_model.py +++ b/vllm/models/deepseek_v4/common/vl_model.py @@ -78,6 +78,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 in case the loader + reuses their buffers. + """ + 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, @@ -329,7 +347,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 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 d67753405ad5..8d9c7a3c3609 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/vl_model.py @@ -41,6 +41,7 @@ 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 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: 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 fused expert weights 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.