@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 639461eda4

[RL][Sleep] Retain frozen weights across level-2 sleep (#57891)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>

aoshen02 committed Sep 22, 2026 at 12:36 UTC 639461eda424debe7bbdf2b1fc577bf83e29c643

5 files changed +56

docs/features/sleep_mode.md

+15

@@ -74,6 +74,21 @@ llm.wake_up(tags=["weights"])

74 llm.wake_up(tags=["kv_cache"])

75 ```

77 +#### Retaining frozen weights during RLHF updates

78 +

79 +Set `sleep_preserve_parameter_names` (CLI: `--sleep-preserve-parameter-names`)

80 +to runtime parameter-name glob

81 +patterns for weights that stay frozen during training. Each pattern must match

82 +`model.named_parameters()`; checkpoint names and `requires_grad` are not used.

83 +

84 +Level-2 sleep backs up selected GPU parameters to pageable CPU memory; weights

85 +wake-up restores them in place and releases the backups. CPU parameters need no

86 +copy. Allow enough host memory for backups. Level-1 behavior is unchanged.

87 +

88 +The trainer must omit these parameters from updates, and reload/post-processing

89 +must preserve their values and storage (including avoiding replacement with meta

90 +tensors). This option does not filter updates and applies only to the target model.

91 +

92 #### Release only KV cache memory

94 `LLM.release_kv_cache_memory()` discards KV cache physical memory while keeping model weights resident. It requires a completed pause and all executor memory to be resident: full sleep, partial wake-up, and repeated release without restoring memory are rejected. Requests retained with `mode="keep"` are recomputed after wake-up.

tests/v1/worker/test_sleep_mode_backend.py

@@ -56,6 +56,7 @@ def test_worker_drives_communicator_suspension(monkeypatch, enable_nccl_comm_sus

57 worker = object.__new__(Worker)

58 worker._sleep_mode_backend = Backend()

59 + worker._sleep_saved_parameters = {}

60 worker._sleep_saved_buffers = {}

61 worker._sleep_saved_draft_buffers = {}

62 worker.vllm_config = SimpleNamespace(

vllm/config/model.py

@@ -340,6 +340,9 @@ class ModelConfig:

340 enable_sleep_mode: bool = False

341 """Enable sleep mode for the engine (only cuda and

342 hip platforms are supported)."""

343 + sleep_preserve_parameter_names: list[str] = field(default_factory=list)

344 + """Parameter-name globs to preserve across level-2 sleep.

345 + The sender must omit these parameters; loaders must preserve their storage."""

346 sleep_mode_backend: str = "cumem"

347 """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"``

348 (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible

vllm/engine/arg_utils.py

@@ -741,6 +741,9 @@ class EngineArgs:

741

742 generation_config: str = ModelConfig.generation_config

743 enable_sleep_mode: bool = ModelConfig.enable_sleep_mode

744 + sleep_preserve_parameter_names: list[str] = get_field(

745 + ModelConfig, "sleep_preserve_parameter_names"

746 + )

747 enable_cumem_allocator: bool = ModelConfig.enable_cumem_allocator

748 enable_nccl_comm_suspend: bool = ModelConfig.enable_nccl_comm_suspend

749 override_generation_config: dict[str, Any] = get_field(

@@ -975,6 +978,10 @@ class EngineArgs:

978 model_group.add_argument(

979 "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"]

980 )

981 + model_group.add_argument(

982 + "--sleep-preserve-parameter-names",

983 + **model_kwargs["sleep_preserve_parameter_names"],

984 + )

985 model_group.add_argument(

986 "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"]

987 )

@@ -1901,6 +1908,7 @@ class EngineArgs:

1908 generation_config=self.generation_config,

1909 override_generation_config=self.override_generation_config,

1910 enable_sleep_mode=self.enable_sleep_mode,

1911 + sleep_preserve_parameter_names=self.sleep_preserve_parameter_names,

1912 enable_cumem_allocator=self.enable_cumem_allocator,

1913 enable_nccl_comm_suspend=self.enable_nccl_comm_suspend,

1914 model_impl=self.model_impl,

vllm/v1/worker/gpu_worker.py

+29

@@ -8,6 +8,7 @@ import time

8 from collections.abc import Callable

9 from contextlib import AbstractContextManager, contextmanager, nullcontext

10 from datetime import timedelta

11 +from fnmatch import filter as fnmatch_filter

12 from types import NoneType

13 from typing import TYPE_CHECKING, Any, cast

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

209 self.worker_sentinel = WorkerSentinel(worker=self)

210 # Buffers saved before sleep

211 self._sleep_saved_buffers: dict[str, torch.Tensor] = {}

212 + self._sleep_saved_parameters: dict[str, torch.Tensor] = {}

213 self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {}

214

215 # Weight transfer engine is created in `load_model` once the model

@@ -243,6 +245,30 @@ class Worker(WorkerBase):

245 )

246 return self._sleep_mode_backend

247

248 + def _save_sleep_parameters(self, model: nn.Module) -> None:

249 + patterns = self.model_config.sleep_preserve_parameter_names

250 + if not patterns:

251 + self._sleep_saved_parameters = {}

252 + return

253 + parameters = dict(model.named_parameters())

254 + names: set[str] = set()

255 + for pattern in patterns:

256 + matches = fnmatch_filter(parameters, pattern)

257 + if not matches:

258 + raise ValueError(f"No parameter matches sleep retention: {pattern}")

259 + names.update(matches)

260 + self._sleep_saved_parameters = {

261 + name: param.detach().to("cpu")

262 + for name, param in parameters.items()

263 + if name in names and not param.is_cpu

264 + }

265 +

266 + @torch.no_grad()

267 + def _restore_sleep_parameters(self, model: nn.Module) -> None:

268 + for name, value in self._sleep_saved_parameters.items():

269 + model.get_parameter(name).copy_(value)

270 + self._sleep_saved_parameters.clear()

271 +

272 def sleep(self, level: int = 1) -> None:

273 torch.accelerator.synchronize()

274 free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]

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

276 # Save the buffers before level 2 sleep

277 if level == 2:

278 model = self.model_runner.model

279 + self._save_sleep_parameters(model)

280 self._sleep_saved_buffers = {

281 name: buffer.cpu().clone() for name, buffer in model.named_buffers()

282 }

@@ -287,6 +314,8 @@ class Worker(WorkerBase):

314

315 # Restore the buffers after level 2 sleep

316 wake_weights = tags is None or "weights" in tags

317 + if wake_weights and self._sleep_saved_parameters:

318 + self._restore_sleep_parameters(self.model_runner.model)

319 if wake_weights and len(self._sleep_saved_buffers):

320 model = self.model_runner.model

321 for name, buffer in model.named_buffers():