diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index 0db82de9d3..bdfef2ecac 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -77,8 +77,15 @@ llm.wake_up(tags=["kv_cache"]) #### Retaining frozen weights during RLHF updates 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 +to runtime parameter-name glob patterns for weights that stay frozen during +training. It covers two use cases: + +- **Frozen Engram tables (original use case):** retain table weights and scales + omitted by the trainer, without retransmitting them after every sleep. +- **Frozen draft weights:** retain draft-private parameters when the trainer + updates only the target model. + +Target patterns 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 @@ -87,7 +94,33 @@ 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. +tensors). This option does not filter updates. + +For frozen Engram tables omitted by the trainer, select the target table weights +and scales. For example, DeepSeek V4.1 uses these runtime parameter names: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve deepseek-ai/DeepSeek-V4.1-Flash --enable-sleep-mode \ + --sleep-preserve-parameter-names '*.engram.embed_tokens.weight*' +``` + +CPU-offloaded tables remain on CPU without an extra backup copy. Engram's +[inference-time CPU offload](engram.md#cpu-offload) is configured separately; +this option preserves the frozen tables across level-2 sleep and weight updates. + +To retain a frozen drafter while updating only target weights, use `draft.*`: + +```bash +VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3.8-27B --enable-sleep-mode \ + --speculative-config '{"method":"mtp","num_speculative_tokens":3}' \ + --sleep-preserve-parameter-names 'draft.*' +``` + +vLLM selects all draft-private parameters internally; you do not need to list +them individually. Use `draft.` to select specific draft parameters. +Parameters shared with the target are excluded from the draft selection, so +target updates remain authoritative. An empty `draft.*` selection is valid. +The Engram and draft patterns can be supplied together. #### Release only KV cache memory diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 6dc20a8177..adb8f52e2e 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -19,6 +19,48 @@ from vllm.v1.worker.startup_plan import ( ) +def test_sleep_preservation_restores_private_draft_not_shared_target(): + target = torch.nn.Linear(2, 2, bias=False) + draft = torch.nn.Module() + draft.private = torch.nn.Linear(2, 2, bias=False) + draft.shared = target.weight + expected = draft.private.weight.detach().clone() + pointer = draft.private.weight.data_ptr() + worker = object.__new__(gpu_worker.Worker) + worker.get_draft_model = lambda: draft + assert set(worker._get_sleep_parameters(target)) == { + "weight", + "draft.private.weight", + } + worker._sleep_saved_parameters = {"draft.private.weight": expected} + with torch.no_grad(): + target.weight.fill_(3.0) + draft.private.weight.zero_() + worker._restore_sleep_parameters(target) + torch.testing.assert_close(draft.private.weight, expected) + assert draft.private.weight.data_ptr() == pointer + assert torch.all(target.weight == 3.0) + assert worker._sleep_saved_parameters == {} + + +@pytest.mark.parametrize("has_shared_draft", [False, True]) +def test_sleep_preservation_allows_empty_draft_but_rejects_target_typo( + has_shared_draft, +): + target = torch.nn.Linear(2, 2, bias=False) + draft = torch.nn.Module() if has_shared_draft else None + if draft is not None: + draft.shared = target.weight + worker = object.__new__(gpu_worker.Worker) + worker.get_draft_model = lambda: draft + worker.model_config = SimpleNamespace(sleep_preserve_parameter_names=["draft.*"]) + worker._save_sleep_parameters(target) + assert worker._sleep_saved_parameters == {} + worker.model_config.sleep_preserve_parameter_names = ["missing"] + with pytest.raises(ValueError, match="No parameter matches sleep retention"): + worker._save_sleep_parameters(target) + + def test_load_model_preserves_compiled_graphs_at_runtime(monkeypatch): """Profiling must use serving's thread count to keep Dynamo guards valid.""" from torch._dynamo.testing import CompileCounter diff --git a/vllm/config/model.py b/vllm/config/model.py index 502f1bbd8d..d4bafad5d6 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -372,6 +372,9 @@ class ModelConfig: hip platforms are supported).""" sleep_preserve_parameter_names: list[str] = field(default_factory=list) """Parameter-name globs to preserve across level-2 sleep. + Use target globs for frozen Engram table weights and scales omitted by the + trainer (e.g. ``*.engram.embed_tokens.weight*``). + Use ``draft.`` for draft-private parameters, or ``draft.*`` for all. 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"`` diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3efe77a3a4..04feb5da02 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -254,16 +254,28 @@ class Worker(WorkerBase): ) return self._sleep_mode_backend + def _get_sleep_parameters(self, model: nn.Module) -> dict[str, nn.Parameter]: + parameters = dict(model.named_parameters()) + draft = self.get_draft_model() + if draft is not None: + shared = {id(parameter) for parameter in parameters.values()} + parameters.update( + (f"draft.{name}", parameter) + for name, parameter in draft.named_parameters() + if id(parameter) not in shared + ) + return parameters + def _save_sleep_parameters(self, model: nn.Module) -> None: patterns = self.model_config.sleep_preserve_parameter_names if not patterns: self._sleep_saved_parameters = {} return - parameters = dict(model.named_parameters()) + parameters = self._get_sleep_parameters(model) names: set[str] = set() for pattern in patterns: matches = fnmatch_filter(parameters, pattern) - if not matches: + if not matches and pattern != "draft.*": raise ValueError(f"No parameter matches sleep retention: {pattern}") names.update(matches) self._sleep_saved_parameters = { @@ -274,8 +286,9 @@ class Worker(WorkerBase): @torch.no_grad() def _restore_sleep_parameters(self, model: nn.Module) -> None: + parameters = self._get_sleep_parameters(model) for name, value in self._sleep_saved_parameters.items(): - model.get_parameter(name).copy_(value) + parameters[name].copy_(value) self._sleep_saved_parameters.clear() def sleep(self, level: int = 1) -> None: