@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b50fdebce0

[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>

Silen Naihin committed Aug 5, 2026 at 18:13 UTC b50fdebce0698c1f2d57ba7ff94211323655557e

11 files changed +252 -4

docs/features/sleep_mode.md

+1 -1

@@ -18,7 +18,7 @@ Key benefits:

19 ## Sleep levels

21 -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.

21 +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")`.

23 ## Usage

tests/basic_correctness/test_mem.py

+93

@@ -209,6 +209,99 @@ def test_deep_sleep():

209 assert output[0].outputs[0].text == output2[0].outputs[0].text

210

211

212 +@create_new_process_for_each_test()

213 +def test_deep_sleep_lora():

214 + """Level-2 sleep/wake/reload with enable_lora=True.

215 +

216 + LoRA wrapping moves parameters under base_layer and adds LoRA

217 + stacked tensors that are plain attributes, not restored by the

218 + reload machinery — reload must forward checkpoint weights through

219 + the wrappers and reset the LoRA state afterwards.

220 + """

221 + model = "hmellor/tiny-random-LlamaForCausalLM"

222 + llm = LLM(

223 + model,

224 + enable_sleep_mode=True,

225 + enable_lora=True,

226 + max_lora_rank=8,

227 + enforce_eager=True,

228 + )

229 + prompt = "How are you?"

230 + sampling_params = SamplingParams(temperature=0, max_tokens=10)

231 + output = llm.generate(prompt, sampling_params)

232 +

233 + # Level-2 sleep discards all GPU memory

234 + llm.sleep(level=2)

235 +

236 + # Reload weights from checkpoint

237 + llm.wake_up(tags=["weights"])

238 + llm.collective_rpc("reload_weights")

239 + llm.wake_up(tags=["kv_cache"])

240 + output2 = llm.generate(prompt, sampling_params)

241 + assert output[0].outputs[0].text == output2[0].outputs[0].text

242 +

243 + # Multiple cycles should not accumulate corruption

244 + for _ in range(3):

245 + llm.sleep(level=2)

246 + llm.wake_up(tags=["weights"])

247 + llm.collective_rpc("reload_weights")

248 + llm.wake_up(tags=["kv_cache"])

249 + output3 = llm.generate(prompt, sampling_params)

250 + assert output[0].outputs[0].text == output3[0].outputs[0].text

251 +

252 +

253 +def _lora_logits_mapping_present(model) -> bool:

254 + from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA

255 +

256 + return any(

257 + isinstance(m, LogitsProcessorWithLoRA)

258 + and m.sharded_to_full_mapping_gpu is not None

259 + for m in model.modules()

260 + )

261 +

262 +

263 +@create_new_process_for_each_test()

264 +def test_deep_sleep_lora_tp2(num_gpus_available, monkeypatch):

265 + """Level-2 sleep/wake/reload with enable_lora=True and TP=2.

266 +

267 + With TP > 1 the LoRA logits processor carries

268 + ``sharded_to_full_mapping_gpu``, a permanent index mapping used to

269 + reorder gathered logits. Like the LoRA stacked tensors it is a plain

270 + attribute allocated in the sleep-mode pool, so level-2 sleep destroys

271 + its contents — it must be restored after reload.

272 + """

273 + if num_gpus_available < 2:

274 + pytest.skip("Requires at least 2 GPUs")

275 +

276 + # Needed for apply_model to reach the multiproc TP workers below.

277 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

278 +

279 + model = "hmellor/tiny-random-LlamaForCausalLM"

280 + llm = LLM(

281 + model,

282 + enable_sleep_mode=True,

283 + enable_lora=True,

284 + max_lora_rank=8,

285 + tensor_parallel_size=2,

286 + enforce_eager=True,

287 + )

288 +

289 + # Guard against this test silently not exercising the TP>1 reindex

290 + # path (e.g. if lm_head wrapping conditions change).

291 + assert all(llm.apply_model(_lora_logits_mapping_present))

292 +

293 + prompt = "How are you?"

294 + sampling_params = SamplingParams(temperature=0, max_tokens=10)

295 + output = llm.generate(prompt, sampling_params)

296 +

297 + llm.sleep(level=2)

298 + llm.wake_up(tags=["weights"])

299 + llm.collective_rpc("reload_weights")

300 + llm.wake_up(tags=["kv_cache"])

301 + output2 = llm.generate(prompt, sampling_params)

302 + assert output[0].outputs[0].text == output2[0].outputs[0].text

303 +

304 +

305 @create_new_process_for_each_test()

306 def test_deep_sleep_async():

307 async def test():

tests/lora/test_layers.py

+47

@@ -83,6 +83,21 @@ NUM_RANDOM_SEEDS = 2

83 VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS = 2

86 +def test_base_layer_with_lora_delegates_load_weights():

87 + class BaseLayer(torch.nn.Module):

88 + def load_weights(self, weights):

89 + self.loaded_weights = list(weights)

90 + return {"weight"}

91 +

92 + base_layer = BaseLayer()

93 + layer = BaseLayerWithLoRA()

94 + layer.base_layer = base_layer

95 + weights = [("weight", torch.ones(1))]

96 +

97 + assert layer.load_weights(weights) == {"weight"}

98 + assert base_layer.loaded_weights == weights

99 +

100 +

101 @pytest.fixture(autouse=True)

102 def clean_cache_reset_device(reset_default_device):

103 # 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(

524 lora_logits_processor.create_lora_weights(max_loras, lora_config)

525

526

527 +@torch.inference_mode()

528 +@pytest.mark.parametrize("device", DEVICES)

529 +def test_lm_head_reset_sharded_to_full_mapping(

530 + default_vllm_config, dist_init, device

531 +) -> None:

532 + if current_platform.is_cuda_alike() or current_platform.is_xpu():

533 + torch.accelerator.set_device_index(device)

534 +

535 + torch.set_default_device(device)

536 + max_loras = 8

537 + vocab_size = 1024

538 + lora_config = LoRAConfig(

539 + max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16

540 + )

541 +

542 + logits_processor = LogitsProcessor(vocab_size)

543 + sharded_to_full_mapping = list(reversed(range(vocab_size)))

544 + lora_logits_processor = LogitsProcessorWithLoRA(

545 + logits_processor, 1024, torch.float16, device, sharded_to_full_mapping

546 + )

547 + lora_logits_processor.create_lora_weights(max_loras, lora_config)

548 +

549 + lora_logits_processor.sharded_to_full_mapping_gpu.fill_(-1)

550 +

551 + lora_logits_processor.reset_sharded_to_full_mapping()

552 +

553 + torch.testing.assert_close(

554 + lora_logits_processor.sharded_to_full_mapping_gpu,

555 + torch.tensor(sharded_to_full_mapping, dtype=torch.long),

556 + )

557 +

558 +

559 @torch.inference_mode()

560 @pytest.mark.parametrize("num_loras", [1, 2, 4])

561 @pytest.mark.parametrize("device", DEVICES)

tests/lora/test_lora_manager.py

@@ -699,6 +699,10 @@ def test_set_adapter_mapping_refreshes_after_slot_reassignment(

699 manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2)))

700 assert punica_wrapper.token_lora_indices.tolist() == [1, 0]

701

702 + manager.remove_all_adapters()

703 + assert manager._last_mapping is None

704 + assert manager._last_slot_layout is None

705 +

706

707 @pytest.mark.parametrize("device", DEVICES)

708 def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):

tests/v1/worker/test_gpu_worker_weight_transfer.py

+40

@@ -8,8 +8,12 @@ session is active. These tests verify that delegation and the session guard.

8 """

10 import pytest

11 +import torch

12 +import torch.nn as nn

14 from vllm.config import VllmConfig, get_current_vllm_config

15 +from vllm.lora.layers import BaseLayerWithLoRA

16 +from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload

17 from vllm.v1.worker.gpu_worker import Worker

@@ -21,6 +25,7 @@ class _RecordingEngine:

25 self.started = False

26 self.finished = False

27 self.reset_count = 0

28 + self.supports_draft_weight_update = False

29 self.update_calls: list[dict] = []

30 self.seen_configs: list[VllmConfig] = []

@@ -48,16 +53,22 @@ class _RecordingEngine:

53 class _RecordingModelRunner:

54 def __init__(self) -> None:

55 self.seen_config: VllmConfig | None = None

56 + self.reset_lora_calls = 0

58 def reload_weights(self) -> None:

59 self.seen_config = get_current_vllm_config()

61 + def reset_lora_state(self) -> None:

62 + self.reset_lora_calls += 1

63 +

65 def _make_worker(engine: _RecordingEngine | None) -> Worker:

66 worker = object.__new__(Worker)

67 worker.vllm_config = VllmConfig()

68 worker.weight_transfer_engine = engine

69 worker._weight_update_active = False

70 + worker._weight_update_is_draft = False

71 + worker.model_runner = _RecordingModelRunner()

72 return worker

@@ -71,6 +82,22 @@ def test_reload_weights_sets_current_config():

82 assert model_runner.seen_config is worker.vllm_config

85 +def test_reload_parameter_lookup_preserves_lora_module_names():

86 + base_layer = nn.Module()

87 + qweight = nn.Parameter(torch.ones(1))

88 + base_layer.register_parameter("qweight", qweight)

89 + wrapper = BaseLayerWithLoRA()

90 + wrapper.base_layer = base_layer

91 + model = nn.Module()

92 + model.proj = wrapper

93 +

94 + named_parameters = dict(model.named_parameters())

95 + assert set(named_parameters) == {"proj.base_layer.qweight"}

96 + assert named_parameters["proj.base_layer.qweight"] is qweight

97 + assert model.get_parameter("proj.base_layer.qweight") is qweight

98 + assert _get_parameter_for_reload(model, "proj.qweight") is qweight

99 +

100 +

101 def test_start_update_finish_delegates_to_engine():

102 engine = _RecordingEngine()

103 worker = _make_worker(engine)

@@ -88,6 +115,19 @@ def test_start_update_finish_delegates_to_engine():

115 assert engine.reset_count == 1

116 assert worker._weight_update_active is False

117 assert engine.seen_configs == [worker.vllm_config] * 3

118 + assert worker.model_runner.reset_lora_calls == 1

119 +

120 +

121 +def test_finish_draft_session_keeps_lora_state():

122 + engine = _RecordingEngine()

123 + engine.supports_draft_weight_update = True

124 + worker = _make_worker(engine)

125 + worker._set_draft_weight_update_target = lambda: None

126 +

127 + Worker.start_draft_weight_update(worker)

128 + Worker.finish_weight_update(worker)

129 +

130 + assert worker.model_runner.reset_lora_calls == 0

131

132

133 def test_double_start_raises():

vllm/lora/layers/base.py

+13

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +from collections.abc import Iterable

5 from typing import TYPE_CHECKING, overload

7 import torch

@@ -14,6 +15,18 @@ if TYPE_CHECKING:

17 class BaseLayerWithLoRA(nn.Module):

18 + def load_weights(

19 + self, weights: Iterable[tuple[str, torch.Tensor]]

20 + ) -> Iterable[str]:

21 + """Load checkpoint weights into the wrapped base layer."""

22 + base_load_weights = getattr(self.base_layer, "load_weights", None)

23 + if callable(base_load_weights):

24 + return base_load_weights(weights)

25 +

26 + from vllm.model_executor.models.utils import AutoWeightsLoader

27 +

28 + return AutoWeightsLoader(self.base_layer).load_weights(weights)

29 +

30 @overload

31 def slice_lora_a(

32 self, lora_a: list[torch.Tensor | None]

vllm/lora/layers/logits_processor.py

+12

@@ -118,6 +118,18 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

118 else:

119 self.sharded_to_full_mapping_gpu = None

120

121 + def reset_sharded_to_full_mapping(self) -> None:

122 + """Restore the TP logits mapping after its GPU memory is reused."""

123 + mapping_gpu = self.sharded_to_full_mapping_gpu

124 + if mapping_gpu is not None:

125 + mapping_gpu.copy_(

126 + torch.tensor(

127 + self.sharded_to_full_mapping,

128 + device=mapping_gpu.device,

129 + dtype=mapping_gpu.dtype,

130 + )

131 + )

132 +

133 def reset_lora(self, index: int):

134 self.lora_a_stacked[index] = 0

135 self.lora_b_stacked[index] = 0

vllm/lora/model_manager.py

@@ -387,6 +387,8 @@ class LoRAModelManager:

387 self._registered_adapters.clear()

388 self.lora_index_to_id = [None] * self.lora_slots

389 self._active_adapters.clear()

390 + self._last_mapping = None

391 + self._last_slot_layout = None

392

393 def _create_lora_modules(self):

394 def _parent_module(module_name: str) -> str:

vllm/v1/worker/gpu_model_runner.py

+18 -3

@@ -56,7 +56,7 @@ from vllm.forward_context import (

56 set_forward_context,

57 )

58 from vllm.logger import init_logger

59 -from vllm.lora.layers import LoRAMapping, LoRAMappingType

59 +from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping, LoRAMappingType

60 from vllm.model_executor.layers.attention import Attention, MLAAttention

61 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

62 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager

@@ -250,6 +250,16 @@ if TYPE_CHECKING:

250

251 logger = init_logger(__name__)

252

253 +

254 +def _get_parameter_for_reload(model: nn.Module, name: str) -> nn.Parameter:

255 + """Resolve checkpoint names without changing the model's module tree."""

256 + module_name, _, parameter_name = name.rpartition(".")

257 + module = model.get_submodule(module_name)

258 + if isinstance(module, BaseLayerWithLoRA):

259 + module = module.base_layer

260 + return module.get_parameter(parameter_name)

261 +

262 +

263 AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata]

264 # list when ubatching is enabled

265 PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict

@@ -5535,7 +5545,10 @@ class GPUModelRunner(

5545 )

5546

5547 model = self.get_model()

5538 - weights_to_load = {name for name, _ in model.named_parameters()}

5548 + weights_to_load = {

5549 + name.replace(".base_layer.", ".") if self.lora_config else name

5550 + for name, _ in model.named_parameters()

5551 + }

5552 counter_before_reloading = time.perf_counter()

5553

5554 # load weights from disk if none are provided

@@ -5572,10 +5585,12 @@ class GPUModelRunner(

5585 )

5586 loaded_weights = set()

5587 for name, loaded_weight in weights_iterator:

5575 - param = model.get_parameter(name) # TODO: buffers?

5588 + param = _get_parameter_for_reload(model, name) # TODO: buffers?

5589 param.copy_(loaded_weight)

5590 loaded_weights.add(name)

5591

5592 + self.reset_lora_state()

5593 +

5594 # logging and validation

5595 counter_after_reloading = time.perf_counter()

5596 diff_seconds = counter_after_reloading - counter_before_reloading

vllm/v1/worker/gpu_worker.py

@@ -160,6 +160,7 @@ class Worker(WorkerBase):

160 # is available, since the engine needs a reference to the model.

161 self.weight_transfer_engine: WeightTransferEngine | None = None

162 self._weight_update_active = False

163 + self._weight_update_is_draft = False

164

165 # Torch/CUDA profiler. Enabled and configured through profiler_config.

166 # Profiler wrapper is created lazily in profile() when start is called,

@@ -1272,6 +1273,7 @@ class Worker(WorkerBase):

1273 self.weight_transfer_engine.reset_weight_update_target()

1274 raise

1275 self._weight_update_active = True

1276 + self._weight_update_is_draft = is_draft

1277

1278 def update_weights(self, update_info: dict) -> None:

1279 """

@@ -1316,6 +1318,10 @@ class Worker(WorkerBase):

1318 self.weight_transfer_engine.reset_weight_update_target()

1319 self._weight_update_active = False

1320

1321 + # Weight transfer bypasses GPUModelRunner.reload_weights().

1322 + if not self._weight_update_is_draft:

1323 + self.model_runner.reset_lora_state()

1324 +

1325 def shutdown(self) -> None:

1326 gc.unfreeze()

1327

vllm/v1/worker/lora_model_runner_mixin.py

+16

@@ -4,6 +4,7 @@

4 Define LoRA functionality mixin for model runners.

5 """

7 +from collections.abc import Callable

8 from contextlib import contextmanager

9 from typing import TypeAlias

@@ -28,6 +29,21 @@ logger = init_logger(__name__)

30 # Defined as a mixin for GPUModelRunner

31 class LoRAModelRunnerMixin:

32 + lora_config: LoRAConfig | None

33 + get_model: Callable[[], nn.Module]

34 +

35 + def reset_lora_state(self) -> None:

36 + """Invalidate LoRA state after base weights are replaced."""

37 + if not self.lora_config:

38 + return

39 +

40 + from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA

41 +

42 + self.lora_manager.remove_all_adapters()

43 + for module in self.get_model().modules():

44 + if isinstance(module, LogitsProcessorWithLoRA):

45 + module.reset_sharded_to_full_mapping()

46 +

47 def load_lora_model(

48 self,

49 model: nn.Module,