Conversation
The content of pull request 4312 as reviewed (head 087c4d1, fourteen commits) rebased onto main after the expert-parallel merge, as one commit: the block attention residual travels between stages as a second payload, the cross-stage cache adapter ships only the blocks a receiver does not hold, the split puts the residual's final aggregation with the head, and the debug flavors and integration cells for the pipeline x virtual-stage matrix. Conflicts were only the debug flavor signatures (main's moe_comm_backend argument and its (config, max_context_len) registry) and the shape-suffix legend, resolved to main's.
The shape-suffix legend goes back to its original lines (the reflow was churn). The block-opening test is named first_layer_in_block, as suggested. The layout inference the adapter calls from layout.py is public now; it was a private name used across files. The pp2 integration cell is enough for CI; the pp8 x vp4 cell goes (its recipe flavor stays for local matrices). (cherry picked from commit e326c70)
(cherry picked from commit 0c0e489)
Three review items on pipeline_adapter.py. The transport switch attn_res_cache is a property of the pipeline transport, not of the model, so it leaves the model config and becomes an argument of pipeline_kimi_k3 (a recipe passes functools.partial(pipeline_kimi_k3, attn_res_cache=False) as the pipelining_fn). The split is a pure function of the config, kimi_k3_module_fqns_per_model_part, that places the AttnRes aggregation on the head's stage and the tower on the embedding's; the entry hands pipeline_llm a replaced parallelism config instead of mutating the user's. The _return_only_new_blocks probe was left over from the experiment model and warned "naive (full-stack) mode" on every K3 run while delta mode was on; delta mode is the layout tables' presence, so the probe goes. (cherry picked from commit 7dda3b8)
The adapter required n_layers to divide by the stage count and both first/last_stage_less_layers to be 0, because a rank sees only its own stages under Interleaved1F1B and the layout inference verified the contiguous equal-split default instead of learning the real map. The map is now one all-gather over the pipeline group of each rank's (layer id -> stage) entries (gather_layer_to_stage), validated for completeness and contiguity, and handed to BlockLayoutTables, which already took an explicit map. Uneven stages, and the 93-layer model over any stage count, route their deltas from the real block boundaries. (cherry picked from commit ca5f34e)
…s attend Review suggestion: the first layer of a block appends the incoming stream to the block stack at the top of the layer, and both attention residuals then read the same shape -- the stack, plus the open block's partial sum when the layer is not the block's first. _apply_attention_residual takes that partial as Optional. The values every softmax sees are the same tensors in the same order as before, so the forward is unchanged bit for bit (checked at dp1 and pp2 x vp2 on a shared compile cache); a first layer builds its stack once instead of twice. (cherry picked from commit eef340d)
… is done with them The cache dropped everything at step end, so the blocks of every micro-batch in flight stayed resident for the whole step. No forward on a rank reads the cache for a micro-batch after the rank's last virtual stage has run, and the backward never reads it (across ranks the gradient rides the autograd graph and the schedule's backward P2P; on a rank it goes through the captured-grad slots), so the adapter now releases the blocks right after that forward and keeps only the slots until the producer's backward pops them. The step-end sweep stays as the safety net. This is the "released as soon as the micro-batch finishes" of the report; the cache holds the micro-batches between their first and last virtual stage on the rank instead of all of them. (cherry picked from commit d72faf3)
Review request. The debug flavor took default arguments and a second,
32-layer flavor existed only because 24 and 12 do not tile the pipeline
shapes. Now there is one debug model: _debugmodel requires its depth, the
MLA layers come from the same rule as the full model
(kimi_k3_full_attention_layers: every fourth layer and the last, which
reproduces range(3, 92, 4) | {92} at 93), and attn_res_block_size is the
model's 12. The depth is 30: the last block is partial (12 + 12 + 6, as the
full model's 7 x 12 + 9), the stack ends on a lone MLA layer after seven
groups, and with the embedding and the head counted as a layer each the 32
units divide into every pipeline shape up to 32 stages, uneven stages
included. The pp8 x vp4 flavor runs it at one layer per stage.
(cherry picked from commit 395fc6b)
torch.distributed.pipelining sizes a stage's gradient receive buffer from the
strides of the next stage's input gradients, computed once with
torch.autograd.grad at metadata inference, and c10d refuses a buffer that is
not non-overlapping and dense. A stage whose first use of an input is a
concatenation gets a view of the concatenation's gradient for that input:
the head's final aggregation on a stage that holds no transformer layer,
which is what one layer per stage over 32 stages produces on this model, and
a block-opening layer. That failed pp8 x vp4 at its first backward receive
("Tensors for P2P must be non-overlapping and dense"), with the transport on
or off. The model now passes its stage inputs through an identity whose
backward returns a contiguous gradient; forward and values are unchanged.
(cherry picked from commit c3df748)
The cross-stage cache adapter wrapped each stage's module and reached the schedule's state through wrappers: a patch on forward_one_chunk and backward_one_chunk with a thread-local for the micro-batch id, a patch on schedule.step for eviction, a tensor grad hook and an autograd Function for the same-rank gradient bridge, a zero-valued touch to keep a received carrier on the graph, and a reshape at shape inference. Every one of them was a seam between a linear stage chain and a model whose block is read by every later stage. AttnResPipelineStage, a PipelineStage subclass, replaces all of it. The hop between stages still carries (hidden, delta). forward_one_chunk assembles the full stack from the rank's store and the received delta, runs the model, and routes the payload; backward_one_chunk hands the gradient of the received columns back as the delta's gradient and deposits the gradient of the stored columns in the store; _retrieve_recv_grads collects the deposits for the blocks the stage committed, and the stage that received a block collects them for that block, before their own backward, which the schedule orders after every later stage's. The chunk id comes with the calls, the blocks of a micro-batch are released after the rank's last stage forward, the routing tables (BlockLayoutTables, now keyed by stage id and built from the schedule's own stage-to-rank map) say how many deposits each block must have, and metadata inference runs the same assembly, returning dense gradients so the previous stage's receive buffers are dense. The model neither returns a delta nor needs the _DenseGradient boundary any more. pipeline_llm takes a stage_class, the one generic hook this needs. With attn_res_cache off the same stage runs the whole stack on every hop, so the two transports differ only in the routing tables. (cherry picked from commit 3af70c9)
pyrefly on main's types: the model forward is declared to return the (hidden, block stack) pair a non-head stage hands on, the schedule's stages are narrowed to the subclass once, and the tests assert the optionals they index.
…yload's gradient is missing, the stage-count source named, comments trimmed
…pipelining entry back in parallelize.py, comments to one line pipeline_llm keeps only the stage_class parameter (no docstring, the file has none) and the pyrefly suppression the hook had dropped; the pipelining entry, the split and the stage lookup live in parallelize.py like every other model's, layout.py and pipeline_stage.py stay; the full-attention helper, the debug registry comment and the residual docstring are one line or upstream's own.
The stage count is the multiple of pp nearest to units / layers_per_stage, so a layer count no shape divides (the 93-layer model's 95 units) still splits into pp x vp stages that differ by a layer; core's ceiling would refuse the shape, so core gets the split alone and not the knob.
…lock Two blocks of 12 and a partial block of 9, as 93 = 7 x 12 + 9; with the embedding and the head 35 units, which no pipeline shape divides, so every split the matrix runs is uneven the way the real one is. 30 layers were 32 units, and 32 divides everything.
…s it With one stage per rank the delta transport degenerates to the whole stack -- no rank ever re-sees a block -- so the plain pp2 recipe exercised none of it.
Whether a layer opens a block is a property of the layer's index and the block size, both fixed at construction; the forward read it back from them on every call.
…can be asked for directly
module_fqns_per_model_part, pipeline_parallel_layers_per_stage and the new
pipeline_parallel_virtual_stages_per_rank all describe the same thing, so
__post_init__ takes at most one of them.
The new field exists because the stage count was reachable only as
ceil(units / layers_per_stage), which skips values a looped schedule may
need: a 33-layer model is 35 units, and that expression takes 35, 18, 12, 9,
7, 6, 5, 4, 3, 2, 1 -- four stages per rank on eight ranks is not among them,
although the split generator produces 32 stages happily. Stating the count
per rank rather than in total keeps it a multiple of the pipeline degree by
construction and matches the vocabulary the schedule checks already use
("stages per rank"); it must be at least 1, and the resulting stage count is
checked against the model's unit count where the model is known.
Core's split places the embedding, the layers and the head. A model whose output is aggregated next to the head has nowhere to say so: the aggregation modules have to land on the stage that holds lm_head. last_stage_modules names them, for the generator and for pipeline_llm; the first stage already has pipeline_with_first_stage_modules, which now forwards the same argument because it hands pipeline_llm a split that is already spelled out. The pinned names count as no layer, so the default -- no pinned modules -- returns what it returned before: checked against the previous function over 11232 shapes (1 to 32 stages, 1 to 39 layers, both weights 0 to 2), every split identical. _get_pipeline_metadata honours pipeline_parallel_virtual_stages_per_rank: the stage count stated per rank, validated against the model's unit count and against what the schedule accepts per rank. Adding the exclusivity check the review asked for in ParallelismConfig surfaced that pipeline_with_first_stage_modules spells a split out while leaving pipeline_parallel_layers_per_stage set, so the knob is silently ignored from that point on -- pipeline_llm prefers the explicit split. The function now clears it in the same replace, which says so and keeps the config's invariant. Kimi K3 pins the vision tower through that entry and its AttnRes aggregation through last_stage_modules, so the model no longer carries a split function of its own. The pp8 x vp4 recipe drops the split it used to spell out by hand -- a spelled-out split silences the pinned modules, which left the vision tower on no stage at all -- and asks for four stages per rank instead.
The split core applies is a pure function of the config -- the FQNs given, or the ones generated from the same metadata with the same pinned modules -- so every rank computes it and reads the layer-to-stage map off it. The all-gather over the pipeline group and the local-stage walk go.
…cks stay on the device One instance per rank, shared by the rank's stages; the name says so. The docstring states that blocks are device-resident and when they are released.
flatten_args has a detach=True overload that returns a tuple; spelling the forward cache's inputs as lists keeps the checker on the list overload. The decoder forward's override ignore names the code the checker reports.
…hared flavor The "debugmodel" entry is what every K3 cell builds -- the multimodal b200 cell, the tensor-parallel cells, whatever a sibling change adds next -- so deepening it to 33 layers made all of them train a deeper model on account of pipeline parallelism alone. It goes back to the 24 layers main has, and the deep shape becomes a flavor of its own that only the pipeline stress recipe asks for, so the depth travels with that recipe. The 33-layer shape itself is unchanged and still the point of the stress cell: two blocks of 12 and the 93-layer model's partial block of 9, 35 units with the embedding and the head, which no pipeline shape divides. A second, cheap pipeline cell runs on the shared model: pp2 x vp2 on two GPUs. 26 units over four stages is uneven, a stage boundary falls inside a block, and two stages per rank is the smallest shape where a rank receives a block it already holds, so the rank cache and the gradient deposits are exercised -- none of which a plain pp2 cell reaches. A unit test pins all three depths.
…try point Kimi K3 needs the split twice: once for core to build the stages from, and once to read the layer-to-stage map off, which is what replaced the all-gather. It was getting the second copy by importing two private functions and rerunning core's derivation, which hard-codes that derivation in a model file and is the cross-file private call the review objected to earlier. llm_split_with_pinned_modules is that derivation as a public entry point. It returns the split with the caller's modules pinned to both ends, and the parallelism config that spells it out: the fields that would otherwise derive a split are cleared there, since they have been read by then and pipeline_llm prefers the explicit split. That keeps the config's one-field invariant for every caller by construction rather than by each remembering to clear the right fields -- pipeline_with_first_stage_modules cleared pipeline_parallel_layers_per_stage and would have had to learn about pipeline_parallel_virtual_stages_per_rank as well. pipeline_with_first_stage_modules keeps its behaviour, shares that code instead of owning a copy, and gains last_stage_module_fqns for symmetry; pipeline_llm's last_stage_modules parameter, added earlier in this branch, has no caller left and goes. Kimi K3 imports one public name, and the split exists in its hands exactly once, so no rank can derive a different one from its peers.
…t derived it pipeline_with_first_stage_modules puts the generated split on the config with dataclasses.replace, which re-runs ParallelismConfig.__post_init__. The test asks for the stage count the legal way, with pipeline_parallel_layers_per_stage, and checks that the config handed to pipeline_llm carries the split and neither knob that derives one. Without the clear, the exclusivity check refuses it.
…own fields Core's pipelining constructs plain PipelineStages again; stage_class leaves pipeline_llm and _pipeline_module_split, so stage construction is main's. After pipeline_llm returns, K3 rebuilds each stage the schedule holds on this rank as an AttnResPipelineStage from the constructed stage's fields -- the same parallelized module, index, stage count, device, group, dw_builder and get_mesh callback -- carries over the stage-to-rank map the schedule wrote, and puts the rebuilt stages back in the schedule's _stage or _stages. Building a stage only reads its group (the per-direction P2P groups are cached per group), so nothing collective runs twice. A CPU test covers the swap for a single-stage and a multi-stage schedule.
The features suite runs its real-PG cells on every pull request on 8 x A10G, where Attention Gym's KDA kernel refuses to run. Kimi K3's other integration cell already lives in the B200 suite for that reason; the two pipeline cells and their recipes move there too.
…'s split pipeline_parallel_virtual_stages_per_rank served one test cell: the pp8 x vp4 recipe, whose 32 stages over 35 units no layers_per_stage reaches. That recipe now hands core's own split for that count to module_fqns_per_model_part, with the vision tower and the AttnRes aggregation where Kimi K3's entry pins them, so the command line gains no option. The pp2 x vp2 recipe needs nothing: looped schedules default to two stages per rank. __post_init__ keeps the check for the two knobs that already existed, and _get_pipeline_metadata is main's again.
…line or two The design argument in pipeline_kimi_k3's docstring and its two comment blocks, and in the two stage-swap helpers, stays in the PR text; the stray blank lines among the imports go. No behaviour change.
… does The module docstrings of pipeline_stage.py and layout.py, the rank cache, stack assembly, deposit-count, table-building and split-reading docstrings, and the pinned-module docstrings in pipeline_parallel.py keep one or two lines of what each does; the design argument stays in the PR text. Docstrings only: the code is unchanged.
…l.py keeps upstream's shape Kimi K3 builds its split bottom up in parallelize.py -- the vision tower with the embedding, the AttnRes aggregation with the head -- hands it to pipeline_llm through module_fqns_per_model_part, and reads the layer-to-stage map off the same list. llm_split_with_pinned_modules, the last-stage parameter of the LLM split generator and pipeline_with_first_stage_modules' last-stage argument are gone. The one core change left clears pipeline_parallel_layers_per_stage where pipeline_with_first_stage_modules spells a split out, which the one-way split check in ParallelismConfig requires. The pp8 x vp4 recipe spells out the same Kimi K3 split. Over 10368 configurations (layer count, pp degree, schedule, layers_per_stage, first / last stage weights, with and without the tower) the split, the cleared knob and every rejection match the previous code exactly.
…'s shape The one-way check in ParallelismConfig.__post_init__ (an explicit split and layers_per_stage not both) is dropped from this PR, and with it the line it forced into pipeline_with_first_stage_modules: on main that entry writes the split while leaving pipeline_parallel_layers_per_stage set, so the check rejected a valid config for every model going through it. Kimi K3's entry no longer clears the knob either; pipeline_llm ignores it once module_fqns_per_model_part is given. The split itself is unchanged (the same 10368 configurations, identical splits and rejections).
…together The one-way check lives in Trainer.Config.__post_init__, next to the num_pp_microbatches check, rather than in ParallelismConfig.__post_init__: dataclasses.replace re-runs the latter, and on main pipeline_with_first_stage_modules writes its generated split into the config while leaving pipeline_parallel_layers_per_stage set, so a check there rejects a valid config for every model going through that entry unless pipeline_parallel.py changes too. The trainer config is built once from the user's settings, so the check sees what the user gave.
…mConfig.__post_init__ An explicit module_fqns_per_model_part and pipeline_parallel_layers_per_stage are refused together in ParallelismConfig.__post_init__ again, instead of Trainer.Config. dataclasses.replace re-runs that check, and main's pipeline_with_first_stage_modules writes its generated split while keeping pipeline_parallel_layers_per_stage, so that entry clears the knob in the same replace (one line); otherwise every model on it would raise whenever layers_per_stage is set. Kimi K3's entry clears it the same way.
A 16-layer model with 4-layer blocks, as 16 and 8 Interleaved1F1B stages on 4 gloo ranks, whose block gradients are small integers: every sum is exact in bf16 and fp32. The rank cache, the whole-stack transport and a single device then agree bitwise on every block gradient and micro-batch loss for three SGD steps, and the gradients equal their closed form.
Adds an uneven 8-stage split whose stages open blocks after their first layer. The test now passes ufmt and pyrefly (the schedule takes a list typed to the base stage class, the last stage passes target and losses explicitly), forces the CPU device through a device_type property as the other CPU DTensor tests do, and ends with unittest.main().
… pinned host memory Between the stage that commits a block and the rank's later stages that read it the block sits on pinned host memory; the copies run on the current stream, so stream order serializes them, and every stored block is detached (its gradient travels through the deposits), so the result is bitwise the on-device store. A recipe turns it on with functools.partial(pipeline_kimi_k3, attn_res_cache_offload=True), like the transport switch. (cherry picked from commit 240c320)
QIU023
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
September 17, 2026 09:17
meta-cla
Bot
added
the
CLA Signed
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
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
1 participant
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.
Summary
Add
attn_res_cache_offloadtopipeline_kimi_k3: with it on, the blocks a rank keeps for its later stages (RankStore) are parked on pinned host memory from the stage that commits them until a later stage on the rank reads them, instead of staying on the device for that whole span.RankStore(offload=...)(kimi_k3/pipeline_stage.py):putcopies each block to a pinned host tensor on the current stream,blockscopies it back to the device it saw when a later stage assembles its stack.pipeline_kimi_k3(attn_res_cache_offload=...)(kimi_k3/parallelize.py): the switch, passed the way the transport switch is,functools.partial(pipeline_kimi_k3, attn_res_cache_offload=True)as thepipelining_fnof a recipe.Design
Every stored block is detached (its gradient travels through the store's deposits), so the round trip changes values nowhere. Both copies are
non_blockingon the current stream, so stream order alone serializes the D2H of the commit and the H2D of the read; nothing changes in the stage,_assemblestill checks the held set against the routing tables and stacks whatblocks()returns. A store on CPU has nowhere to park and stores the tensor as it is.The switch lives on
pipeline_kimi_k3next toattn_res_cachebecause it is a property of the rank store the pipelining entry builds, not of the model.Test plan
pytest tests/unit_tests/cpu/test_kimi_k3_pp_stage.py tests/unit_tests/cpu/test_kimi_k3_pp_layout.py tests/unit_tests/cpu/test_kimi_k3_stage_swap.py tests/unit_tests/cpu/test_kimi_k3_pp_exact_block_grads.py -q(17 passed)