[Bugfix] Fix level-2 sleep/wake/reload with enable_lora=True (#39935)

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>
Co-authored-by: Dakai An <dakaian108@gmail.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

docs/features/sleep_mode.md

@@ -18,7 +18,7 @@ Key benefits:
 
 ## Sleep levels
 
-Level 1 sleep will offload the model weights and discard the KV cache. The content of KV cache is forgotten. Level 1 sleep is good for sleeping and waking up the engine to run the same model again. The model weights are backed up in CPU memory. Please make sure there's enough CPU memory to store the model weights. Level 2 sleep will discard both the model weights and the KV cache (while the model's buffers are kept in CPU, like rope scaling tensors). The content of both the model weights and KV cache is forgotten. Level 2 sleep is good for sleeping and waking up the engine to run a different model or update the model, where previous model weights are not needed, e.g. RLHF weight update.
+Level 1 sleep will offload the model weights and discard the KV cache. The content of KV cache is forgotten. Level 1 sleep is good for sleeping and waking up the engine to run the same model again. The model weights are backed up in CPU memory. Please make sure there's enough CPU memory to store the model weights. Level 2 sleep will discard both the model weights and the KV cache (while the model's buffers are kept in CPU, like rope scaling tensors). The content of both the model weights and KV cache is forgotten. Level 2 sleep is good for sleeping and waking up the engine to run a different model or update the model, where previous model weights are not needed, e.g. RLHF weight update. Level 2 sleep is also useful when there is not enough CPU memory to back up the model weights, e.g. when a colocated trainer already uses CPU memory for offloading its own state; since nothing is backed up, restore the weights after waking up with `collective_rpc("reload_weights")`.
 
 ## Usage
 

tests/basic_correctness/test_mem.py

@@ -209,6 +209,99 @@ def test_deep_sleep():
     assert output[0].outputs[0].text == output2[0].outputs[0].text
 
 
+@create_new_process_for_each_test()
+def test_deep_sleep_lora():
+    """Level-2 sleep/wake/reload with enable_lora=True.
+
+    LoRA wrapping moves parameters under base_layer and adds LoRA
+    stacked tensors that are plain attributes, not restored by the
+    reload machinery — reload must forward checkpoint weights through
+    the wrappers and reset the LoRA state afterwards.
+    """
+    model = "hmellor/tiny-random-LlamaForCausalLM"
+    llm = LLM(
+        model,
+        enable_sleep_mode=True,
+        enable_lora=True,
+        max_lora_rank=8,
+        enforce_eager=True,
+    )
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    output = llm.generate(prompt, sampling_params)
+
+    # Level-2 sleep discards all GPU memory
+    llm.sleep(level=2)
+
+    # Reload weights from checkpoint
+    llm.wake_up(tags=["weights"])
+    llm.collective_rpc("reload_weights")
+    llm.wake_up(tags=["kv_cache"])
+    output2 = llm.generate(prompt, sampling_params)
+    assert output[0].outputs[0].text == output2[0].outputs[0].text
+
+    # Multiple cycles should not accumulate corruption
+    for _ in range(3):
+        llm.sleep(level=2)
+        llm.wake_up(tags=["weights"])
+        llm.collective_rpc("reload_weights")
+        llm.wake_up(tags=["kv_cache"])
+    output3 = llm.generate(prompt, sampling_params)
+    assert output[0].outputs[0].text == output3[0].outputs[0].text
+
+
+def _lora_logits_mapping_present(model) -> bool:
+    from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
+
+    return any(
+        isinstance(m, LogitsProcessorWithLoRA)
+        and m.sharded_to_full_mapping_gpu is not None
+        for m in model.modules()
+    )
+
+
+@create_new_process_for_each_test()
+def test_deep_sleep_lora_tp2(num_gpus_available, monkeypatch):
+    """Level-2 sleep/wake/reload with enable_lora=True and TP=2.
+
+    With TP > 1 the LoRA logits processor carries
+    ``sharded_to_full_mapping_gpu``, a permanent index mapping used to
+    reorder gathered logits. Like the LoRA stacked tensors it is a plain
+    attribute allocated in the sleep-mode pool, so level-2 sleep destroys
+    its contents — it must be restored after reload.
+    """
+    if num_gpus_available < 2:
+        pytest.skip("Requires at least 2 GPUs")
+
+    # Needed for apply_model to reach the multiproc TP workers below.
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+    model = "hmellor/tiny-random-LlamaForCausalLM"
+    llm = LLM(
+        model,
+        enable_sleep_mode=True,
+        enable_lora=True,
+        max_lora_rank=8,
+        tensor_parallel_size=2,
+        enforce_eager=True,
+    )
+
+    # Guard against this test silently not exercising the TP>1 reindex
+    # path (e.g. if lm_head wrapping conditions change).
+    assert all(llm.apply_model(_lora_logits_mapping_present))
+
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    output = llm.generate(prompt, sampling_params)
+
+    llm.sleep(level=2)
+    llm.wake_up(tags=["weights"])
+    llm.collective_rpc("reload_weights")
+    llm.wake_up(tags=["kv_cache"])
+    output2 = llm.generate(prompt, sampling_params)
+    assert output[0].outputs[0].text == output2[0].outputs[0].text
+
+
 @create_new_process_for_each_test()
 def test_deep_sleep_async():
     async def test():

tests/lora/test_layers.py

@@ -83,6 +83,21 @@ NUM_RANDOM_SEEDS = 2
 VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS = 2
 
 
+def test_base_layer_with_lora_delegates_load_weights():
+    class BaseLayer(torch.nn.Module):
+        def load_weights(self, weights):
+            self.loaded_weights = list(weights)
+            return {"weight"}
+
+    base_layer = BaseLayer()
+    layer = BaseLayerWithLoRA()
+    layer.base_layer = base_layer
+    weights = [("weight", torch.ones(1))]
+
+    assert layer.load_weights(weights) == {"weight"}
+    assert base_layer.loaded_weights == weights
+
+
 @pytest.fixture(autouse=True)
 def clean_cache_reset_device(reset_default_device):
     # Release any memory we might be holding on to. CI runs OOMs otherwise.
@@ -509,6 +524,38 @@ def test_lm_head_logits_processor_invalid_vocab_size(
         lora_logits_processor.create_lora_weights(max_loras, lora_config)
 
 
+@torch.inference_mode()
+@pytest.mark.parametrize("device", DEVICES)
+def test_lm_head_reset_sharded_to_full_mapping(
+    default_vllm_config, dist_init, device
+) -> None:
+    if current_platform.is_cuda_alike() or current_platform.is_xpu():
+        torch.accelerator.set_device_index(device)
+
+    torch.set_default_device(device)
+    max_loras = 8
+    vocab_size = 1024
+    lora_config = LoRAConfig(
+        max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
+    )
+
+    logits_processor = LogitsProcessor(vocab_size)
+    sharded_to_full_mapping = list(reversed(range(vocab_size)))
+    lora_logits_processor = LogitsProcessorWithLoRA(
+        logits_processor, 1024, torch.float16, device, sharded_to_full_mapping
+    )
+    lora_logits_processor.create_lora_weights(max_loras, lora_config)
+
+    lora_logits_processor.sharded_to_full_mapping_gpu.fill_(-1)
+
+    lora_logits_processor.reset_sharded_to_full_mapping()
+
+    torch.testing.assert_close(
+        lora_logits_processor.sharded_to_full_mapping_gpu,
+        torch.tensor(sharded_to_full_mapping, dtype=torch.long),
+    )
+
+
 @torch.inference_mode()
 @pytest.mark.parametrize("num_loras", [1, 2, 4])
 @pytest.mark.parametrize("device", DEVICES)

tests/lora/test_lora_manager.py

@@ -699,6 +699,10 @@ def test_set_adapter_mapping_refreshes_after_slot_reassignment(
     manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2)))
     assert punica_wrapper.token_lora_indices.tolist() == [1, 0]
 
+    manager.remove_all_adapters()
+    assert manager._last_mapping is None
+    assert manager._last_slot_layout is None
+
 
 @pytest.mark.parametrize("device", DEVICES)
 def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):

tests/v1/worker/test_gpu_worker_weight_transfer.py

@@ -8,8 +8,12 @@ session is active. These tests verify that delegation and the session guard.
 """
 
 import pytest
+import torch
+import torch.nn as nn
 
 from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.lora.layers import BaseLayerWithLoRA
+from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload
 from vllm.v1.worker.gpu_worker import Worker
 
 
@@ -21,6 +25,7 @@ class _RecordingEngine:
         self.started = False
         self.finished = False
         self.reset_count = 0
+        self.supports_draft_weight_update = False
         self.update_calls: list[dict] = []
         self.seen_configs: list[VllmConfig] = []
 
@@ -48,16 +53,22 @@ class _RecordingEngine:
 class _RecordingModelRunner:
     def __init__(self) -> None:
         self.seen_config: VllmConfig | None = None
+        self.reset_lora_calls = 0
 
     def reload_weights(self) -> None:
         self.seen_config = get_current_vllm_config()
 
+    def reset_lora_state(self) -> None:
+        self.reset_lora_calls += 1
+
 
 def _make_worker(engine: _RecordingEngine | None) -> Worker:
     worker = object.__new__(Worker)
     worker.vllm_config = VllmConfig()
     worker.weight_transfer_engine = engine
     worker._weight_update_active = False
+    worker._weight_update_is_draft = False
+    worker.model_runner = _RecordingModelRunner()
     return worker
 
 
@@ -71,6 +82,22 @@ def test_reload_weights_sets_current_config():
     assert model_runner.seen_config is worker.vllm_config
 
 
+def test_reload_parameter_lookup_preserves_lora_module_names():
+    base_layer = nn.Module()
+    qweight = nn.Parameter(torch.ones(1))
+    base_layer.register_parameter("qweight", qweight)
+    wrapper = BaseLayerWithLoRA()
+    wrapper.base_layer = base_layer
+    model = nn.Module()
+    model.proj = wrapper
+
+    named_parameters = dict(model.named_parameters())
+    assert set(named_parameters) == {"proj.base_layer.qweight"}
+    assert named_parameters["proj.base_layer.qweight"] is qweight
+    assert model.get_parameter("proj.base_layer.qweight") is qweight
+    assert _get_parameter_for_reload(model, "proj.qweight") is qweight
+
+
 def test_start_update_finish_delegates_to_engine():
     engine = _RecordingEngine()
     worker = _make_worker(engine)
@@ -88,6 +115,19 @@ def test_start_update_finish_delegates_to_engine():
     assert engine.reset_count == 1
     assert worker._weight_update_active is False
     assert engine.seen_configs == [worker.vllm_config] * 3
+    assert worker.model_runner.reset_lora_calls == 1
+
+
+def test_finish_draft_session_keeps_lora_state():
+    engine = _RecordingEngine()
+    engine.supports_draft_weight_update = True
+    worker = _make_worker(engine)
+    worker._set_draft_weight_update_target = lambda: None
+
+    Worker.start_draft_weight_update(worker)
+    Worker.finish_weight_update(worker)
+
+    assert worker.model_runner.reset_lora_calls == 0
 
 
 def test_double_start_raises():

vllm/lora/layers/base.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterable
 from typing import TYPE_CHECKING, overload
 
 import torch
@@ -14,6 +15,18 @@ if TYPE_CHECKING:
 
 
 class BaseLayerWithLoRA(nn.Module):
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[str]:
+        """Load checkpoint weights into the wrapped base layer."""
+        base_load_weights = getattr(self.base_layer, "load_weights", None)
+        if callable(base_load_weights):
+            return base_load_weights(weights)
+
+        from vllm.model_executor.models.utils import AutoWeightsLoader
+
+        return AutoWeightsLoader(self.base_layer).load_weights(weights)
+
     @overload
     def slice_lora_a(
         self, lora_a: list[torch.Tensor | None]

vllm/lora/layers/logits_processor.py

@@ -118,6 +118,18 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
         else:
             self.sharded_to_full_mapping_gpu = None
 
+    def reset_sharded_to_full_mapping(self) -> None:
+        """Restore the TP logits mapping after its GPU memory is reused."""
+        mapping_gpu = self.sharded_to_full_mapping_gpu
+        if mapping_gpu is not None:
+            mapping_gpu.copy_(
+                torch.tensor(
+                    self.sharded_to_full_mapping,
+                    device=mapping_gpu.device,
+                    dtype=mapping_gpu.dtype,
+                )
+            )
+
     def reset_lora(self, index: int):
         self.lora_a_stacked[index] = 0
         self.lora_b_stacked[index] = 0

vllm/lora/model_manager.py

@@ -387,6 +387,8 @@ class LoRAModelManager:
         self._registered_adapters.clear()
         self.lora_index_to_id = [None] * self.lora_slots
         self._active_adapters.clear()
+        self._last_mapping = None
+        self._last_slot_layout = None
 
     def _create_lora_modules(self):
         def _parent_module(module_name: str) -> str:

vllm/v1/worker/gpu_model_runner.py

@@ -56,7 +56,7 @@ from vllm.forward_context import (
     set_forward_context,
 )
 from vllm.logger import init_logger
-from vllm.lora.layers import LoRAMapping, LoRAMappingType
+from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType
 from vllm.model_executor.layers.attention import Attention, MLAAttention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
@@ -250,6 +250,16 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+
+def _get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter:
+    """Resolve checkpoint names without changing the model's module tree."""
+    module_name, _, parameter_name = name.rpartition(".")
+    module = model.get_submodule(module_name)
+    if isinstance(module, BaseLayerWithLoRA):
+        module = module.base_layer
+    return module.get_parameter(parameter_name)
+
+
 AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata]
 # list when ubatching is enabled
 PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict
@@ -5535,7 +5545,10 @@ class GPUModelRunner(
             )
 
         model = self.get_model()
-        weights_to_load = {name for name, _ in model.named_parameters()}
+        weights_to_load = {
+            name.replace(".base_layer.", ".") if self.lora_config else name
+            for name, _ in model.named_parameters()
+        }
         counter_before_reloading = time.perf_counter()
 
         # load weights from disk if none are provided
@@ -5572,10 +5585,12 @@ class GPUModelRunner(
             )
             loaded_weights = set()
             for name, loaded_weight in weights_iterator:
-                param = model.get_parameter(name)  # TODO: buffers?
+                param = _get_parameter_for_reload(model, name)  # TODO: buffers?
                 param.copy_(loaded_weight)
                 loaded_weights.add(name)
 
+        self.reset_lora_state()
+
         # logging and validation
         counter_after_reloading = time.perf_counter()
         diff_seconds = counter_after_reloading - counter_before_reloading

vllm/v1/worker/gpu_worker.py

@@ -160,6 +160,7 @@ class Worker(WorkerBase):
         # is available, since the engine needs a reference to the model.
         self.weight_transfer_engine: WeightTransferEngine | None = None
         self._weight_update_active = False
+        self._weight_update_is_draft = False
 
         # Torch/CUDA profiler. Enabled and configured through profiler_config.
         # Profiler wrapper is created lazily in profile() when start is called,
@@ -1272,6 +1273,7 @@ class Worker(WorkerBase):
             self.weight_transfer_engine.reset_weight_update_target()
             raise
         self._weight_update_active = True
+        self._weight_update_is_draft = is_draft
 
     def update_weights(self, update_info: dict) -> None:
         """
@@ -1316,6 +1318,10 @@ class Worker(WorkerBase):
             self.weight_transfer_engine.reset_weight_update_target()
             self._weight_update_active = False
 
+        # Weight transfer bypasses GPUModelRunner.reload_weights().
+        if not self._weight_update_is_draft:
+            self.model_runner.reset_lora_state()
+
     def shutdown(self) -> None:
         gc.unfreeze()
 

vllm/v1/worker/lora_model_runner_mixin.py

@@ -4,6 +4,7 @@
 Define LoRA functionality mixin for model runners.
 """
 
+from collections.abc import Callable
 from contextlib import contextmanager
 from typing import TypeAlias
 
@@ -28,6 +29,21 @@ logger = init_logger(__name__)
 
 # Defined as a mixin for GPUModelRunner
 class LoRAModelRunnerMixin:
+    lora_config: LoRAConfig | None
+    get_model: Callable[[], nn.Module]
+
+    def reset_lora_state(self) -> None:
+        """Invalidate LoRA state after base weights are replaced."""
+        if not self.lora_config:
+            return
+
+        from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
+
+        self.lora_manager.remove_all_adapters()
+        for module in self.get_model().modules():
+            if isinstance(module, LogitsProcessorWithLoRA):
+                module.reset_sharded_to_full_mapping()
+
     def load_lora_model(
         self,
         model: nn.Module,