Repository navigation
[Bugfix][EPLB] Start the async rearrangement interval after the last layer commits - #64485
Conversation
…layer commits In async mode the rearrangement counter was only reset when a cycle started. A cycle longer than step_interval drove "steps until the next rearrangement" negative and started the next rearrangement in the same step the last layer committed, so its physical-slot load window still held rows recorded under the old placement. Reset the counter again once the last layer of the last in-flight model commits, log committed-layer progress instead of a negative count while a cycle is in flight, and warn on rank 0 every step_interval steps a cycle is still running. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bowen Wang <bowen@inferact.ai>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| logger.warning( | ||
| "Async EPLB rearrangement still in progress after %d steps: %s", | ||
| self.expert_rearrangement_step, | ||
| ", ".join( | ||
| f"{ms.model_name} {ms.num_committed_layers}/" | ||
| f"{ms.model.num_moe_layers} layers committed" | ||
| for ms in self.model_states.values() | ||
| ), | ||
| ) |
There was a problem hiding this comment.
since it only fires on rank 0, it can't tell us which rank is stuck. can we find a way to show which rank is behind? but maybe takes some extra communication cost, we can use debug flags open it
There was a problem hiding this comment.
Good point, done in 4a71a57.
- The async worker now records how far it has got in the current rearrangement: which layer, and whether that layer's result has been published.
- When the warning fires, every EP rank writes its progress into its own slot of an int64 tensor of EP size, and a sum all-reduce over the EP CPU group gathers them.
- Rank 0 then logs the ranks that are furthest behind and where they are:
Async EPLB rearrangement still in progress after 30 steps: 1/40 layers of <model> committed, slowest EP ranks [0, 1]: transferring layer 1
I didn't put it behind a flag:
- The all-reduce only runs when the warning fires, which is every
step_intervalsteps while a cycle is overrunning. - While a cycle is in flight, async mode already does a CPU all-reduce on the EP group every step (
_all_ranks_result_ready). - In a GB200 run with
step_interval=600, the warning never fired.
What it can and cannot tell you:
- A rank stuck outside a collective is named in every warning while it stays stuck. For example, it may still be computing the new placement, or still finishing its own transfer after the others have published. In a two-rank gloo simulation, the warning named the delayed rank.
- When the ranks wait for each other inside a layer's transfer, they all report the same layer, so the warning lists all of them. It still names the layer. These waits are:
- the coordinated-stop all-reduce at the start of each layer;
- the communicator's exchange: the gloo/XCCL send/recv waits, the NIXL post-read barrier, or the symm_mem barriers.
- On GB200 (DP2, gloo communicator) I ran with
step_interval=30, so every cycle overran but no rank was actually slow.- 91 of 94 warnings listed both ranks.
- The other 3 named one rank because of a short race: that rank had not yet returned from the commit of layer N, while the other was already waiting for it at layer N+1.
|
This pull request has merge conflicts that must be resolved before it can be |
Carry the rearrangement status through the deferred balancedness log added in vllm-project#62009. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bowen Wang <bowen@inferact.ai>
The in-progress warning only reported committed layers, so a stall did not show which rank's async worker was behind. The worker now records its progress through each layer (transferring, published). Every EP rank reaches the warning in the same step, so one small CPU all-reduce there gathers each rank's progress, and rank 0 names the ranks that are furthest behind and where they are. This only runs when a cycle overruns step_interval. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bowen Wang <bowen@inferact.ai>
|
This pull request has merge conflicts that must be resolved before it can be |
Move the worker progress marks into _transfer_cycle, which vllm-project#46370 split out of transfer_run_periodically. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bowen Wang <bowen@inferact.ai>
ilmarkov
left a comment
ilmarkov
left a comment
There was a problem hiding this comment.
Thanks for the logging improvement! Layer progress makes much more sense rather than negative steps. Just small nits.
| ep_rank=ep_group.rank(), | ||
| ) | ||
| if not self._async_rearrangement_in_progress(): | ||
| # Count the interval from when the new placement is fully |
There was a problem hiding this comment.
Could you improve the wording a bit? "Start the next interval after the final commit, allowing the load window to fill under the new placement"
| model_state.rebalanced for model_state in self.model_states.values() | ||
| ) | ||
|
|
||
| def _warn_async_rearrangement_in_progress(self) -> None: |
There was a problem hiding this comment.
Could you add a two-rank Gloo test for an overrun? It would verify that both ranks enter the new collective, only rank 0 logs, and the warning identifies the rank with the lowest worker progress. The PR mentions a manual simulation of this, but the added automated test covers only restart timing.
There was a problem hiding this comment.
Added in e9fb25b: test_async_rearrangement_overrun_warning_names_slowest_rank.
- Two gloo processes run the real
EplbState.step()at the step where the cycle overrunsstep_interval. - Rank 1's worker is one stage behind rank 0: rank 0 has published layer 1, rank 1 is still transferring it. So rank 0 can only name rank 1 through the all-reduce.
- It checks that rank 0 logs once and names rank 1 and
transferring layer 1, and that rank 1 logs nothing. - If one rank skips the all-reduce, the other times out after 30 s and the test fails instead of hanging.
- It builds a gloo
GroupCoordinatorwithout a device communicator, liketest_cpu_all_reduce.py, so it also runs on the single-GPU "EPLB Algorithm" step.
It fails with main's eplb_state.py, and also when the all-reduce runs only on rank 0, without the all-reduce, with max instead of min, or when every rank logs.
Add a two-rank gloo test: while an async cycle overruns step_interval, both ranks enter the progress all-reduce, only rank 0 logs, and the warning names the rank that is behind. Reword the comment on the interval reset. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Bowen Wang <bowen@inferact.ai>
|
/ci run |
|
✅ Triggered Buildkite CI #101816 for commit ℹ️ This PR is 12 commits behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #102056 for commit ℹ️ This PR is 1 commit behind upstream |
Port vllm-project#64485 onto the async runner: the rearrangement interval restarts when the cycle's last layer commits, and the overrun warning reads the running layer and the ranks still busy from the runner instead of a worker-written progress field. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Shuojie Wang <latioswang@gmail.com>
Overview
Async EPLB resets the rearrangement counter only when a cycle starts. A cycle
longer than
step_intervaltherefore logs a negative "steps until the nextrearrangement" and starts the next rearrangement in the same step its last
layer commits. This PR counts the interval from the last layer's commit and
logs in-flight progress instead. If a cycle overruns, a warning names the
slowest EP ranks; when the ranks wait for each other inside a transfer, it
names all of them.
Related: #57883 (stall that shows the negative count), #60685 / #60686 (load
window credited to the wrong experts), #57905.
Claims
when the last layer of the last in-flight model commits.
step_interval >= window_size(default 3000 / 1000), every window rowthat the next rearrangement uses is recorded after that commit, however long
the cycle ran.
step_interval - window_sizeleavespre-commit rows in the window. [Bug]: EPLB counts load for the wrong expert when step_interval < window_size #60685 notes that async EPLB has its problem
per layer. This also hits the defaults (3000 / 1000) once a cycle runs past
2000 steps, which [Bug]: EPLB counts load for the wrong expert when step_interval < window_size #60685 does not cover.
step_interval < window_sizecase of [Bug]: EPLB counts load for the wrong expert when step_interval < window_size #60685 is not fixed here. Norare ranks that ran dummy steps (see Limitations).
async rearrangement in progress: k/L layers committed.step_intervalsteps while a cycle is still running, rank 0 warns.The warning names the EP ranks whose async worker is furthest behind,
and the layer and stage they are at. A stall like [Bug]: Async EPLB silently stops making progress — rearrangement step counter runs away negative forever, no warning/error (EP16 cross-node, use_async=True, communicator=nixl) #57883 shows up even with
log_balancednessoff.Validation
GPU end-to-end (GB200, Qwen3.5-35B-A3B, 2 GPUs with EP)
Setup:
num_redundant_experts=2,log_balancedness_interval=1.torch_gloo(CPU-staged), because NIXL is not installed inthis venv. The cycle lengths below are for that path.
GSM8K row instead used the fa3b9d1 wheel.
PYTHONPATHthatdiffers from the wheel only in this PR's files. The md5 of
eplb_state.pyis logged before each server start.window_size=step_interval=30. A cycle takes 7.3-8.9 s(140-370 steps), so every cycle overruns the interval.
Cadence (target model):
cycle, which was still in flight when the run ended.
window_size=50,step_interval=600, TP2With the drafter, both models report the in-progress status. The drafter's
single layer commits first. The counter resets once the target's last layer
commits.
Sample PR lines (DP2):
Every cycle in this run overran, but no rank was actually slow. Of the 94
DP2 warnings:
in that layer's coordinated-stop all-reduce or gloo exchange.
published layer N, either waiting for the commitor just past it.
from the commit of layer N, while the other was already waiting for it at
layer N+1.
GSM8K accuracy, one value per round:
step_interval=600, TP2(MTP main: 0.8324 to 0.8976), so a single-round difference is noise.
weights are moved or how tokens are routed.
Which PR revision each arm ran:
merge of main. Its warning names no ranks.
step_interval=600, MTP and DP2 rows ran 4a71a57. Later commits onlymerge main ([Feature] Initial support for fault tolerant ep using scale-down #46370 moved the worker loop into
_transfer_cycle, and the twoprogress writes moved with it), reword a comment and add the second unit test.
Unit tests
tests/distributed/test_eplb_utils.pyis run by the "EPLB Algorithm" CI step.test_async_rearrangement_interval_starts_after_last_layer_commituses a realsingle-rank parallel state (
dist_init), the realEplbState.step()/rearrange()/_move_to_workspace(), and a cycle longer thanstep_interval. It plays the async worker, publishing one layer result everytwo steps.
test_async_rearrangement_overrun_warning_names_slowest_rankruns two glooprocesses through the real
EplbState.step()at the overrun step, with rank1's worker one stage behind rank 0. Rank 0 must log once and name rank 1
("transferring layer 1"), and rank 1 must not log. A rank that skips the
all-reduce makes the other one time out after 30 s.
On main, the second rearrangement starts at step 6, the step that commits the
last layer. With this PR it starts at step 10 (= 6 +
step_interval).The warning test also fails if the all-reduce runs only on rank 0 (timeout),
without the all-reduce, with
maxinstead ofmin, or if every rank logs.CPU simulation of the window misattribution
Setup:
EplbState, with 4 layers and 440 steps per layer.step_interval=1500,window_size=1000.transfer_run_periodically. The router splits eachlogical expert's demand evenly over its replicas.
Other checks:
step_interval >= window_sizecase showed 0% misattribution with thisPR. The sweep covered
log_balancednesson and off, log intervals1 / 250 / 2000, 20 layers, and a last commit at counter I-1, I and I+1.
delays and dummy steps:
collective, and listed both ranks when they were held inside one.
Details
Root cause. In
EplbState.step(),expert_rearrangement_stepis resetjust before
rearrange()starts a cycle.step_intervalwhile a cycle is still in flight,the bail-out returns without resetting it. The counter keeps growing, so the
log goes negative. A production log showed
-250withstep_interval=1500.step()that commits the last layer then seescounter >= step_intervalwith nothing in flight, and starts the next cycleat once.
map. So rows recorded before each layer's commit are credited to the wrong
logical experts.
Change.
_move_to_workspace, once no model is still rebalancing, set thecounter back to 0. The
_all_ranks_result_readyall-reduce gates the commit,so it happens in the same step on every rank and the counter stays
consistent.
num_committed_layersfor the progress log.rearrange()setsworker_progressto -1. The async worker then writes2 * layerwhile transferring a layer, and2 * layer + 1once its resultis published.
zero int64 tensor of EP size. A sum all-reduce over the EP CPU group then
gathers them. This costs one small CPU all-reduce, only when the warning
fires. While a cycle is in flight, async mode already does a CPU all-reduce
on that group every step (
_all_ranks_result_ready).Trade-off. In async mode, the period goes from max(
step_interval, D) tostep_interval + D, where D is the cycle length. Here D was 7.3-8.9 s withthe gloo communicator, or 140-370 steps under GSM8K load. With the default
interval of 3000 steps, that makes the period 5-12% longer: about 10% without
the drafter, and about 5% with MTP. The
step_intervaldocstring now says whenthe interval starts.
Not done: clearing the window on commit.
step_interval >= window_size, clearing changes nothing on rankswithout dummy steps, because the window is fully overwritten after the
commit.
[Bugfix][EPLB] Record the load window by logical expert #60686's logical window fixes that without throwing data away.
step_interval < window_size, clearing shrinks the effective windowto
step_interval, which [Bugfix][EPLB] Record the load window by logical expert #60686's review objected to.Limitations.
step_interval < window_size, some window rows still come from beforethe commit.
not advance the window ring.
both cases.
of them. It still names the layer.
Not a duplicate.
keeps the counter overrun, the negative log and the back-to-back cycles.
cycle_in_flightflag, setand cleared together with
rebalanced.the ranks to disagree on
rebalanced.lines in
step().merges main and computes the in-progress status when the entry is queued.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
AI assistance (Claude Code) was used for the investigation, the code and this description.
🤖 Generated with Claude Code