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 addresses a critical issue in RL disaggregated training where the MTP draft model failed to synchronize weights with the target model. By ensuring both models are updated simultaneously during the weight update cycle, the MTP acceptance rate is restored, significantly improving performance and training stability on NPU hardware. 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. |
![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 weight updates for speculative draft models on AscendSuggested PR Summary:
### What this PR does / why we need it?
This PR adds support for updating weights of MTP/Eagle draft models alongside the target model on Ascend NPUs. It introduces `_get_draft_model` to retrieve and unwrap the draft model, and updates `start_weight_update`, `update_weights`, and `finish_weight_update` to handle draft models in both checkpoint and direct loading formats.
Feedback:
- In `_get_draft_model`, the check `hasattr(drafter, "get_model")` is incorrect and should be performed on `draft_model` instead of `drafter` to correctly unwrap the `ACLGraphWrapper`.
- In direct weight loading mode, the `_loader` function should run under `torch.no_grad()` to avoid tracking gradients and causing unnecessary memory overhead.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Not specified.| if hasattr(drafter, "get_model"): | ||
| with suppress(Exception): | ||
| draft_model = drafter.get_model() |
There was a problem hiding this comment.
The check hasattr(drafter, "get_model") and subsequent call drafter.get_model() are incorrect because drafter is the speculative proposer (e.g., AscendEagleProposer), whereas the draft model itself (draft_model) is the one that may be wrapped in ACLGraphWrapper. This typo will prevent the draft model from being unwrapped, which can cause hasattr(draft_model, "load_weights") to return False and disable weight updates for the draft model entirely.
| if hasattr(drafter, "get_model"): | |
| with suppress(Exception): | |
| draft_model = drafter.get_model() | |
| if hasattr(draft_model, "get_model"): | |
| with suppress(Exception): | |
| draft_model = draft_model.get_model() |
| def _loader(weights): | ||
| for name, weight in weights: |
There was a problem hiding this comment.
In direct mode, copying weights and invoking weight_loader should be performed under torch.no_grad() to prevent autograd from tracking these operations and constructing unnecessary gradient graphs, which can lead to memory overhead or errors.
| def _loader(weights): | |
| for name, weight in weights: | |
| def _loader(weights): | |
| with torch.no_grad(): | |
| for name, weight in weights: |
CalvinXKY
changed the title
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
NPUWorker's weight update methods (start/update/finish) only updated the target model, leaving the MTP draft model at HF checkpoint init values. This caused 0% MTP acceptance rate and garbage rollout output in RL training disaggregated mode. Changes: - Add _get_draft_model() to retrieve MTP draft model from AscendEagleProposer - start_weight_update: also initialize layerwise_reload for draft model - update_weights: load weights into both target and draft model; replace naive param.copy_() with TP-sharding-aware loader that handles packed params (qkv_proj, gate_up_proj) and TP rank sharding - finish_weight_update: also finalize layerwise_reload for draft model Non-MTP models are unaffected (_get_draft_model returns None). Signed-off-by: kaiyuan <kyxiezju@163.com>
Signed-off-by: kaiyuan <kyxiezju@163.com>
e2579b1 to
88d54d3
Compare
Replace engine getattr side-channel with explicit set_draft_model API, tighten draft unwrap logic, and cover worker/engine draft sync paths. Signed-off-by: kaiyuan <kyxiezju@163.com>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
1 participant
What this PR does / why we need it?
In RL disaggregated training, weight sync only updated the target model and skipped the MTP draft (
model_runner.drafter). After the first update, draft stayed at HF init while target moved, so MTP acceptance dropped to 0% and rollout became garbage.After #12876, weight load lives in transfer engines; this PR wires draft sync through that path:
NPUWorker:_get_draft_model()/_bind_draft_model_to_engine(); onstart_weight_updatebind draft viaengine.set_draft_model(...)HCCLWeightTransferEngine/NPUIPCWeightTransferEngine: explicitset_draft_model; start/finish/_load_weights_with_draftalso update draft when presentis_checkpoint_formatkept on workerstart_weight_updatefor RL-client compat but ignored (engine owns format handling)Non-MTP models unchanged (
_get_draft_model()returnsNone).Companion vime fix: vllm-project/vime#375
Does this PR introduce any user-facing change?
No. Fixes RL weight-sync for MTP on NPU; non-MTP / non-RL paths unchanged.
How was this patch tested?
Unit tests
tests/ut/worker/a2/test_worker_v1.py: draft unwrap +set_draft_modelbind on starttests/ut/distributed/weight_transfer/test_npu_ipc_engine.py: draft start/finish/loadtests/ut/distributed/weight_transfer/test_hccl_engine_draft.py: HCCL draft helpersE2E with vime disaggregated RL on 8x Ascend 910B1, GLM-4.7-Flash (30B-A3B MoE + MTP), 4 train + 4 rollout, TP=4 (together with vime#375):