diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 9ab167ab9a..dcb6ad8c94 100644 --- a/docs/features/sleep_mode.md +++ b/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 diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index e4b3fb6584..0618562b3e 100644 --- a/tests/basic_correctness/test_mem.py +++ b/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(): diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 95964e9f28..ab09514274 100644 --- a/tests/lora/test_layers.py +++ b/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) diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 834273b8c5..43fa1975bd 100644 --- a/tests/lora/test_lora_manager.py +++ b/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): diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index 6a97d64c6b..23c0b4bf61 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/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(): diff --git a/vllm/lora/layers/base.py b/vllm/lora/layers/base.py index 26d2fb46d1..85efbff1b0 100644 --- a/vllm/lora/layers/base.py +++ b/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] diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index 496413c21f..03f1b7ffdd 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/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 diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index ba545bc9c7..bd5e91cfa5 100644 --- a/vllm/lora/model_manager.py +++ b/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: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 64fefd7a6a..cab2ce96c4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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 diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index ed11c0def9..bee85c6eca 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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() diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 3a14abfc35..27e850b66a 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/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,