[Bugfix] DeepSeek-V4 DP+PIECEWISE: force aiter unreg collective capture path#1596
Open
ZhangLirong-amd wants to merge 1 commit into
Open
[Bugfix] DeepSeek-V4 DP+PIECEWISE: force aiter unreg collective capture path#1596ZhangLirong-amd wants to merge 1 commit into
ZhangLirong-amd wants to merge 1 commit into
Conversation
…re path
DP-attention + PIECEWISE cudagraph corrupts long generations (repeated-token
garbage; GSM8K ~0.87 vs 0.94 baseline). Root cause is in aiter's custom
all_gather/reduce_scatter 'registered' capture path used by the DP-MoE
all_gather/reduce_scatter collectives.
The registered path lets each collective kernel directly read every peer's
ORIGINAL input pointer, cross-registered at register_graph_buffers(). This is
only safe under a single whole-forward FULL cudagraph, whose global
read/overwrite ordering of every collective's input holds across all ranks.
PIECEWISE splits the forward into many small per-piece graphs with eager
sections between them, losing that ordering: the caching allocator recycles the
same input address across layers, so a fast rank can overwrite its input via a
later piece while a slow peer is still reading it -> stale cross-rank reads ->
progressive hidden corruption -> repeated-token garbage. Only DP hits this
(TP uses in-place all_reduce; DP+EP uses all2all); FULL and eager are clean.
Fix: when the runner uses PIECEWISE, set enable_register_for_capturing=False on
the aiter tp/dp/ep ca_comm before capture, forcing the copy-in ('unreg') path
that snapshots the input into a pre-registered pool before each collective and
is therefore order-independent (same path as eager / gfx1250). FULL is
unaffected and keeps the faster direct-read path.
Verified: DP+PIECEWISE GSM8K back to 0.94+ (from 0.87), no repeated-token
collapse; TP and FULL paths unchanged.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR mitigates a DeepSeek-V4 correctness bug seen specifically under DP-attention + --cudagraph-mode PIECEWISE by forcing aiter’s DP/TP/EP collective communicator to use the unregistered (“copy-in”) cudagraph capture path, avoiding cross-rank stale reads caused by allocator pointer reuse across piecewise captures.
Changes:
- Add
ModelRunner._force_aiter_unreg_capture_for_piecewise()to disableca_comm.enable_register_for_capturingwhen PIECEWISE cudagraph mode is active. - Invoke that helper early in
ModelRunner.capture_cudagraph()when_piecewise_cg_active()is true.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+2438
to
+2445
| seen = set() | ||
| for getter in ("get_tp_group", "get_dp_group", "get_ep_group"): | ||
| try: | ||
| from aiter.dist import parallel_state as _ps | ||
|
|
||
| group = getattr(_ps, getter)() | ||
| except Exception: | ||
| continue |
valarLip
approved these changes
Jul 15, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
DeepSeek-V4 DP+PIECEWISE corrupts long outputs (repeated tokens, GSM8K 0.87 vs 0.94). Root cause: aiter custom ag/rs registered read-path relies on FULL whole-graph cross-rank ordering, which PIECEWISE breaks — pool-recycled input is overwritten before peers read it. Fix: force the unreg copy-in path under PIECEWISE. Restores 0.94+.