Repository navigation
[BugFix][Worker] Align MRV2 dummy LM-head TP before EPLB - #17233
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 collective-order timeout issue in the MRV2 worker configuration. By reordering the dummy LM-head computation to precede the EPLB step, it ensures that idle DP ranks correctly participate in required collectives, preventing hangs in multi-node NPU deployments. 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 [BugFix] 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! |
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Worker][BugFix] Join LM-head TP collectives before EPLB step in dummy runsSuggested PR Summary:
### What this PR does / why we need it?
This PR moves the LM-head TP collective joining logic from `_dummy_run` to `execute_model` during dummy runs. This ensures that idle ranks join the LM-head TP collectives before the EPLB step can block them, preventing potential hangs. Additionally, the unit tests have been refactored and parametrized to verify this execution order.
Feedback: A potential `IndexError` was identified in `execute_model` if `hidden_states` has zero rows. We should check the shape of `hidden_states` before indexing to prevent crashes on idle ranks.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The changes are tested using the updated unit test `test_dummy_lmhead_collective_precedes_eplb` in `tests/ut/worker/test_model_runner_v2_finegrained_tp.py`.Defer the upstream dummy EPLB step until after the existing Ascend LM-head collective by using skip_eplb on the parent and applying its decorator to the Ascend wrapper. Preserve the parent dummy-forward and speculator ordering. Signed-off-by: freyfwt <freytian1996@gmail.com>
84b9067 to
7398afe
Compare
|
/nightly QWEN3_235B_PD QWEN3_235B_PD_3_5K_1_5k
|
cf6514c to
7398afe
Compare
…t#17233) ### What this PR does / why we need it? Fixes vllm-project#17232. When a decode DP rank is idle, the MRV2 dummy path currently enters the parent `_dummy_run()` EPLB step before it joins the LM-head TP collective. An active rank samples through LM-head TP first and advances EPLB afterward. In the QWEN3_235B_PD nightly configuration, this opposite collective order produces an HCCL timeout. Keep the existing dummy `compute_logits()` call after the parent's dummy forward and speculator. Pass `skip_eplb=True` to the parent's decorated `_dummy_run()`, then apply the same upstream `step_eplb_after(is_dummy=True)` decorator to the Ascend override. The resulting order is forward, dummy LM-head, EPLB, without unwrapping or changing upstream methods. Preserve `skip_eplb`, profiling, non-last-stage, and fixed logits capacity behavior. Add regression tests for the collective order and skip paths. The observed reproduction, including the pinned YAML and required PD, DP, EP, EPLB, and LM-head TP settings, is in vllm-project#17232. This change overlaps the target-side dummy LM-head ordering part of vllm-project#15079; that PR also covers speculative draft alignment. ### Does this PR introduce _any_ user-facing change? Yes. The affected MRV2 PD deployment can complete requests without this collective-order timeout. ### How was this patch tested? - `PYTHONPATH=<this checkout>:<vLLM checkout at 84030bbe3d74d99bad477a3d2e37a973ccd8865c> python -m pytest -q -p no:cacheprovider tests/ut/worker/test_model_runner_v2_finegrained_tp.py tests/ut/worker/test_adaptive_verification.py` — 33 passed. - `bash format.sh ci` — passed. - Full mypy for Python 3.10, 3.11, and 3.12 in the project lint environment — passed. - Real HCCL A/B reproduction on one Ascend A3 node with two NPUs: the unpatched runner enters EPLB `AllReduce` on the idle rank while the active rank enters LM-head `AllGather`, then times out; this PR completes both collectives in the same order on both ranks. - Model-backed A/B on the same two NPUs with local Qwen3-30B-A3B-W8A8 weights and the same vLLM commit: MRV2, DP=2, TP=1, EP enabled, LM-head TP=2, EPLB enabled (`window_size=4000`, `step_interval=200`, `num_redundant_experts=0`, async and balancedness logging enabled), eager mode, async scheduling, `max_model_len=2048`, `max_num_batched_tokens=256`, `max_num_seqs=4`. With the parent commit `5f8a1286a`, the first completion request returns HTTP 500; the idle worker's upstream dummy EPLB step times out in HCCL `AllReduce` after 60 seconds. With this PR commit `7398afeae`, the same request returns HTTP 200 and 16 output tokens while EPLB keeps stepping. The exact three-node NPU nightly deployment and its graph mode have not been rerun with this patch; the linked failing job provides the original hardware trace. - vLLM main: vllm-project/vllm@84030bb Signed-off-by: freyfwt <freytian1996@gmail.com>
What this PR does / why we need it?
Fixes #17232.
When a decode DP rank is idle, the MRV2 dummy path currently enters the parent
_dummy_run()EPLB step before it joins the LM-head TP collective. An active rank samples through LM-head TP first and advances EPLB afterward. In the QWEN3_235B_PD nightly configuration, this opposite collective order produces an HCCL timeout.Keep the existing dummy
compute_logits()call after the parent's dummy forward and speculator. Passskip_eplb=Trueto the parent's decorated_dummy_run(), then apply the same upstreamstep_eplb_after(is_dummy=True)decorator to the Ascend override. The resulting order is forward, dummy LM-head, EPLB, without unwrapping or changing upstream methods. Preserveskip_eplb, profiling, non-last-stage, and fixed logits capacity behavior. Add regression tests for the collective order and skip paths.The observed reproduction, including the pinned YAML and required PD, DP, EP, EPLB, and LM-head TP settings, is in #17232. This change overlaps the target-side dummy LM-head ordering part of #15079; that PR also covers speculative draft alignment.
Does this PR introduce any user-facing change?
Yes. The affected MRV2 PD deployment can complete requests without this collective-order timeout.
How was this patch tested?
PYTHONPATH=<this checkout>:<vLLM checkout at 84030bbe3d74d99bad477a3d2e37a973ccd8865c> python -m pytest -q -p no:cacheprovider tests/ut/worker/test_model_runner_v2_finegrained_tp.py tests/ut/worker/test_adaptive_verification.py— 33 passed.bash format.sh ci— passed.AllReduceon the idle rank while the active rank enters LM-headAllGather, then times out; this PR completes both collectives in the same order on both ranks.window_size=4000,step_interval=200,num_redundant_experts=0, async and balancedness logging enabled), eager mode, async scheduling,max_model_len=2048,max_num_batched_tokens=256,max_num_seqs=4. With the parent commit5f8a1286a, the first completion request returns HTTP 500; the idle worker's upstream dummy EPLB step times out in HCCLAllReduceafter 60 seconds. With this PR commit7398afeae, the same request returns HTTP 200 and 16 output tokens while EPLB keeps stepping.The exact three-node NPU nightly deployment and its graph mode have not been rerun with this patch; the linked failing job provides the original hardware trace.