From 91563b1e1d4bc454007dc5918fab8aad583a7227 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 07:12:58 +0000 Subject: [PATCH 1/9] [RL][Sleep] Retain frozen weights across level-2 sleep Co-authored-by: Codex Signed-off-by: aoshen02 --- docs/features/sleep_mode.md | 20 +++ .../model_loader/test_reload.py | 70 +++++++++++ tests/v1/worker/test_sleep_mode_backend.py | 1 + vllm/config/weight_transfer.py | 8 ++ .../model_loader/reload/frozen.py | 116 ++++++++++++++++++ .../model_loader/reload/layerwise.py | 4 +- vllm/v1/worker/gpu_worker.py | 15 +++ 7 files changed, 233 insertions(+), 1 deletion(-) create mode 100644 vllm/model_executor/model_loader/reload/frozen.py diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 2616f374ea5a..86710dad1824 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -74,6 +74,26 @@ llm.wake_up(tags=["weights"]) llm.wake_up(tags=["kv_cache"]) ``` +#### Retaining frozen weights during RLHF updates + +Set `weight_transfer_config.frozen_weight_modules` to runtime module-name glob +patterns for submodules whose parameters never change during training. These +are vLLM module names, not checkpoint keys; do not infer them from +`requires_grad=False`, which is also used for ordinary inference weights. + +With level 2 sleep, their GPU parameters are copied to pageable CPU memory once. +`wake_up(tags=["weights"])` restores them in place and keeps the CPU copies for +later cycles. Already CPU-resident parameters are reused without another copy. +Budget host memory for these retained local shards alongside trainer offload. +Level 1 retains its existing full-weight backup behavior. + +The sender must omit the corresponding checkpoint tensors. Layerwise reload +leaves the selected modules intact, and their parameter loaders reject incoming +updates. Use the model's checkpoint-name mapping and validate the sender/receiver +exclusion sets before transfer. Restart the instance to change frozen weights; +direct writes that bypass weight loaders are unsupported. Selection applies to +the target model, not a separately instantiated draft model. + #### Release only KV cache memory `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. diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 155957160fb6..0a211185a317 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -20,6 +20,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import QKVParallelLinear from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.model_loader.reload.frozen import FrozenWeights from vllm.model_executor.model_loader.reload.layerwise import ( finalize_layerwise_reload, initialize_layerwise_reload, @@ -69,6 +70,75 @@ def __init__(self): ) +def test_frozen_module_stays_materialized_during_reload(): + """Omitting a frozen table must not move it to meta or accept updates.""" + model = torch.nn.ModuleDict( + {"table": torch.nn.Embedding(4, 2), "projection": torch.nn.Linear(2, 2)} + ) + record_metadata_for_reloading(model) + frozen = FrozenWeights(model, ["table"]) + table = model.table.weight + expected = table.detach().clone() + frozen.save() + assert frozen.backups == {} # Host-resident weights need no second CPU copy. + initialize_layerwise_reload(model) + assert model.table.weight is table + torch.testing.assert_close(table, expected) + assert model.projection.weight.is_meta + with pytest.raises(ValueError, match="Cannot reload a frozen weight"): + table.weight_loader(table, torch.zeros_like(table)) + + +def test_frozen_module_rejects_alias_to_mutable_weight(): + """An update through a tied mutable name would invalidate the frozen copy.""" + model = torch.nn.ModuleDict( + {"table": torch.nn.Embedding(4, 2), "projection": torch.nn.Linear(2, 4)} + ) + model.projection.weight = model.table.weight + with pytest.raises(ValueError, match="aliases mutable parameter"): + FrozenWeights(model, ["table"]) + + +def test_frozen_module_rejects_misspelled_selection(): + model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) + with pytest.raises(ValueError, match="matched nothing"): + FrozenWeights(model, ["missing_table"]) + + +def test_frozen_module_rejects_mutable_buffer_alias(): + model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) + model.register_buffer("alias", model.table.weight.detach()) + with pytest.raises(ValueError, match="aliases mutable buffer"): + FrozenWeights(model, ["table"]) + + +def test_frozen_module_rejects_storage_rebinding(): + model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) + frozen = FrozenWeights(model, ["table"]) + model.table.weight.data = model.table.weight.detach().clone() + with pytest.raises(RuntimeError, match="storage was replaced"): + frozen.restore() + + +def test_frozen_checkpoint_names_use_model_mapper(): + from vllm.model_executor.models.utils import WeightsMapper + + model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) + model.hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"checkpoint_table.": "table.", "ignored.": None} + ) + frozen = FrozenWeights(model, ["table"]) + assert frozen.checkpoint_names( + ["checkpoint_table.weight", "projection.weight", "ignored.weight"] + ) == {"checkpoint_table.weight"} + + +def test_frozen_module_rejects_whole_model_pattern(): + model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) + with pytest.raises(ValueError, match="entire model"): + FrozenWeights(model, ["*"]) + + class _ParentAliasedChildBufferLayer(torch.nn.Module): def __init__(self): super().__init__() diff --git a/tests/v1/worker/test_sleep_mode_backend.py b/tests/v1/worker/test_sleep_mode_backend.py index e2514f84e162..773357825126 100644 --- a/tests/v1/worker/test_sleep_mode_backend.py +++ b/tests/v1/worker/test_sleep_mode_backend.py @@ -56,6 +56,7 @@ def resume(self, tags: list[str] | None = None) -> None: worker = object.__new__(Worker) worker._sleep_mode_backend = Backend() + worker._frozen_weights = None worker._sleep_saved_buffers = {} worker._sleep_saved_draft_buffers = {} worker.vllm_config = SimpleNamespace( diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index 3e849753c2e6..ed8cf7b8ac03 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Literal +from pydantic import Field + from vllm.config.utils import config @@ -13,3 +15,9 @@ class WeightTransferConfig: """The backend to use for weight transfer. Validated against the `WeightTransferEngineFactory` registry at engine creation time. """ + + frozen_weight_modules: list[str] = Field(default_factory=list) + """Runtime module-name glob patterns excluded from weight reloads. Their + parameters are immutable after initial loading and must be omitted by the + sender. L2 sleep keeps a persistent CPU copy; CPU-resident weights are reused. + """ diff --git a/vllm/model_executor/model_loader/reload/frozen.py b/vllm/model_executor/model_loader/reload/frozen.py new file mode 100644 index 000000000000..cbffbf01402c --- /dev/null +++ b/vllm/model_executor/model_loader/reload/frozen.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable +from fnmatch import fnmatchcase + +import torch + + +def _reject_frozen_update(*args, **kwargs): + raise ValueError("Cannot reload a frozen weight; restart to change frozen weights") + + +class FrozenWeights: + """Retain explicitly frozen modules across L2 sleep and partial reloads. + + Patterns name runtime modules, not checkpoint tensors. Configure after initial + loading. Frozen modules must keep their parameter objects and storage intact. + """ + + def __init__(self, model: torch.nn.Module, patterns: list[str]): + modules = dict(model.named_modules()) + selected: set[str] = set() + for pattern in patterns: + matches = {name for name in modules if fnmatchcase(name, pattern)} + if not matches: + raise ValueError( + f"Frozen weight module pattern matched nothing: {pattern}" + ) + if "" in matches: + raise ValueError("Select frozen submodules, not the entire model") + selected.update(matches) + names = { + name + for name in modules + if any( + not root or name == root or name.startswith(root + ".") + for root in selected + ) + } + self.model = model + self.module_names = names + self.parameters = { + name: param + for name, param in model.named_parameters(remove_duplicate=False) + if name.rpartition(".")[0] in names + } + if patterns and not self.parameters: + raise ValueError("Frozen weight modules contain no parameters") + frozen_storage = { + (param.device, param.untyped_storage().data_ptr()) + for param in self.parameters.values() + if param.numel() + } + for name, param in model.named_parameters(remove_duplicate=False): + if ( + name not in self.parameters + and param.numel() + and (param.device, param.untyped_storage().data_ptr()) in frozen_storage + ): + raise ValueError(f"Frozen storage aliases mutable parameter: {name}") + for name, buffer in model.named_buffers(remove_duplicate=False): + if ( + name.rpartition(".")[0] not in names + and buffer.numel() + and (buffer.device, buffer.untyped_storage().data_ptr()) + in frozen_storage + ): + raise ValueError(f"Frozen storage aliases mutable buffer: {name}") + for name in names: + modules[name]._vllm_frozen_weights = True + for param in self.parameters.values(): + param.weight_loader = _reject_frozen_update + self.backups: dict[str, torch.Tensor] = {} + self.layouts = { + name: (param.data_ptr(), param.shape, param.stride(), param.dtype) + for name, param in self.parameters.items() + } + + def checkpoint_names(self, names: Iterable[str]) -> set[str]: + """Resolve a sender's checkpoint names using the model's own mapper.""" + mapper = getattr(self.model, "hf_to_vllm_mapper", None) + result = set() + for name in names: + mapped = mapper.apply_list([name]) if mapper is not None else [name] + if mapped and all( + any( + not module or key.startswith(module + ".") + for module in self.module_names + ) + for key in mapped + ): + result.add(name) + return result + + def _validate_storage(self) -> None: + for name, param in self.parameters.items(): + if self.model.get_parameter(name) is not param or self.layouts[name] != ( + param.data_ptr(), + param.shape, + param.stride(), + param.dtype, + ): + raise RuntimeError(f"Frozen parameter storage was replaced: {name}") + + @torch.no_grad() + def save(self) -> None: + self._validate_storage() + for name, param in self.parameters.items(): + if param.device.type != "cpu" and name not in self.backups: + self.backups[name] = param.to(device="cpu", copy=True) + + @torch.no_grad() + def restore(self) -> None: + self._validate_storage() + for name, backup in self.backups.items(): + self.parameters[name].copy_(backup) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 062bfa60906b..1f3db1d987ff 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -101,7 +101,7 @@ def initialize_layerwise_reload(model: torch.nn.Module): info = get_layerwise_info(layer) # Skip if the layer has already been initialized - if info.can_load(): + if getattr(layer, "_vllm_frozen_weights", False) or info.can_load(): continue # Save current tensors for later copying @@ -241,6 +241,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = [] for layer in model.modules(): + if getattr(layer, "_vllm_frozen_weights", False): + continue info = get_layerwise_info(layer) if not info.can_load(): info.reset() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index d914bd2deab3..17e05a8ebde5 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -101,6 +101,9 @@ from .gpu.warmup import warmup_kernels from .utils import request_memory +if TYPE_CHECKING: + from vllm.model_executor.model_loader.reload.frozen import FrozenWeights + logger = init_logger(__name__) @@ -208,6 +211,7 @@ def __init__( self.worker_sentinel = WorkerSentinel(worker=self) # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} + self._frozen_weights: FrozenWeights | None = None self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {} # Weight transfer engine is created in `load_model` once the model @@ -249,6 +253,8 @@ def sleep(self, level: int = 1) -> None: # Save the buffers before level 2 sleep if level == 2: + if self._frozen_weights is not None: + self._frozen_weights.save() model = self.model_runner.model self._sleep_saved_buffers = { name: buffer.cpu().clone() for name, buffer in model.named_buffers() @@ -287,6 +293,8 @@ def wake_up(self, tags: list[str] | None = None) -> None: # Restore the buffers after level 2 sleep wake_weights = tags is None or "weights" in tags + if wake_weights and self._frozen_weights is not None: + self._frozen_weights.restore() if wake_weights and len(self._sleep_saved_buffers): model = self.model_runner.model for name, buffer in model.named_buffers(): @@ -514,6 +522,13 @@ def load_model(self, *, load_dummy_weights: bool = False) -> None: get_ec_transfer().start_worker_services() if self.vllm_config.weight_transfer_config is not None: + if self.vllm_config.weight_transfer_config.frozen_weight_modules: + from vllm.model_executor.model_loader.reload.frozen import FrozenWeights + + self._frozen_weights = FrozenWeights( + self.model_runner.get_model(), + self.vllm_config.weight_transfer_config.frozen_weight_modules, + ) self.weight_transfer_engine = WeightTransferEngineFactory.create_engine( self.vllm_config.weight_transfer_config, self.vllm_config, From d76d7f38fc9ccc495632cdebfc8fc6bdb7cc9b0f Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 07:19:41 +0000 Subject: [PATCH 2/9] Scope frozen weight retention to sleep and wake Co-authored-by: Codex Signed-off-by: aoshen02 --- docs/features/sleep_mode.md | 12 ++-- .../model_loader/test_reload.py | 70 ------------------- vllm/config/weight_transfer.py | 7 +- .../model_loader/reload/layerwise.py | 4 +- .../frozen.py => v1/worker/frozen_weights.py} | 33 ++------- vllm/v1/worker/gpu_worker.py | 4 +- 6 files changed, 19 insertions(+), 111 deletions(-) rename vllm/{model_executor/model_loader/reload/frozen.py => v1/worker/frozen_weights.py} (76%) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 86710dad1824..8dc9f6ad2b7b 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -87,12 +87,12 @@ later cycles. Already CPU-resident parameters are reused without another copy. Budget host memory for these retained local shards alongside trainer offload. Level 1 retains its existing full-weight backup behavior. -The sender must omit the corresponding checkpoint tensors. Layerwise reload -leaves the selected modules intact, and their parameter loaders reject incoming -updates. Use the model's checkpoint-name mapping and validate the sender/receiver -exclusion sets before transfer. Restart the instance to change frozen weights; -direct writes that bypass weight loaders are unsupported. Selection applies to -the target model, not a separately instantiated draft model. +This option only controls sleep/wake; it does not modify weight loaders or +filter incoming updates. The trainer (for example, VERL) must omit these +parameters and leave their values and storage unchanged during reload and +post-processing. Restart the instance to change frozen weights. Reload paths +that replace selected parameters with meta tensors are not supported by this +option alone. Selection applies to the target, not a separate draft model. #### Release only KV cache memory diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 0a211185a317..155957160fb6 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -20,7 +20,6 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import QKVParallelLinear from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase -from vllm.model_executor.model_loader.reload.frozen import FrozenWeights from vllm.model_executor.model_loader.reload.layerwise import ( finalize_layerwise_reload, initialize_layerwise_reload, @@ -70,75 +69,6 @@ def __init__(self): ) -def test_frozen_module_stays_materialized_during_reload(): - """Omitting a frozen table must not move it to meta or accept updates.""" - model = torch.nn.ModuleDict( - {"table": torch.nn.Embedding(4, 2), "projection": torch.nn.Linear(2, 2)} - ) - record_metadata_for_reloading(model) - frozen = FrozenWeights(model, ["table"]) - table = model.table.weight - expected = table.detach().clone() - frozen.save() - assert frozen.backups == {} # Host-resident weights need no second CPU copy. - initialize_layerwise_reload(model) - assert model.table.weight is table - torch.testing.assert_close(table, expected) - assert model.projection.weight.is_meta - with pytest.raises(ValueError, match="Cannot reload a frozen weight"): - table.weight_loader(table, torch.zeros_like(table)) - - -def test_frozen_module_rejects_alias_to_mutable_weight(): - """An update through a tied mutable name would invalidate the frozen copy.""" - model = torch.nn.ModuleDict( - {"table": torch.nn.Embedding(4, 2), "projection": torch.nn.Linear(2, 4)} - ) - model.projection.weight = model.table.weight - with pytest.raises(ValueError, match="aliases mutable parameter"): - FrozenWeights(model, ["table"]) - - -def test_frozen_module_rejects_misspelled_selection(): - model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) - with pytest.raises(ValueError, match="matched nothing"): - FrozenWeights(model, ["missing_table"]) - - -def test_frozen_module_rejects_mutable_buffer_alias(): - model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) - model.register_buffer("alias", model.table.weight.detach()) - with pytest.raises(ValueError, match="aliases mutable buffer"): - FrozenWeights(model, ["table"]) - - -def test_frozen_module_rejects_storage_rebinding(): - model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) - frozen = FrozenWeights(model, ["table"]) - model.table.weight.data = model.table.weight.detach().clone() - with pytest.raises(RuntimeError, match="storage was replaced"): - frozen.restore() - - -def test_frozen_checkpoint_names_use_model_mapper(): - from vllm.model_executor.models.utils import WeightsMapper - - model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) - model.hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={"checkpoint_table.": "table.", "ignored.": None} - ) - frozen = FrozenWeights(model, ["table"]) - assert frozen.checkpoint_names( - ["checkpoint_table.weight", "projection.weight", "ignored.weight"] - ) == {"checkpoint_table.weight"} - - -def test_frozen_module_rejects_whole_model_pattern(): - model = torch.nn.ModuleDict({"table": torch.nn.Embedding(4, 2)}) - with pytest.raises(ValueError, match="entire model"): - FrozenWeights(model, ["*"]) - - class _ParentAliasedChildBufferLayer(torch.nn.Module): def __init__(self): super().__init__() diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index ed8cf7b8ac03..54347f888f5d 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -17,7 +17,8 @@ class WeightTransferConfig: """ frozen_weight_modules: list[str] = Field(default_factory=list) - """Runtime module-name glob patterns excluded from weight reloads. Their - parameters are immutable after initial loading and must be omitted by the - sender. L2 sleep keeps a persistent CPU copy; CPU-resident weights are reused. + """Runtime module-name glob patterns retained across level-2 sleep. Their + parameters must remain immutable and be omitted by the sender. This only + controls sleep/wake, not weight loading. GPU parameters keep a CPU backup; + CPU-resident parameters are reused. """ diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 1f3db1d987ff..062bfa60906b 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -101,7 +101,7 @@ def initialize_layerwise_reload(model: torch.nn.Module): info = get_layerwise_info(layer) # Skip if the layer has already been initialized - if getattr(layer, "_vllm_frozen_weights", False) or info.can_load(): + if info.can_load(): continue # Save current tensors for later copying @@ -241,8 +241,6 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = [] for layer in model.modules(): - if getattr(layer, "_vllm_frozen_weights", False): - continue info = get_layerwise_info(layer) if not info.can_load(): info.reset() diff --git a/vllm/model_executor/model_loader/reload/frozen.py b/vllm/v1/worker/frozen_weights.py similarity index 76% rename from vllm/model_executor/model_loader/reload/frozen.py rename to vllm/v1/worker/frozen_weights.py index cbffbf01402c..44781fba63cc 100644 --- a/vllm/model_executor/model_loader/reload/frozen.py +++ b/vllm/v1/worker/frozen_weights.py @@ -1,17 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable from fnmatch import fnmatchcase import torch -def _reject_frozen_update(*args, **kwargs): - raise ValueError("Cannot reload a frozen weight; restart to change frozen weights") - - class FrozenWeights: - """Retain explicitly frozen modules across L2 sleep and partial reloads. + """Retain explicitly frozen parameters across level-2 sleep. Patterns name runtime modules, not checkpoint tensors. Configure after initial loading. Frozen modules must keep their parameter objects and storage intact. @@ -38,7 +33,6 @@ def __init__(self, model: torch.nn.Module, patterns: list[str]): ) } self.model = model - self.module_names = names self.parameters = { name: param for name, param in model.named_parameters(remove_duplicate=False) @@ -66,32 +60,13 @@ def __init__(self, model: torch.nn.Module, patterns: list[str]): in frozen_storage ): raise ValueError(f"Frozen storage aliases mutable buffer: {name}") - for name in names: - modules[name]._vllm_frozen_weights = True - for param in self.parameters.values(): - param.weight_loader = _reject_frozen_update self.backups: dict[str, torch.Tensor] = {} + self._needs_restore = False self.layouts = { name: (param.data_ptr(), param.shape, param.stride(), param.dtype) for name, param in self.parameters.items() } - def checkpoint_names(self, names: Iterable[str]) -> set[str]: - """Resolve a sender's checkpoint names using the model's own mapper.""" - mapper = getattr(self.model, "hf_to_vllm_mapper", None) - result = set() - for name in names: - mapped = mapper.apply_list([name]) if mapper is not None else [name] - if mapped and all( - any( - not module or key.startswith(module + ".") - for module in self.module_names - ) - for key in mapped - ): - result.add(name) - return result - def _validate_storage(self) -> None: for name, param in self.parameters.items(): if self.model.get_parameter(name) is not param or self.layouts[name] != ( @@ -108,9 +83,13 @@ def save(self) -> None: for name, param in self.parameters.items(): if param.device.type != "cpu" and name not in self.backups: self.backups[name] = param.to(device="cpu", copy=True) + self._needs_restore = True @torch.no_grad() def restore(self) -> None: + if not self._needs_restore: + return self._validate_storage() for name, backup in self.backups.items(): self.parameters[name].copy_(backup) + self._needs_restore = False diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 17e05a8ebde5..e56b320829ca 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -102,7 +102,7 @@ from .utils import request_memory if TYPE_CHECKING: - from vllm.model_executor.model_loader.reload.frozen import FrozenWeights + from vllm.v1.worker.frozen_weights import FrozenWeights logger = init_logger(__name__) @@ -523,7 +523,7 @@ def load_model(self, *, load_dummy_weights: bool = False) -> None: if self.vllm_config.weight_transfer_config is not None: if self.vllm_config.weight_transfer_config.frozen_weight_modules: - from vllm.model_executor.model_loader.reload.frozen import FrozenWeights + from vllm.v1.worker.frozen_weights import FrozenWeights self._frozen_weights = FrozenWeights( self.model_runner.get_model(), From 40a263bd36f831c7dd63f2d08608bfb30b1accd2 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 07:24:13 +0000 Subject: [PATCH 3/9] Simplify retained parameters to sleep and wake snapshots Co-authored-by: Codex Signed-off-by: aoshen02 --- docs/features/sleep_mode.md | 9 +- tests/v1/worker/test_sleep_mode_backend.py | 2 +- vllm/config/weight_transfer.py | 4 +- vllm/v1/worker/frozen_weights.py | 95 ---------------------- vllm/v1/worker/gpu_worker.py | 40 +++++---- 5 files changed, 33 insertions(+), 117 deletions(-) delete mode 100644 vllm/v1/worker/frozen_weights.py diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 8dc9f6ad2b7b..9bccdbd869ef 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -81,10 +81,11 @@ patterns for submodules whose parameters never change during training. These are vLLM module names, not checkpoint keys; do not infer them from `requires_grad=False`, which is also used for ordinary inference weights. -With level 2 sleep, their GPU parameters are copied to pageable CPU memory once. -`wake_up(tags=["weights"])` restores them in place and keeps the CPU copies for -later cycles. Already CPU-resident parameters are reused without another copy. -Budget host memory for these retained local shards alongside trainer offload. +With level 2 sleep, their GPU parameters are copied to pageable CPU memory. +`wake_up(tags=["weights"])` restores them in place and releases the CPU copies. +Each sleep/wake cycle saves and restores the current values, just like buffers. +Already CPU-resident parameters need no copy. Budget host memory for these +backups alongside trainer offload. Level 1 retains its existing full-weight backup behavior. This option only controls sleep/wake; it does not modify weight loaders or diff --git a/tests/v1/worker/test_sleep_mode_backend.py b/tests/v1/worker/test_sleep_mode_backend.py index 773357825126..686331f2d700 100644 --- a/tests/v1/worker/test_sleep_mode_backend.py +++ b/tests/v1/worker/test_sleep_mode_backend.py @@ -56,7 +56,7 @@ def resume(self, tags: list[str] | None = None) -> None: worker = object.__new__(Worker) worker._sleep_mode_backend = Backend() - worker._frozen_weights = None + worker._sleep_saved_parameters = {} worker._sleep_saved_buffers = {} worker._sleep_saved_draft_buffers = {} worker.vllm_config = SimpleNamespace( diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index 54347f888f5d..6a450f5eadae 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -19,6 +19,6 @@ class WeightTransferConfig: frozen_weight_modules: list[str] = Field(default_factory=list) """Runtime module-name glob patterns retained across level-2 sleep. Their parameters must remain immutable and be omitted by the sender. This only - controls sleep/wake, not weight loading. GPU parameters keep a CPU backup; - CPU-resident parameters are reused. + controls sleep/wake, not weight loading. GPU parameters are backed up to CPU + at sleep and restored at wake; CPU-resident parameters need no copy. """ diff --git a/vllm/v1/worker/frozen_weights.py b/vllm/v1/worker/frozen_weights.py deleted file mode 100644 index 44781fba63cc..000000000000 --- a/vllm/v1/worker/frozen_weights.py +++ /dev/null @@ -1,95 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from fnmatch import fnmatchcase - -import torch - - -class FrozenWeights: - """Retain explicitly frozen parameters across level-2 sleep. - - Patterns name runtime modules, not checkpoint tensors. Configure after initial - loading. Frozen modules must keep their parameter objects and storage intact. - """ - - def __init__(self, model: torch.nn.Module, patterns: list[str]): - modules = dict(model.named_modules()) - selected: set[str] = set() - for pattern in patterns: - matches = {name for name in modules if fnmatchcase(name, pattern)} - if not matches: - raise ValueError( - f"Frozen weight module pattern matched nothing: {pattern}" - ) - if "" in matches: - raise ValueError("Select frozen submodules, not the entire model") - selected.update(matches) - names = { - name - for name in modules - if any( - not root or name == root or name.startswith(root + ".") - for root in selected - ) - } - self.model = model - self.parameters = { - name: param - for name, param in model.named_parameters(remove_duplicate=False) - if name.rpartition(".")[0] in names - } - if patterns and not self.parameters: - raise ValueError("Frozen weight modules contain no parameters") - frozen_storage = { - (param.device, param.untyped_storage().data_ptr()) - for param in self.parameters.values() - if param.numel() - } - for name, param in model.named_parameters(remove_duplicate=False): - if ( - name not in self.parameters - and param.numel() - and (param.device, param.untyped_storage().data_ptr()) in frozen_storage - ): - raise ValueError(f"Frozen storage aliases mutable parameter: {name}") - for name, buffer in model.named_buffers(remove_duplicate=False): - if ( - name.rpartition(".")[0] not in names - and buffer.numel() - and (buffer.device, buffer.untyped_storage().data_ptr()) - in frozen_storage - ): - raise ValueError(f"Frozen storage aliases mutable buffer: {name}") - self.backups: dict[str, torch.Tensor] = {} - self._needs_restore = False - self.layouts = { - name: (param.data_ptr(), param.shape, param.stride(), param.dtype) - for name, param in self.parameters.items() - } - - def _validate_storage(self) -> None: - for name, param in self.parameters.items(): - if self.model.get_parameter(name) is not param or self.layouts[name] != ( - param.data_ptr(), - param.shape, - param.stride(), - param.dtype, - ): - raise RuntimeError(f"Frozen parameter storage was replaced: {name}") - - @torch.no_grad() - def save(self) -> None: - self._validate_storage() - for name, param in self.parameters.items(): - if param.device.type != "cpu" and name not in self.backups: - self.backups[name] = param.to(device="cpu", copy=True) - self._needs_restore = True - - @torch.no_grad() - def restore(self) -> None: - if not self._needs_restore: - return - self._validate_storage() - for name, backup in self.backups.items(): - self.parameters[name].copy_(backup) - self._needs_restore = False diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index e56b320829ca..0a7e4f881f58 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -8,6 +8,7 @@ from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta +from fnmatch import fnmatchcase from types import NoneType from typing import TYPE_CHECKING, Any, cast @@ -101,9 +102,6 @@ from .gpu.warmup import warmup_kernels from .utils import request_memory -if TYPE_CHECKING: - from vllm.v1.worker.frozen_weights import FrozenWeights - logger = init_logger(__name__) @@ -211,7 +209,7 @@ def __init__( self.worker_sentinel = WorkerSentinel(worker=self) # Buffers saved before sleep self._sleep_saved_buffers: dict[str, torch.Tensor] = {} - self._frozen_weights: FrozenWeights | None = None + self._sleep_saved_parameters: dict[str, torch.Tensor] = {} self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {} # Weight transfer engine is created in `load_model` once the model @@ -253,9 +251,24 @@ def sleep(self, level: int = 1) -> None: # Save the buffers before level 2 sleep if level == 2: - if self._frozen_weights is not None: - self._frozen_weights.save() model = self.model_runner.model + transfer_config = self.vllm_config.weight_transfer_config + patterns = transfer_config.frozen_weight_modules if transfer_config else [] + modules = dict(model.named_modules()) + retained = {} + for pattern in patterns: + matches = [name for name in modules if fnmatchcase(name, pattern)] + if not matches: + raise ValueError(f"No module matches sleep retention: {pattern}") + for name in matches: + for key, param in modules[name].named_parameters(): + if param.device.type != "cpu": + retained[f"{name}.{key}" if name else key] = param + for name, param in retained.items(): + if name not in self._sleep_saved_parameters: + self._sleep_saved_parameters[name] = param.detach().to( + device="cpu", copy=True + ) self._sleep_saved_buffers = { name: buffer.cpu().clone() for name, buffer in model.named_buffers() } @@ -293,8 +306,12 @@ def wake_up(self, tags: list[str] | None = None) -> None: # Restore the buffers after level 2 sleep wake_weights = tags is None or "weights" in tags - if wake_weights and self._frozen_weights is not None: - self._frozen_weights.restore() + if wake_weights and self._sleep_saved_parameters: + model = self.model_runner.model + with torch.no_grad(): + for name, saved in self._sleep_saved_parameters.items(): + model.get_parameter(name).copy_(saved) + self._sleep_saved_parameters.clear() if wake_weights and len(self._sleep_saved_buffers): model = self.model_runner.model for name, buffer in model.named_buffers(): @@ -522,13 +539,6 @@ def load_model(self, *, load_dummy_weights: bool = False) -> None: get_ec_transfer().start_worker_services() if self.vllm_config.weight_transfer_config is not None: - if self.vllm_config.weight_transfer_config.frozen_weight_modules: - from vllm.v1.worker.frozen_weights import FrozenWeights - - self._frozen_weights = FrozenWeights( - self.model_runner.get_model(), - self.vllm_config.weight_transfer_config.frozen_weight_modules, - ) self.weight_transfer_engine = WeightTransferEngineFactory.create_engine( self.vllm_config.weight_transfer_config, self.vllm_config, From d219c7164fb0155c2ee2bd3e480839600a82d210 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 07:27:33 +0000 Subject: [PATCH 4/9] Match retained sleep weights by parameter name Co-authored-by: Codex Signed-off-by: aoshen02 --- docs/features/sleep_mode.md | 7 ++++--- vllm/config/weight_transfer.py | 6 +++--- vllm/v1/worker/gpu_worker.py | 27 ++++++++++++--------------- 3 files changed, 19 insertions(+), 21 deletions(-) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 9bccdbd869ef..5a4bbfae56b7 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -76,10 +76,11 @@ llm.wake_up(tags=["kv_cache"]) #### Retaining frozen weights during RLHF updates -Set `weight_transfer_config.frozen_weight_modules` to runtime module-name glob -patterns for submodules whose parameters never change during training. These -are vLLM module names, not checkpoint keys; do not infer them from +Set `weight_transfer_config.frozen_weight_names` to runtime parameter-name glob +patterns, such as `*.engram.*`, for weights that never change during training. +These match `model.named_parameters()`, not checkpoint keys; do not infer them from `requires_grad=False`, which is also used for ordinary inference weights. +Each pattern must match at least one parameter, including CPU-resident parameters. With level 2 sleep, their GPU parameters are copied to pageable CPU memory. `wake_up(tags=["weights"])` restores them in place and releases the CPU copies. diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index 6a450f5eadae..71005597a40e 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -16,9 +16,9 @@ class WeightTransferConfig: `WeightTransferEngineFactory` registry at engine creation time. """ - frozen_weight_modules: list[str] = Field(default_factory=list) - """Runtime module-name glob patterns retained across level-2 sleep. Their - parameters must remain immutable and be omitted by the sender. This only + frozen_weight_names: list[str] = Field(default_factory=list) + """Runtime parameter-name glob patterns retained across level-2 sleep. + These parameters must remain immutable and be omitted by the sender. This only controls sleep/wake, not weight loading. GPU parameters are backed up to CPU at sleep and restored at wake; CPU-resident parameters need no copy. """ diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 0a7e4f881f58..816e63fc6ee6 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -8,7 +8,7 @@ from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta -from fnmatch import fnmatchcase +from fnmatch import filter as fnmatch_filter from types import NoneType from typing import TYPE_CHECKING, Any, cast @@ -253,22 +253,19 @@ def sleep(self, level: int = 1) -> None: if level == 2: model = self.model_runner.model transfer_config = self.vllm_config.weight_transfer_config - patterns = transfer_config.frozen_weight_modules if transfer_config else [] - modules = dict(model.named_modules()) - retained = {} + patterns = transfer_config.frozen_weight_names if transfer_config else [] + parameters = dict(model.named_parameters()) + retained_names: set[str] = set() for pattern in patterns: - matches = [name for name in modules if fnmatchcase(name, pattern)] + matches = fnmatch_filter(parameters, pattern) if not matches: - raise ValueError(f"No module matches sleep retention: {pattern}") - for name in matches: - for key, param in modules[name].named_parameters(): - if param.device.type != "cpu": - retained[f"{name}.{key}" if name else key] = param - for name, param in retained.items(): - if name not in self._sleep_saved_parameters: - self._sleep_saved_parameters[name] = param.detach().to( - device="cpu", copy=True - ) + raise ValueError(f"No parameter matches sleep retention: {pattern}") + retained_names.update(matches) + self._sleep_saved_parameters = { + name: param.detach().to("cpu") + for name, param in parameters.items() + if name in retained_names and not param.is_cpu + } self._sleep_saved_buffers = { name: buffer.cpu().clone() for name, buffer in model.named_buffers() } From 4f5fc441e0ec76934889424bfef14b606926fe41 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 14:08:44 +0000 Subject: [PATCH 5/9] Extract sleep parameter backup and restore helpers Co-authored-by: Codex Signed-off-by: aoshen02 --- vllm/v1/worker/gpu_worker.py | 23 +++++------------------ vllm/v1/worker/sleep.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 18 deletions(-) create mode 100644 vllm/v1/worker/sleep.py diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 816e63fc6ee6..a8ceecd0ac56 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -8,7 +8,6 @@ from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta -from fnmatch import filter as fnmatch_filter from types import NoneType from typing import TYPE_CHECKING, Any, cast @@ -89,6 +88,7 @@ ) from vllm.v1.utils import compute_iteration_details, report_usage_stats from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel +from vllm.v1.worker.sleep import restore_sleep_parameters, save_sleep_parameters from vllm.v1.worker.startup_plan import ( maybe_apply_startup_plan, maybe_save_startup_plan, @@ -254,18 +254,7 @@ def sleep(self, level: int = 1) -> None: model = self.model_runner.model transfer_config = self.vllm_config.weight_transfer_config patterns = transfer_config.frozen_weight_names if transfer_config else [] - parameters = dict(model.named_parameters()) - retained_names: set[str] = set() - for pattern in patterns: - matches = fnmatch_filter(parameters, pattern) - if not matches: - raise ValueError(f"No parameter matches sleep retention: {pattern}") - retained_names.update(matches) - self._sleep_saved_parameters = { - name: param.detach().to("cpu") - for name, param in parameters.items() - if name in retained_names and not param.is_cpu - } + self._sleep_saved_parameters = save_sleep_parameters(model, patterns) self._sleep_saved_buffers = { name: buffer.cpu().clone() for name, buffer in model.named_buffers() } @@ -304,11 +293,9 @@ def wake_up(self, tags: list[str] | None = None) -> None: # Restore the buffers after level 2 sleep wake_weights = tags is None or "weights" in tags if wake_weights and self._sleep_saved_parameters: - model = self.model_runner.model - with torch.no_grad(): - for name, saved in self._sleep_saved_parameters.items(): - model.get_parameter(name).copy_(saved) - self._sleep_saved_parameters.clear() + restore_sleep_parameters( + self.model_runner.model, self._sleep_saved_parameters + ) if wake_weights and len(self._sleep_saved_buffers): model = self.model_runner.model for name, buffer in model.named_buffers(): diff --git a/vllm/v1/worker/sleep.py b/vllm/v1/worker/sleep.py new file mode 100644 index 000000000000..2fc03b81f3ee --- /dev/null +++ b/vllm/v1/worker/sleep.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from fnmatch import filter as fnmatch_filter + +import torch +from torch import nn + + +def save_sleep_parameters( + model: nn.Module, patterns: list[str] +) -> dict[str, torch.Tensor]: + """Back up selected device parameters before level-2 sleep.""" + if not patterns: + return {} + parameters = dict(model.named_parameters()) + names: set[str] = set() + for pattern in patterns: + matches = fnmatch_filter(parameters, pattern) + if not matches: + raise ValueError(f"No parameter matches sleep retention: {pattern}") + names.update(matches) + return { + name: param.detach().to("cpu") + for name, param in parameters.items() + if name in names and not param.is_cpu + } + + +@torch.no_grad() +def restore_sleep_parameters(model: nn.Module, saved: dict[str, torch.Tensor]) -> None: + """Restore parameters in place and release their CPU backups.""" + for name, value in saved.items(): + model.get_parameter(name).copy_(value) + saved.clear() From 09c2c4c4bc222b7368f7cfd41b502bbfc05ac2d6 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 14:09:54 +0000 Subject: [PATCH 6/9] Reuse worker utilities for sleep parameter helpers Co-authored-by: Codex Signed-off-by: aoshen02 --- vllm/v1/worker/gpu_worker.py | 7 +++++-- vllm/v1/worker/sleep.py | 34 ---------------------------------- vllm/v1/worker/utils.py | 31 +++++++++++++++++++++++++++++++ 3 files changed, 36 insertions(+), 36 deletions(-) delete mode 100644 vllm/v1/worker/sleep.py diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a8ceecd0ac56..8b99b9502180 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -88,12 +88,15 @@ ) from vllm.v1.utils import compute_iteration_details, report_usage_stats from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel -from vllm.v1.worker.sleep import restore_sleep_parameters, save_sleep_parameters from vllm.v1.worker.startup_plan import ( maybe_apply_startup_plan, maybe_save_startup_plan, ) -from vllm.v1.worker.utils import is_residual_scattered_for_sp +from vllm.v1.worker.utils import ( + is_residual_scattered_for_sp, + restore_sleep_parameters, + save_sleep_parameters, +) from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager diff --git a/vllm/v1/worker/sleep.py b/vllm/v1/worker/sleep.py deleted file mode 100644 index 2fc03b81f3ee..000000000000 --- a/vllm/v1/worker/sleep.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from fnmatch import filter as fnmatch_filter - -import torch -from torch import nn - - -def save_sleep_parameters( - model: nn.Module, patterns: list[str] -) -> dict[str, torch.Tensor]: - """Back up selected device parameters before level-2 sleep.""" - if not patterns: - return {} - parameters = dict(model.named_parameters()) - names: set[str] = set() - for pattern in patterns: - matches = fnmatch_filter(parameters, pattern) - if not matches: - raise ValueError(f"No parameter matches sleep retention: {pattern}") - names.update(matches) - return { - name: param.detach().to("cpu") - for name, param in parameters.items() - if name in names and not param.is_cpu - } - - -@torch.no_grad() -def restore_sleep_parameters(model: nn.Module, saved: dict[str, torch.Tensor]) -> None: - """Restore parameters in place and release their CPU backups.""" - for name, value in saved.items(): - model.get_parameter(name).copy_(value) - saved.clear() diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 6c2a05a6260d..9c637f4fe865 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -4,6 +4,7 @@ from collections import defaultdict from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field +from fnmatch import filter as fnmatch_filter from itertools import product as iprod from typing import Any @@ -44,6 +45,36 @@ logger = init_logger(__name__) +def save_sleep_parameters( + model: torch.nn.Module, patterns: list[str] +) -> dict[str, torch.Tensor]: + """Back up selected device parameters before level-2 sleep.""" + if not patterns: + return {} + parameters = dict(model.named_parameters()) + names: set[str] = set() + for pattern in patterns: + matches = fnmatch_filter(parameters, pattern) + if not matches: + raise ValueError(f"No parameter matches sleep retention: {pattern}") + names.update(matches) + return { + name: param.detach().to("cpu") + for name, param in parameters.items() + if name in names and not param.is_cpu + } + + +@torch.no_grad() +def restore_sleep_parameters( + model: torch.nn.Module, saved: dict[str, torch.Tensor] +) -> None: + """Restore parameters in place and release their CPU backups.""" + for name, value in saved.items(): + model.get_parameter(name).copy_(value) + saved.clear() + + def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None: if not any(num_nans_in_logits.values()): return From 81a00bee21e374a741189c36ba595ae8fd3d0089 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 21 Sep 2026 14:10:59 +0000 Subject: [PATCH 7/9] Shorten frozen-weight retention documentation Co-authored-by: Codex Signed-off-by: aoshen02 --- docs/features/sleep_mode.md | 28 ++++++++++------------------ vllm/config/weight_transfer.py | 6 ++---- 2 files changed, 12 insertions(+), 22 deletions(-) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 5a4bbfae56b7..dd490dde825b 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -77,24 +77,16 @@ llm.wake_up(tags=["kv_cache"]) #### Retaining frozen weights during RLHF updates Set `weight_transfer_config.frozen_weight_names` to runtime parameter-name glob -patterns, such as `*.engram.*`, for weights that never change during training. -These match `model.named_parameters()`, not checkpoint keys; do not infer them from -`requires_grad=False`, which is also used for ordinary inference weights. -Each pattern must match at least one parameter, including CPU-resident parameters. - -With level 2 sleep, their GPU parameters are copied to pageable CPU memory. -`wake_up(tags=["weights"])` restores them in place and releases the CPU copies. -Each sleep/wake cycle saves and restores the current values, just like buffers. -Already CPU-resident parameters need no copy. Budget host memory for these -backups alongside trainer offload. -Level 1 retains its existing full-weight backup behavior. - -This option only controls sleep/wake; it does not modify weight loaders or -filter incoming updates. The trainer (for example, VERL) must omit these -parameters and leave their values and storage unchanged during reload and -post-processing. Restart the instance to change frozen weights. Reload paths -that replace selected parameters with meta tensors are not supported by this -option alone. Selection applies to the target, not a separate draft model. +patterns for weights that stay frozen during training. Each pattern must match +`model.named_parameters()`; checkpoint names and `requires_grad` are not used. + +Level-2 sleep backs up selected GPU parameters to pageable CPU memory; weights +wake-up restores them in place and releases the backups. CPU parameters need no +copy. Allow enough host memory for backups. Level-1 behavior is unchanged. + +The trainer must omit these parameters from updates, and reload/post-processing +must preserve their values and storage (including avoiding replacement with meta +tensors). This option does not filter updates and applies only to the target model. #### Release only KV cache memory diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index 71005597a40e..a1b1733282b6 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -17,8 +17,6 @@ class WeightTransferConfig: """ frozen_weight_names: list[str] = Field(default_factory=list) - """Runtime parameter-name glob patterns retained across level-2 sleep. - These parameters must remain immutable and be omitted by the sender. This only - controls sleep/wake, not weight loading. GPU parameters are backed up to CPU - at sleep and restored at wake; CPU-resident parameters need no copy. + """Parameter-name globs to preserve across level-2 sleep. + The sender must omit these parameters; loaders must preserve their storage. """ From 4f63af01961c88ccb34b400d5b8b91f74b8ceef2 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 22 Sep 2026 00:50:00 +0000 Subject: [PATCH 8/9] Keep sleep parameter retention in GPU worker Co-authored-by: Codex Signed-off-by: aoshen02 --- vllm/v1/worker/gpu_worker.py | 40 ++++++++++++++++++++++++++---------- vllm/v1/worker/utils.py | 31 ---------------------------- 2 files changed, 29 insertions(+), 42 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 8b99b9502180..1528423a22a8 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -8,6 +8,7 @@ from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta +from fnmatch import filter as fnmatch_filter from types import NoneType from typing import TYPE_CHECKING, Any, cast @@ -92,11 +93,7 @@ maybe_apply_startup_plan, maybe_save_startup_plan, ) -from vllm.v1.worker.utils import ( - is_residual_scattered_for_sp, - restore_sleep_parameters, - save_sleep_parameters, -) +from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager @@ -248,6 +245,31 @@ def sleep_mode_backend(self) -> "SleepModeBackend": ) return self._sleep_mode_backend + def _save_sleep_parameters(self, model: nn.Module) -> None: + transfer_config = self.vllm_config.weight_transfer_config + patterns = transfer_config.frozen_weight_names if transfer_config else [] + if not patterns: + self._sleep_saved_parameters = {} + return + parameters = dict(model.named_parameters()) + names: set[str] = set() + for pattern in patterns: + matches = fnmatch_filter(parameters, pattern) + if not matches: + raise ValueError(f"No parameter matches sleep retention: {pattern}") + names.update(matches) + self._sleep_saved_parameters = { + name: param.detach().to("cpu") + for name, param in parameters.items() + if name in names and not param.is_cpu + } + + @torch.no_grad() + def _restore_sleep_parameters(self, model: nn.Module) -> None: + for name, value in self._sleep_saved_parameters.items(): + model.get_parameter(name).copy_(value) + self._sleep_saved_parameters.clear() + def sleep(self, level: int = 1) -> None: torch.accelerator.synchronize() free_bytes_before_sleep = torch.accelerator.get_memory_info()[0] @@ -255,9 +277,7 @@ def sleep(self, level: int = 1) -> None: # Save the buffers before level 2 sleep if level == 2: model = self.model_runner.model - transfer_config = self.vllm_config.weight_transfer_config - patterns = transfer_config.frozen_weight_names if transfer_config else [] - self._sleep_saved_parameters = save_sleep_parameters(model, patterns) + self._save_sleep_parameters(model) self._sleep_saved_buffers = { name: buffer.cpu().clone() for name, buffer in model.named_buffers() } @@ -296,9 +316,7 @@ def wake_up(self, tags: list[str] | None = None) -> None: # Restore the buffers after level 2 sleep wake_weights = tags is None or "weights" in tags if wake_weights and self._sleep_saved_parameters: - restore_sleep_parameters( - self.model_runner.model, self._sleep_saved_parameters - ) + self._restore_sleep_parameters(self.model_runner.model) if wake_weights and len(self._sleep_saved_buffers): model = self.model_runner.model for name, buffer in model.named_buffers(): diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 9c637f4fe865..6c2a05a6260d 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -4,7 +4,6 @@ from collections import defaultdict from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field -from fnmatch import filter as fnmatch_filter from itertools import product as iprod from typing import Any @@ -45,36 +44,6 @@ logger = init_logger(__name__) -def save_sleep_parameters( - model: torch.nn.Module, patterns: list[str] -) -> dict[str, torch.Tensor]: - """Back up selected device parameters before level-2 sleep.""" - if not patterns: - return {} - parameters = dict(model.named_parameters()) - names: set[str] = set() - for pattern in patterns: - matches = fnmatch_filter(parameters, pattern) - if not matches: - raise ValueError(f"No parameter matches sleep retention: {pattern}") - names.update(matches) - return { - name: param.detach().to("cpu") - for name, param in parameters.items() - if name in names and not param.is_cpu - } - - -@torch.no_grad() -def restore_sleep_parameters( - model: torch.nn.Module, saved: dict[str, torch.Tensor] -) -> None: - """Restore parameters in place and release their CPU backups.""" - for name, value in saved.items(): - model.get_parameter(name).copy_(value) - saved.clear() - - def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None: if not any(num_nans_in_logits.values()): return From c07be86df8f0e8a390561d1555c8689e80e2c90b Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 22 Sep 2026 00:56:21 +0000 Subject: [PATCH 9/9] Configure parameter retention alongside sleep mode Co-authored-by: Codex Signed-off-by: aoshen02 --- docs/features/sleep_mode.md | 3 ++- vllm/config/model.py | 3 +++ vllm/config/weight_transfer.py | 7 ------- vllm/engine/arg_utils.py | 8 ++++++++ vllm/v1/worker/gpu_worker.py | 3 +-- 5 files changed, 14 insertions(+), 10 deletions(-) diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index dd490dde825b..3e419707f219 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -76,7 +76,8 @@ llm.wake_up(tags=["kv_cache"]) #### Retaining frozen weights during RLHF updates -Set `weight_transfer_config.frozen_weight_names` to runtime parameter-name glob +Set `sleep_preserve_parameter_names` (CLI: `--sleep-preserve-parameter-names`) +to runtime parameter-name glob patterns for weights that stay frozen during training. Each pattern must match `model.named_parameters()`; checkpoint names and `requires_grad` are not used. diff --git a/vllm/config/model.py b/vllm/config/model.py index 6b36afd91068..7cdba9f78fcc 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -340,6 +340,9 @@ class ModelConfig: enable_sleep_mode: bool = False """Enable sleep mode for the engine (only cuda and hip platforms are supported).""" + sleep_preserve_parameter_names: list[str] = field(default_factory=list) + """Parameter-name globs to preserve across level-2 sleep. + The sender must omit these parameters; loaders must preserve their storage.""" sleep_mode_backend: str = "cumem" """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"`` (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index a1b1733282b6..3e849753c2e6 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Literal -from pydantic import Field - from vllm.config.utils import config @@ -15,8 +13,3 @@ class WeightTransferConfig: """The backend to use for weight transfer. Validated against the `WeightTransferEngineFactory` registry at engine creation time. """ - - frozen_weight_names: list[str] = Field(default_factory=list) - """Parameter-name globs to preserve across level-2 sleep. - The sender must omit these parameters; loaders must preserve their storage. - """ diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 555f1aebb67f..a6df0d723062 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -741,6 +741,9 @@ class EngineArgs: generation_config: str = ModelConfig.generation_config enable_sleep_mode: bool = ModelConfig.enable_sleep_mode + sleep_preserve_parameter_names: list[str] = get_field( + ModelConfig, "sleep_preserve_parameter_names" + ) enable_cumem_allocator: bool = ModelConfig.enable_cumem_allocator enable_nccl_comm_suspend: bool = ModelConfig.enable_nccl_comm_suspend override_generation_config: dict[str, Any] = get_field( @@ -975,6 +978,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: model_group.add_argument( "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"] ) + model_group.add_argument( + "--sleep-preserve-parameter-names", + **model_kwargs["sleep_preserve_parameter_names"], + ) model_group.add_argument( "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"] ) @@ -1901,6 +1908,7 @@ def create_model_config(self) -> ModelConfig: generation_config=self.generation_config, override_generation_config=self.override_generation_config, enable_sleep_mode=self.enable_sleep_mode, + sleep_preserve_parameter_names=self.sleep_preserve_parameter_names, enable_cumem_allocator=self.enable_cumem_allocator, enable_nccl_comm_suspend=self.enable_nccl_comm_suspend, model_impl=self.model_impl, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 9b121296df26..df1d9de6ae40 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -246,8 +246,7 @@ def sleep_mode_backend(self) -> "SleepModeBackend": return self._sleep_mode_backend def _save_sleep_parameters(self, model: nn.Module) -> None: - transfer_config = self.vllm_config.weight_transfer_config - patterns = transfer_config.frozen_weight_names if transfer_config else [] + patterns = self.model_config.sleep_preserve_parameter_names if not patterns: self._sleep_saved_parameters = {} return