[QeRL] Layerwise Reloading (#32133)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

tests/conftest.py

@@ -27,7 +27,7 @@ import threading
 from collections.abc import Generator
 from contextlib import nullcontext
 from enum import Enum
-from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING
+from typing import Any, Callable, TypedDict, TypeVar, cast, TYPE_CHECKING, Optional
 
 import numpy as np
 import pytest
@@ -1024,7 +1024,9 @@ class VllmRunner:
             **kwargs,
         )
 
-    def generate_prompt_perplexity(self, prompts: list[str]) -> list[float]:
+    def generate_prompt_perplexity(
+        self, prompts: list[str], mask: Optional[list[str]] = None
+    ) -> list[float]:
         """
         Return the perplexity score associated with generating the prompts
 
@@ -1035,13 +1037,20 @@ class VllmRunner:
             prompts, max_tokens=1, num_logprobs=None, num_prompt_logprobs=0
         )
 
+        mask_prefix_lens = (
+            [len(self.llm.get_tokenizer()(prefix)["input_ids"]) for prefix in mask]
+            if mask is not None
+            else [0 for _ in range(len(prompts))]
+        )
+
         perplexities = []
-        for output in outputs:
+        for output, mask_prefix_len in zip(outputs, mask_prefix_lens):
             output = cast(TokensTextLogprobsPromptLogprobs, output)
             token_datas = cast(list[dict[int, Logprob] | None], output[3])
             assert token_datas[0] is None
+
             token_log_probs = []
-            for token_data in token_datas[1:]:
+            for token_data in token_datas[mask_prefix_len + 1 :]:
                 assert token_data is not None
                 assert len(token_data) == 1
                 token_log_prob = list(token_data.values())[0].logprob
@@ -1122,6 +1131,9 @@ class VllmRunner:
     def get_llm(self) -> LLM:
         return self.llm
 
+    def collective_rpc(self, *args, **kwargs):
+        return self.llm.collective_rpc(*args, **kwargs)
+
     def __enter__(self):
         return self
 
@@ -1532,3 +1544,9 @@ def use_fresh_inductor_cache():
     """
     with fresh_cache():
         yield
+
+
+@pytest.fixture(scope="function")
+def enable_pickle(monkeypatch):
+    """`LLM.apply_model` requires pickling a function."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

tests/model_executor/model_loader/test_reload.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import gc
+import inspect
+from weakref import WeakKeyDictionary, ref
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm.model_executor.model_loader.reload.meta import (
+    capture_layer_to_meta,
+    get_numel_loaded,
+    materialize_layer,
+    materialize_meta_tensor,
+    restore_layer_on_meta,
+    to_meta_tensor,
+)
+from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo
+from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import cuda_device_count_stateless
+
+
+def test_move_metatensors():
+    tensor = torch.empty((1, 2, 3))
+    meta_tensor = to_meta_tensor(tensor)
+    materialized_tensor = materialize_meta_tensor(meta_tensor)
+
+    assert meta_tensor.device.type == "meta"
+    assert tensor.device == materialized_tensor.device
+
+    assert tensor.dtype == meta_tensor.dtype == materialized_tensor.dtype
+    assert tensor.shape == meta_tensor.shape == materialized_tensor.shape
+    assert tensor.__class__ == meta_tensor.__class__ == materialized_tensor.__class__
+    assert tensor.__dict__ == meta_tensor.__dict__ == materialized_tensor.__dict__
+
+
+def test_reload_lifecycle():
+    layer = torch.nn.Linear(2, 3)
+    info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))
+
+    restore_layer_on_meta(layer, info)
+    for name, tensor in get_layer_tensors(layer).items():
+        meta_tensor = getattr(layer, name)
+        assert tensor.dtype == meta_tensor.dtype
+        assert tensor.shape == meta_tensor.shape
+        assert tensor.__class__ == meta_tensor.__class__
+        assert tensor.__dict__ == meta_tensor.__dict__
+
+    materialize_layer(layer)
+    for name, tensor in get_layer_tensors(layer).items():
+        materialized_tensor = getattr(layer, name)
+        assert tensor.dtype == materialized_tensor.dtype
+        assert tensor.shape == materialized_tensor.shape
+        assert tensor.__class__ == materialized_tensor.__class__
+        assert tensor.__dict__ == materialized_tensor.__dict__
+
+
+def test_model_cleanup(dist_init, default_vllm_config):
+    layer = QKVParallelLinear(2, 3, 4)
+    assert layer.weight.weight_loader.__self__ is layer
+    info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))
+
+    mock_info_dict: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
+        WeakKeyDictionary()
+    )
+    mock_info_dict[layer] = info
+    layer_ref = ref(layer)
+
+    del layer
+    gc.collect()
+
+    assert layer_ref() is None
+    assert len(mock_info_dict) == 0
+
+
+def test_get_numel_loaded():
+    param = torch.empty(10, device="meta")
+    loaded_weight = torch.empty(10)
+
+    def complex_weight_loader(param, loaded_weight):
+        param[:3] = loaded_weight[:3]
+        param[5:8] = loaded_weight[5:8]
+        return "value"
+
+    args = inspect.signature(complex_weight_loader).bind(param, loaded_weight)
+    num_loaded, ret = get_numel_loaded(complex_weight_loader, args)
+    assert num_loaded == 6
+    assert ret == "value"
+
+
+@pytest.mark.parametrize("tp_size", [2])
+@pytest.mark.parametrize(
+    "base_model,mul_model,add_model",
+    [
+        (
+            "Qwen/Qwen3-0.6B",
+            "inference-optimization/Qwen3-0.6B-debug-multiply",
+            "inference-optimization/Qwen3-0.6B-debug-add",
+        ),
+        (
+            "inference-optimization/Qwen3-0.6B-FP8_BLOCK",
+            "inference-optimization/Qwen3-0.6B-debug-multiply-FP8_BLOCK",
+            "inference-optimization/Qwen3-0.6B-debug-add-FP8_BLOCK",
+        ),
+        (
+            "inference-optimization/Qwen3-0.6B-W4A16-G128",
+            "inference-optimization/Qwen3-0.6B-debug-multiply-W4A16-G128",
+            "inference-optimization/Qwen3-0.6B-debug-add-W4A16-G128",
+        ),
+        (
+            "inference-optimization/DeepSeek-V3-debug-empty",
+            "inference-optimization/DeepSeek-V3-debug-multiply",
+            "inference-optimization/DeepSeek-V3-debug-add",
+        ),
+        (
+            "inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC",
+            "inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC",
+            "inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC",
+        ),
+        (
+            "inference-optimization/DeepSeek-V3-debug-empty-NVFP4A16",
+            "inference-optimization/DeepSeek-V3-debug-multiply-NVFP4A16",
+            "inference-optimization/DeepSeek-V3-debug-add-NVFP4A16",
+        ),
+    ],
+)
+def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
+    if cuda_device_count_stateless() < tp_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    if "FP8" in base_model and not current_platform.supports_fp8():
+        pytest.skip(reason="Requires FP8 support")
+
+    with vllm_runner(
+        model_name=base_model,
+        tensor_parallel_size=tp_size,
+        enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
+        enable_prefix_caching=False,
+    ) as llm:
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
+        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 mul_perp < add_perp
+
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})
+        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

tests/quantization/test_torchao.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import importlib.metadata
 import importlib.util
 
 import pytest
 import torch
 
+from vllm.model_executor.model_loader import get_model_loader
 from vllm.platforms import current_platform
 
 DTYPE = ["bfloat16"]
@@ -105,8 +105,8 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):
 
 
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
-def test_online_quant_config_dict_json(vllm_runner):
-    """Testing on the fly quantization, load_weights integration point,
+def test_online_quant_config_dict_json(vllm_runner, enable_pickle):
+    """Testing online quantization, load_weights integration point,
     with config dict serialized to json string
     """
     torch._dynamo.reset()
@@ -135,7 +135,18 @@ def test_online_quant_config_dict_json(vllm_runner):
     ) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
 
-        assert output
+        load_config = llm.llm.llm_engine.vllm_config.load_config
+        model_config = llm.llm.llm_engine.vllm_config.model_config
+
+        def load_weights(model):
+            model_loader = get_model_loader(load_config)
+            weights_iterator = model_loader.get_all_weights(model_config, model)
+            model.load_weights(weights_iterator)
+
+        llm.apply_model(load_weights)
+
+        reload_output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
+        assert output[0][0] == reload_output[0][0]
 
 
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

tests/v1/worker/test_gpu_model_runner.py

@@ -543,7 +543,7 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
 
 
 def test_reload_weights_before_load_model(model_runner):
-    with pytest.raises(AssertionError):
+    with pytest.raises(ValueError):
         model_runner.reload_weights()