[Bugfix] Fix level-2 sleep/wake/reload with enable_lora=True

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 @@
 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 @@
 """
 
 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 @@ def __init__(self, raise_on_update: bool = False):
         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 @@ def reset_weight_update_target(self) -> None:
 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 @@
 
 
 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 @@ def create_lora_weights(
         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 @@ def remove_all_adapters(self):
         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

@@ -57,7 +57,7 @@
     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
@@ -251,6 +251,16 @@
 
 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
@@ -5541,7 +5551,10 @@ def reload_weights(
             )
 
         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
@@ -5575,10 +5588,12 @@ def reload_weights(
             )
             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