Repository navigation
[Refactor] Declare pipeline neighbours from a shared declaration function - #43459
Merged
Merged
Conversation
ch-wan
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
JustinTong0323,
Ying1123,
ispobock,
merrymercy,
sogalin,
wisclmy0611 and
zijiexia
as code owners
October 10, 2026 00:47
This was referenced Oct 10, 2026
ch-wan
force-pushed
the
cheng/refactor/pp-neighbour-declarations
branch
from
October 10, 2026 03:23
ca7d821 to
78ab7b5
Compare
ch-wan
force-pushed
the
cheng/refactor/exit-decisions-at-construction
branch
2 times, most recently
from
October 10, 2026 06:21
0d9d24a to
6225fa8
Compare
ch-wan
force-pushed
the
cheng/refactor/pp-neighbour-declarations
branch
2 times, most recently
from
October 10, 2026 06:28
8425495 to
a39d224
Compare
ch-wan
force-pushed
the
cheng/refactor/exit-decisions-at-construction
branch
from
October 10, 2026 19:42
6225fa8 to
06cc9bc
Compare
Base automatically changed from
cheng/refactor/exit-decisions-at-construction
to
main
October 10, 2026 19:42
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Automations to automatically generate PRs for you. |
A stage the stack binds next to, on another pipeline rank, was the declaration of a layer built on the meta device to read it, and every bound stage kept its producer's declaration. A declaration holds the stage's read, update and output transform: objects that reference the layer's modules and state. So the neighbour layers built for a rank, and on this rank each stage's predecessor, stayed referenced by the bound stages for the life of the model, though binding reads only the capabilities those objects declare. facts_of() reduces a declaration to those capabilities (UpdateFacts, ReadFacts, TransformFacts) without the declarations it is chained to. A read's kernels are not among them: binding takes kernels only from the reads of the stages it binds, so ReadFacts supplies none. The neighbouring stages and the producer each stage records are now such facts, which also replaces the history-free copy the stack made of each producer. A rank whose first stage would run the output transform of an attention on the previous rank, which this rank does not have, is rejected; no model declares one.
To bind the first and last stages of a pipeline rank's stack, make_layers built the neighbouring layers again on the meta device and read the stages they declared. Building a layer can have effects that the meta device does not stop (streams, communicators, explicit devices), and the layer's choices had to be re-derived from modules nobody runs. make_layers and make_pp_layers take stage_facts: the model's shared declaration function, which returns the stages the layer at a global index declares without building it, and which the layer uses for its own declarations too, so both ranks derive a boundary from one source. A neighbour query takes what the function returns (as DeclaredStages). Under pipeline parallelism each layer a rank builds is checked against the function, since the other ranks bind its stages from the function alone. A model on a transitional list may still omit it and have its neighbours built; any other model that declares stages under pipeline parallelism must give one. Nemotron-H, whose declarations already came from its pattern alone, gives one and leaves the list.
StageBoundary.snapshot repeated residual_batch.snapshot with one caller, LongCat-Flash, which now calls residual_batch.snapshot like the other models. StagePlan.produced forwarded to the path's output for one caller. The package exported ExitDecision, FfnExit, LayerResidualOps, NORM_READOUT, StageContract, moe_gathers_over_moe_cp and token_axis_sizes, and imported VariantEdges, though nothing imports any of them from the package; their modules still define them.
…y are The layer boundary guide said a pipeline rank builds the neighbouring layers on the meta device to read their stages, and that a stage's declaration depends only on the config and the layer index. It now describes the shared declaration function a model gives make_layers, the check of each built layer against it, and the transitional list. It gains a table of the exit policies (FFN, operation-scheduled, mixer, stack end, pipeline handoff), with which of their decisions are fixed when a plan is built and which are made per batch. ProducerReduction says that each output has one completer, the boundary, except where an FFN declares output_complete.
Seven output moves accepted **kwargs that no caller fills: an exit, a two-batch-overlap split and a test call them with the output, the residual and the batch, and bind the rest by name. Their signatures now say exactly what they take.
ch-wan
force-pushed
the
cheng/refactor/pp-neighbour-declarations
branch
from
October 10, 2026 19:42
a39d224 to
be00038
Compare
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.
This PR is part of a stack (oldest at bottom):
Motivation
To bind the first and last stages of a pipeline rank's stack,
make_layersbuilt the neighbouring layers again on the meta device and read the stages they declared. Building a layer can have effects the meta device does not stop (streams, communicators, explicit devices), its choices had to be re-derived from modules nobody runs, and every bound stage kept its neighbour's and its producer's declaration, holding their reads, updates and transforms (and so their modules) for the life of the model.Modifications
facts_of()reduces a declaration to the capabilities binding reads (UpdateFacts,ReadFacts,TransformFacts) without the declarations it is chained to; neighbours and recorded producers are such facts. A read's kernels are not among them: binding takes kernels only from the reads of the stages it binds, soReadFactssupplies none. A rank whose first stage would have to run the output transform of an attention on the previous rank is rejected (no model declares one); a test covers it.make_layers/make_pp_layerstakestage_facts: the model's shared declaration function, which returns the stages the layer at a global index declares, and which the layer also uses for its own declarations. Under PP each layer a rank builds is checked against it; a mismatch names the layer, the stage and each differing field with both values. Models not yet migrated stay on a transitional list that keeps the meta-device path; Nemotron-H, whose declarations already came from its pattern, gives a function and leaves the list. [Refactor] Declare every model's stages from a shared function and drop the meta-device neighbour build #43460 migrates every other model and removes the list.Accuracy Tests
How the stack is checked:
test_aiter_allreduce_fusion_amdcases that fail on NVIDIA on main too; the test files that changed most recently pass on every commit. The fulltest/registered/unitsuite (926 files, one process per file) shows no failure on the stack's head that main does not have.516fd1c1b5unless marked as on the merge candidate; rebasing onto current main had no conflicts, and the CPU checks were run again after it.961bcf481d, H200, asymmetric dummy weights unless noted): DeepSeek-V2-Lite input-scattered PP2 is bit-identical to PP1 (main crashes on it); Nemotron-H cut between a mixer and its FFN differs from PP1 only at the rounding level of the FFN-to-mixer cuts (main: 0.93 and 2.21); Llama, Qwen3-MoE, GPT-OSS and DeepSeek-V2-Lite are bit-identical to main at PP1 and PP2; GSM8K (real DeepSeek-V2-Lite weights, 250 questions, TP4 with attention DP2 and the dense FFN fully data-parallel) gives 0.360 at PP2, 0.364 at PP1 and 0.368 on main at PP1, medians of three runs whose run-to-run spread is 0.012–0.016. The last revision of the stack changed only comments, docstrings, an equivalent condition and explicitly set attributes after these runs; its bound plans are identical.For this PR: the plan comparison differs only in how a pipeline rank's neighbour appears in a bound plan (a
ReadFactsinstead of the read object), on an edge whose read does not run on the sending rank.Speed Tests and Profiling
No forward-path change.
Checklist
CI States
Latest PR Test (Base): 🚫 Run #38080873880
Latest PR Test (Extra): 🚫 Run #38080873610
Latest PR Test (AMD ROCm 10): 🚫 Run #38080873786