[Model Loader] Run the model-level post-load hook once at the end of every load and reload

.buildkite/test-amd.yaml

@@ -4840,15 +4840,15 @@ steps:
   - tests/models/test_hyv4_rocm.py
   - tests/models/test_vision.py
   - tests/models/test_deepseek_v4_vl_rocm.py
-  - tests/models/test_deepseek_v4_vl_stream_weights.py
+  - tests/models/test_deepseek_v4_vl.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_v4_vl_stream_weights.py
+    models/test_deepseek_v4_vl.py
   - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
   - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge
 

.buildkite/test_areas/models_basic.yaml

@@ -70,13 +70,13 @@ steps:
   - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   - tests/models/test_deepseek_v4_vl_rocm.py
+  - tests/models/test_deepseek_v4_vl.py
   - tests/models/test_deepseek_v41_decoder_replay_layers.py
   - tests/models/test_deepseek_v41_replay_batch.py
   - tests/models/test_deepseek_v41_pcp.py
   - tests/models/test_deepseek_v41_pipeline.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
   - tests/models/test_muse_spark_rocm_loading.py
   - tests/models/test_muse_spark_logits_local.py
   - tests/models/test_muse_spark_rocm_shared_rope.py
@@ -87,13 +87,13 @@ steps:
       models/transformers/test_backend.py
       models/test_registry.py
       models/test_deepseek_v4_vl_rocm.py
+      models/test_deepseek_v4_vl.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_v41_pipeline.py
-      models/test_deepseek_v4_vl_stream_weights.py
       models/test_muse_spark_rocm_loading.py
       models/test_muse_spark_logits_local.py
       models/test_muse_spark_rocm_shared_rope.py
@@ -110,6 +110,7 @@ steps:
       - vllm/models/deepseek_v4/
       - vllm/models/deepseek_v41/
       - tests/models/test_deepseek_v4_vl_rocm.py
+      - tests/models/test_deepseek_v4_vl.py
       - tests/models/test_hyv4_rocm.py
       - tests/models/transformers/test_layer_registry_e2e.py
       - tests/models/test_deepseek_v41_replay_start.py
@@ -122,7 +123,6 @@ steps:
       - tests/models/test_dsv41_mono_decode.py
       - tests/models/test_muse_spark_rocm_loading.py
       - tests/models/test_dsv4_mono_decode.py
-      - tests/models/test_deepseek_v4_vl_stream_weights.py
       - tests/models/hy_v4/test_triton_ihc.py
       - vllm/models/muse_spark/
       - tests/models/test_muse_spark_logits_local.py
@@ -134,6 +134,7 @@ steps:
         models/transformers/test_backend.py
         models/test_registry.py
         models/test_deepseek_v4_vl_rocm.py
+        models/test_deepseek_v4_vl.py
         models/transformers/test_layer_registry_e2e.py
         models/test_deepseek_v41_replay_start.py
         models/test_deepseek_v41_decoder_replay_layers.py
@@ -144,7 +145,6 @@ steps:
         models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
         models/test_dsv41_mono_decode.py
         models/test_muse_spark_rocm_loading.py
-        models/test_deepseek_v4_vl_stream_weights.py
         models/test_muse_spark_logits_local.py
         models/test_muse_spark_rocm_shared_rope.py
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'

tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py

@@ -8,6 +8,7 @@
 import pathlib
 import subprocess
 import sys
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import Mock
 
@@ -17,6 +18,7 @@
 import vllm.model_executor.model_loader.tensorizer
 from tests.utils import VLLM_PATH, RemoteOpenAIServer
 from vllm import SamplingParams
+from vllm.config.load import LoadConfig
 from vllm.engine.arg_utils import EngineArgs
 from vllm.model_executor.model_loader.tensorizer import (
     TensorizerConfig,
@@ -27,6 +29,7 @@
 )
 from vllm.model_executor.model_loader.tensorizer_loader import (
     BLACKLISTED_TENSORIZER_ARGS,
+    TensorizerLoader,
 )
 from vllm.utils.import_utils import PlaceholderModule
 
@@ -303,6 +306,64 @@ def test_vllm_tensorized_model_has_same_outputs(
         assert outputs == deserialized_outputs
 
 
+@torch.inference_mode()
+def test_deserialized_hf_weights_are_ready_for_forward(tmp_path, monkeypatch):
+    """Raw weights need CPU linear post-processing before the first forward."""
+    from vllm.model_executor.layers import linear
+    from vllm.model_executor.layers import utils as layer_utils
+    from vllm.model_executor.model_loader import base_loader
+    from vllm.platforms.cpu import CpuPlatform
+
+    # Exercise the CPU method on accelerator CI hosts too.
+    for module in (linear, layer_utils, base_loader):
+        monkeypatch.setattr(module, "current_platform", CpuPlatform())
+
+    class TinyModel(torch.nn.Sequential):
+        def __init__(self):
+            super().__init__(torch.nn.Linear(3, 2, bias=False))
+            self[0].quant_method = linear.UnquantizedLinearMethod()
+            self[0]._cpu_skip_gemm_dispatch = True
+
+        def load_weights(self, weights):
+            self.load_state_dict(dict(weights))
+
+        def forward(self, inputs):
+            return self[0].quant_method.apply(self[0], inputs)
+
+    source = torch.nn.Sequential(torch.nn.Linear(3, 2, bias=False))
+    source[0].weight.copy_(torch.tensor([[1.0, 2.0, 3.0], [-2.0, 0.5, 4.0]]))
+    inputs = torch.tensor([[1.0, 2.0, 3.0], [-1.0, 0.0, 2.0]])
+    checkpoint = tmp_path / "raw.tensors"
+    with open_stream(checkpoint, "wb+") as stream:
+        serializer = TensorSerializer(stream)
+        serializer.write_module(source)
+        serializer.close()
+
+    load_config = LoadConfig(
+        load_format="tensorizer",
+        model_loader_extra_config={
+            "tensorizer_config": {"tensorizer_uri": str(checkpoint)}
+        },
+    )
+    model_config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization=None,
+        word_embeddings_untied_by_checkpoint=False,
+    )
+    config = SimpleNamespace(
+        model_config=model_config,
+        load_config=load_config,
+        device_config=SimpleNamespace(device=torch.device("cpu")),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        quant_config=None,
+    )
+    monkeypatch.setattr(base_loader, "initialize_model", lambda **kwargs: TinyModel())
+
+    loaded = TensorizerLoader(load_config).load_model(config, model_config)
+
+    torch.testing.assert_close(loaded(inputs), source(inputs))
+
+
 def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
     # For backwards compatibility, ensure Tensorizer can be still be loaded
     # for inference by passing the model reference name, not a local/S3 dir,

tests/model_executor/model_loader/test_expert_substitution.py

@@ -202,27 +202,40 @@ def test_layerwise_reload_is_rejected_and_keeps_weights():
 @pytest.mark.parametrize("missing_value", [False, True])
 def test_tensorizer_finalizes_local_substitution_values(monkeypatch, missing_value):
     """Raw Tensorizer loading must reject an unloaded local constant row."""
-    from vllm.model_executor.model_loader import tensorizer_loader
+    from vllm.model_executor.model_loader import base_loader, tensorizer_loader
 
     model = _cpu_model(local_layers=(0,))
     weights = [(_approx_value(0, 1), torch.zeros(2))]
     if not missing_value:
         weights.append((_approx_value(0, 2), torch.ones(2)))
     finalize = Mock(wraps=model.process_weights_after_loading)
     monkeypatch.setattr(model, "process_weights_after_loading", finalize)
-    monkeypatch.setattr(tensorizer_loader, "initialize_model", lambda **kwargs: model)
-    loader = object.__new__(tensorizer_loader.TensorizerLoader)
+    monkeypatch.setattr(base_loader, "initialize_model", lambda **kwargs: model)
+    monkeypatch.setattr(tensorizer_loader, "is_vllm_tensorized", lambda _: False)
+    load_config = LoadConfig(
+        load_format="tensorizer",
+        model_loader_extra_config={"tensorizer_config": {"tensorizer_uri": "raw"}},
+    )
+    loader = tensorizer_loader.TensorizerLoader(load_config)
     monkeypatch.setattr(loader, "_get_weights_iterator", lambda: iter(weights))
+    model_config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization=None,
+        word_embeddings_untied_by_checkpoint=False,
+    )
     config = SimpleNamespace(
-        model_config=SimpleNamespace(dtype=torch.float32),
-        device_config=SimpleNamespace(device="cpu"),
+        model_config=model_config,
+        load_config=load_config,
+        device_config=SimpleNamespace(device=torch.device("cpu")),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        quant_config=None,
     )
 
     if missing_value:
         with pytest.raises(ValueError, match=r"missing approx_value .* \[2\]"):
-            loader._load_model_serialized_cpu(config)
+            loader.load_model(config, model_config)
     else:
-        assert loader._load_model_serialized_cpu(config) is model
+        assert loader.load_model(config, model_config) is model
         torch.testing.assert_close(
             model.model.layers[0].mlp.experts.values,
             torch.tensor([[0.0, 0.0], [1.0, 1.0]]),

tests/model_executor/model_loader/test_reload.py

@@ -18,11 +18,18 @@
 import vllm.model_executor.model_loader.reload.layerwise as reload_layerwise
 import vllm.model_executor.model_loader.reload.meta as reload_meta
 from vllm.config import ModelConfig
+from vllm.distributed.weight_transfer.sparse_nccl_engine import (
+    SparseNCCLWeightTransferEngine,
+)
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.layers.attention import MMEncoderAttention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+from vllm.model_executor.model_loader.checkpoint_weight_patch import (
+    CheckpointWeightPatch,
+    load_checkpoint_weight_patches,
+)
 from vllm.model_executor.model_loader.reload.layerwise import (
     finalize_layerwise_reload,
     initialize_layerwise_reload,
@@ -93,6 +100,7 @@ def test_attention_sink_model_loading_shards_once(
     model = cls.__new__(cls)
     torch.nn.Module.__init__(model)
     model.hf_to_vllm_mapper = WeightsMapper()
+    model._context_kv_proj = None  # set by DFlashQwen3Model.__init__
     model.model = torch.nn.Module()
     model.model.mtp = torch.nn.Module()
     sink = torch.nn.Parameter(torch.zeros(4), requires_grad=False)
@@ -1542,3 +1550,68 @@ def test_online_quantize_reload(
         mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
         add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
         assert add_perp < mul_perp
+
+
+class _TwoLayerModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.a = torch.nn.Linear(2, 2, bias=False)
+        self.b = torch.nn.Linear(2, 2, bias=False)
+        self.hook_saw: list[tuple[float, float]] = []
+
+    def load_weights(self, weights) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            param = params[name]
+            getattr(param, "weight_loader", default_weight_loader)(param, tensor)
+            loaded.add(name)
+        return loaded
+
+    def process_weights_after_loading(self) -> None:
+        a, b = self.a.weight.mean().item(), self.b.weight.mean().item()
+        self.hook_saw.append((a, b))
+
+
+def _full(value: float) -> torch.Tensor:
+    return torch.full((2, 2), value)
+
+
+def _loaded_model() -> _TwoLayerModel:
+    model = _TwoLayerModel()
+    model.load_weights([("a.weight", _full(1.0)), ("b.weight", _full(2.0))])
+    return model
+
+
+class TestModelHookOnSplitLoad:
+    def test_layerwise_reload(self):
+        model = _loaded_model()
+        record_metadata_for_reloading(model)
+        config = SimpleNamespace(dtype=torch.float32)
+
+        initialize_layerwise_reload(model)
+        model.load_weights([("a.weight", _full(3.0))])
+        model.load_weights([("b.weight", _full(4.0))])
+        assert model.hook_saw == []
+        finalize_layerwise_reload(model, config)
+        assert model.hook_saw == [(3.0, 4.0)]
+
+        # Partial update: the untouched layer is restored before the hook runs.
+        initialize_layerwise_reload(model)
+        model.load_weights([("a.weight", _full(5.0))])
+        finalize_layerwise_reload(model, config)
+        assert model.hook_saw == [(3.0, 4.0), (5.0, 4.0)]
+
+    def test_sparse_update(self):
+        model = _loaded_model()
+        engine = object.__new__(SparseNCCLWeightTransferEngine)
+        engine.model = model
+
+        engine.start_weight_update()
+        for name, value in (("a.weight", 5.0), ("b.weight", 6.0)):
+            values, indices = torch.full((4,), value), torch.arange(4)
+            patch = CheckpointWeightPatch(name, (2, 2), torch.float32, values, indices)
+            load_checkpoint_weight_patches(model, [patch])
+        assert model.hook_saw == []
+        engine.finish_weight_update()
+        assert model.hook_saw == [(5.0, 6.0)]

tests/models/test_deepseek_v4_vl.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""VL wrapper finalization contract with the language model."""
+
+import importlib
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.models.utils import WeightsMapper
+
+
+class _FakeLanguageModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.tensor_a = nn.Parameter(torch.zeros(1))
+        self.tensor_c = nn.Parameter(torch.zeros(1))
+        self.finalized_values: list[tuple[float, float]] = []
+
+    def process_weights_after_loading(self) -> None:
+        self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item()))
+
+
+@pytest.mark.parametrize(
+    "module,cls_name",
+    [
+        (
+            "vllm.models.deepseek_v4.common.vl_model",
+            "DeepseekV4ForConditionalGeneration",
+        ),
+        ("vllm.models.deepseek_v41.nvidia.vl_model", "DeepseekV41ForCausalLM"),
+        ("vllm.models.deepseek_v41.amd.vl_model", "DeepseekV41ForCausalLM"),
+    ],
+    ids=["v4", "v41-nvidia", "v41-amd"],
+)
+def test_vl_wrapper_delegates_finalization_to_hook(module, cls_name) -> None:
+    wrapper_cls = getattr(importlib.import_module(module), cls_name)
+
+    model = object.__new__(wrapper_cls)
+    nn.Module.__init__(model)
+    model.language_model = _FakeLanguageModel()
+    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])
+        assert model.language_model.tensor_a.item() == 1.0
+        yield "vision.tensor_b", torch.tensor([2.0])
+        assert model.vision.tensor_b.item() == 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.finalized_values == []
+
+    model.process_weights_after_loading()
+
+    assert model.language_model.finalized_values == [(1.0, 3.0)]
+
+    # A reload reruns the hook: no one-shot guard may skip it.
+    model.process_weights_after_loading()
+    assert model.language_model.finalized_values == [(1.0, 3.0)] * 2