[Sleep] Retain draft-private weights with existing parameter selector (#61046)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>

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.<name>` 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
 

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

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.<name>`` 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"``

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: