Repository navigation
[Feature][Sleep] Preserve selected target and draft-private state across level-2 sleep - #18624
Silas-Zeng wants to merge 4 commits into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces robust state preservation for Ascend NPU workers during level-2 sleep cycles. By porting upstream retention semantics, it ensures that critical target and draft model parameters and buffers are correctly snapshotted and restored, maintaining model state integrity across sleep and wake operations without disrupting existing allocator or cleanup flows. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [Feature] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Support parameter and draft buffer preservation during sleep/wake cyclesSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces parameter and draft buffer preservation during level-2 sleep and wake-up cycles in the NPU worker. It saves model parameters and draft model buffers to CPU memory before sleep and restores them upon wake-up, ensuring state preservation. It also adds comprehensive unit tests to verify these behaviors.
Feedback on the changes includes addressing a runtime AttributeError due to the use of non-existent is_cpu on PyTorch tensors, defensively handling missing get_draft_model attributes on the model runner, safely retrieving configuration patterns, and adding validation checks for missing parameters during wake-up.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New unit tests were added in tests/ut/worker/test_sleep_wake_parameter_preservation.py and existing tests were updated.| self._sleep_saved_parameters = { | ||
| name: parameter.detach().to("cpu") | ||
| for name, parameter in parameters.items() | ||
| if name in names and not parameter.is_cpu | ||
| } |
There was a problem hiding this comment.
In PyTorch, torch.Tensor (and nn.Parameter) does not have an is_cpu attribute. Attempting to access parameter.is_cpu will raise an AttributeError at runtime when saving sleep parameters. Check the tensor's device type instead using parameter.device.type != "cpu".
| self._sleep_saved_parameters = { | |
| name: parameter.detach().to("cpu") | |
| for name, parameter in parameters.items() | |
| if name in names and not parameter.is_cpu | |
| } | |
| self._sleep_saved_parameters = { | |
| name: parameter.detach().to("cpu") | |
| for name, parameter in parameters.items() | |
| if name in names and parameter.device.type != "cpu" | |
| } |
| def get_draft_model(self) -> nn.Module | None: | ||
| return self.model_runner.get_draft_model() |
There was a problem hiding this comment.
Defensively check if self.model_runner implements get_draft_model before calling it. This prevents potential AttributeError crashes if a custom or mocked model runner is used.
| def get_draft_model(self) -> nn.Module | None: | |
| return self.model_runner.get_draft_model() | |
| def get_draft_model(self) -> nn.Module | None: | |
| get_draft_model_fn = getattr(self.model_runner, "get_draft_model", None) | |
| if get_draft_model_fn is not None: | |
| return get_draft_model_fn() | |
| return None |
| @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(): | ||
| parameters[name].copy_(value) | ||
| self._sleep_saved_parameters.clear() |
There was a problem hiding this comment.
Add a check to ensure that all saved parameters are present in the model during wake-up, raising a descriptive RuntimeError if any are missing. This is consistent with how model and draft buffers are handled and prevents cryptic KeyError crashes.
| @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(): | |
| parameters[name].copy_(value) | |
| self._sleep_saved_parameters.clear() | |
| @torch.no_grad() | |
| def _restore_sleep_parameters(self, model: nn.Module) -> None: | |
| parameters = self._get_sleep_parameters(model) | |
| missing = set(self._sleep_saved_parameters) - set(parameters) | |
| if missing: | |
| raise RuntimeError(f"Saved model parameters are missing during wake-up: {sorted(missing)}") | |
| for name, value in self._sleep_saved_parameters.items(): | |
| parameters[name].copy_(value) | |
| self._sleep_saved_parameters.clear() |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
9bdd5e5 to
2561108
Compare
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Silas-Zeng <1784012765@qq.com>
|
Hi @Ronald1995, could you take a look at this sleep/wake change, or help route it to the appropriate reviewer? I encountered this issue during RL rollout. This PR ports upstream #57891 and #61046 to Ascend, preserving selected target parameters and draft-private state across level-2 sleep and target-only weight updates. I validated Qwen3-8B + DFlash on Ascend 910B2 with both MRV1 and MRV2; the workflow and results are described above. Thanks! |
1 participant
What this PR does / why we need it?
Fixes #18641
Level-2 sleep discards weights, but a target-only weight update does not reload the draft model. This PR ports the following merged upstream changes to Ascend so selected target parameters and draft-private state survive this workflow:
The worker uses
model_runner.get_draft_model()for MRV1 and MRV2, de-duplicates shared parameters, and restores saved parameters and buffers in place when waking weights. Snapshots remain available across KV-cache-only wakes and are cleared after tensor restoration succeeds.Does this PR introduce any user-facing change?
Yes. Ascend now honors
sleep_preserve_parameter_namesduring level-2 sleep, includingdraft.*for draft-private parameters. This allows users to retain selected weights while updating only the target model.How was this patch tested?
Tested Qwen3-8B + DFlash serving with both MRV1 (
VLLM_USE_V2_MODEL_RUNNER=0) and MRV2 on one Ascend 910B2 with CANN 9.1.For each runner, generated a baseline response, ran
sleep(level=2)→wake(weights)→ target-onlycollective_rpc(reload_weights)→wake(kv_cache), then generated again. Both runs completed successfully with matching output, resumed draft generation, and/is_sleepingreturningfalse.