Skip to content

[BugFix][Worker] Align MRV2 dummy LM-head TP before EPLB - #17233

Merged
yiz-liu merged 1 commit into
vllm-project:mainfrom
freyfwt:fix/mrv2-lmhead-eplb-order
Sep 23, 2026
Merged

yiz-liu merged 1 commit into
vllm-project:mainfrom
freyfwt:fix/mrv2-lmhead-eplb-order

Conversation

@freyfwt

@freyfwt freyfwt commented Sep 22, 2026 •

Copy link
Copy Markdown
Collaborator

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. 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 #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.
  • 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.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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

  • Collective Order Alignment: Moved the dummy LM-head computation call to the end of the execute_model method to ensure it occurs before the EPLB step, preventing collective-order timeouts on idle DP ranks.
  • Regression Testing: Added a new parameterized test case in test_model_runner_v2_finegrained_tp.py to verify the correct execution order of forward, LM-head, and EPLB steps under various configurations.
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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

@github-actions

Copy link
Copy Markdown
Contributor

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

  • A PR should do only one thing, smaller PRs enable faster reviews.
  • Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
  • Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.


Tip

💡 Consider Linking a Related Issue or RFC

Your 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:

  • Fixes #<issue_number>
  • Closes #<issue_number>
  • Resolves #<issue_number>
  • Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

Suggested PR Title:

[Worker][BugFix] Join LM-head TP collectives before EPLB step in dummy runs

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

Comment thread vllm_ascend/worker/v2/model_runner.py Outdated
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>
@freyfwt
freyfwt force-pushed the fix/mrv2-lmhead-eplb-order branch from 84b9067 to 7398afe Compare September 22, 2026 10:07
@yjyang62

yjyang62 commented Sep 22, 2026 •

Copy link
Copy Markdown
Collaborator

/nightly QWEN3_235B_PD QWEN3_235B_PD_3_5K_1_5k
nightly command triggered.

@freyfwt freyfwt added the ready-precise run selected e2e test for pr label Sep 22, 2026
@freyfwt
freyfwt force-pushed the fix/mrv2-lmhead-eplb-order branch from cf6514c to 7398afe Compare September 22, 2026 13:25
@yiz-liu
yiz-liu merged commit ccf46c5 into vllm-project:main Sep 23, 2026
41 of 42 checks passed
tangdafu pushed a commit to tangdafu/vllm-ascend that referenced this pull request Sep 23, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module:tests ready-precise run selected e2e test for pr

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] MRV2 LM-head TP and EPLB collective order hangs idle DP ranks

3 participants