[Kimi K3] Pipeline Parallelism support (text LLM side): PP Virtual Stage Cache Adapter for Attention Residual by QIU023 · Pull Request #4312 · pytorch/torchtitan

QIU023 · 2026-08-25T01:46:45Z

Summary

Adds pipeline parallelism to the Kimi K3 text decoder. Before this change parallelize.py rejects pipeline_parallel_degree > 1.

Core's pipeline_llm splits the model at layer boundaries and carries one hidden state tensor per hop, which cannot express Block Attention Residuals: every later stage needs every earlier block's residual, and the final aggregation (output_res_proj, then output_res_norm) must run only on the stage that owns lm_head.

After it, pipeline_kimi_k3 (kimi_k3/pipeline_parallel/__init__.py) takes core's split through pipeline_with_first_last_stage_modules, which now pins modules to the last stage as well as the first: the vision tower with the embedding, the aggregation with the head.

The schedule is built on AttnResPipelineStage, a torch.distributed.pipelining.PipelineStage subclass. A hop carries (hidden, delta), delta being the block residuals the receiving rank has not seen yet; each rank keeps the blocks it has seen in one store shared by its virtual stages, and the backward returns every block's gradient along the same routes.

Step 1 is identical to the single GPU reference in every cell. With the total grad norm taken in fp32, the whole stack pipeline cells (pp2, pp4, pp2 / pp4 x vp2 / vp4 with the cache off) stay identical to the reference for 100 steps in the bf16 cells, the loss on every step and the grad norm on every step but one last-digit print in two of the six.

In bf16 that norm, and the clip factor it sets every step at this scale, depend on how the pipeline groups the parameters (pytorch/pytorch#194033).

The cached cells differ only in the order in which a cached block's gradient contributions are added. A CPU test with integer valued block gradients shows the routing itself exact, bitwise against a single device closed form, cache on and off, over three optimizer steps.

Design

What forces the protocol: every layer attends over all earlier blocks plus the running partial block. So (R1) the block stack must cross every stage boundary with the hidden state, (R2) the final aggregation runs only on the stage that owns lm_head, and (R3) the stack grows with depth, with a boundary inside a block putting a partial block on the wire.

Sending the whole stack every hop satisfies all three and costs bytes that grow with the stage index. The delta transport sends only what the receiver lacks, which needs the routing tables, a micro batch key that survives P2P, and a way home for the gradient of a block a stage read from its own rank.

Why the rank store is enough: the schedule assigns stages $S = v \cdot P + R$, so a micro batch returns to the same rank every $P$ stages, and that rank already holds every block committed at stages $\le S-P$.

A freshly committed block is therefore new on the wire for $P-1$ hops and no longer. The same rank consumers of a block are exactly the stages congruent to its producer modulo $P$, and their number is what the tables expect as gradient deposits.

The stage protocol, in the subclass (kimi_k3/pipeline_parallel/stage.py)
- forward_one_chunk assembles the full block stack the model expects from the rank's store plus the received delta, runs the stage, keeps the blocks the stage committed, and sends on only what the next rank lacks. The model takes and returns the whole stack and knows nothing of the transport; the chunk id comes with the call.
- backward_one_chunk reads the gradient of the assembled stack, an autograd leaf the stage owns: the received columns go back as the delta's gradient, dense and in wire order, the stored columns are deposited in the rank store. The leaf is detached, so whether the previous stage wants a delta gradient is read off the receive metadata, which the stage checks describes (hidden, delta).
- The stage that committed or received a block collects those deposits into its own incoming gradient (_retrieve_recv_grads) before its backward. The store names no schedule class: it relies only on orderings every schedule gives by data dependency, a rank's stages running forward in increasing stage order and backward in decreasing, and the deposit count check fails loudly if one is missing or extra.
- A micro batch's blocks are released after the rank's last stage forward for it, so the store holds only the in flight micro batches.
- Forward only passes: schedule.eval, which the validator runs, calls backward_one_chunk with the backward disabled; the stage returns and drops the forward's bookkeeping instead of reading a gradient cache it never filled. A payload whose blocks have nothing trainable upstream (a frozen embedding under LoRA) gets no gradient, and its deposits are discarded.
- Metadata inference runs the same assembly (_compute_outputs); _compute_input_grads returns dense gradients, which is where the P2P buffer finding below is handled.
The routing tables (kimi_k3/pipeline_parallel/layout.py)
- BlockLayoutTables simulates one micro batch's forward in stage order over the split the trainer actually applied and tabulates, per stage, the blocks it commits, the blocks its rank already holds, and the blocks its P2P must carry; sender and receiver compute the same tables, so nothing but the delta travels.
- The layer to stage map is read off the split, which every rank computes, with no collective; the stage to rank map is the schedule's own stage_index_to_group_rank. Uneven stages are allowed; a block boundary inside a stage is a partial block on the wire.
- Why the delta is bounded: with $P$ ranks a block committed at stage $S$ is fresh on the wire for $P-1$ hops; from $S+P$ on every receiving rank already holds it, because its previous virtual stage was $S-P$. The per hop payload is bounded by the commits of the last $P-1$ stages, independent of depth.
- attn_res_cache=False, a parameter of pipeline_kimi_k3, sends the whole stack on every hop; the two transports differ only in the tables, which makes them the A/B in the results. With the cache on, pipeline_kimi_k3 checks the schedule's stage-to-rank map for the loop-style assignment the rank store assumes (stage s on rank s % pp) and refuses any other, naming the schedule and the first stage off it; v-shaped schedules are out of scope.
- Plain 1F1B is the naive transport by construction, so the two are bitwise there; with more than one stage per rank they are not, since the cached path sums the same gradient contributions in a different association.
The split is core's. pipeline_with_first_last_stage_modules takes last_stage_module_fqns and appends the ones the model has to the last stage, the mirror of what it does for the first; Kimi K3 pins the vision tower to the first stage, output_res_proj and output_res_norm to the last.
- return_split=True hands the applied split back to a model that routes along it, so Kimi K3 reads it off the call rather than deriving it again; pipeline_llm and the trainer's contract are untouched, and the five upstream models that hand the helper to the trainer as their pipelining_fn only follow the rename.
The stages: core's pipeline_llm constructs plain PipelineStages, as on main; K3 rebuilds each one the schedule holds as an AttnResPipelineStage from the constructed stage's own fields and puts it back in the schedule.
- The seam that is missing is a stage_class argument on pipeline_llm; until it exists the rebuild reads the schedule's _stage / _stages and the stage's mesh callback, and the stage imports flatten_args from pipelining._utils.
The model (model.py): the first layer of a block joins the stack before its sub layers attend, so a stage boundary at a block start needs nothing special and the stack a stage receives is exactly the stack the layers read; the head owning stage alone runs the aggregation.
Core: module_fqns_per_model_part becomes pipeline_parallel_module_fqns_per_model_part, the prefix every other pipeline field carries.
- The helper refuses layers_per_stage at its top, since the split it derives is sized by the schedule's default stage count; the copy it hands pipeline_llm carries that split and nothing else of the user's config changes. A different stage count is spelled out, as the pp4 x vp4 recipe does.
The shared debugmodel flavor is 17 layers, four runs of (3 KDA + 1 MLA) and a closing MLA as the 93 layer model has, in blocks of 4 (four full blocks and a partial one): dim 256 with 4 heads, 8 experts with top 2, a 2 layer vision tower, vocabulary 2048 against the 2020 token test tokenizer, 25.8 M parameters from 1.4 B.
- Why 17: nine layers already have every path the transport takes, and 17 is the least depth in that pattern whose 19 units split into the 16 stages of pp4 x vp4, which the results below run.
- The CI cell is pp4 x vp4 on its own recipe (kimi_k3_debugmodel_pp4_vp4, torchtitan_recipes/tests/b200.py), typechecking off as in every pipeline recipe upstream: sixteen stages, one layer per stage from layer 5 on, the head alone on the last, the second block opening inside stage 2; every path of the transport and every gradient deposit on four GPUs.
- The pp2 x vp2, pp4 x vp2 and pp2 x vp4 rows of the results run from local flavors (core's default split for the vp2 shapes, an 8 stage split spelled out for pp2 x vp4); nothing of them is in the PR.
What this replaced: the first reviewed version carried the same protocol in a 1228 line adapter that wrapped forward_one_chunk, backward_one_chunk and step, kept a thread local micro batch id, and bridged the same rank gradient path with a tensor grad hook and an autograd Function. The subclass implements it once, on the stage's own methods.

Results

4 x H100, kimi_k3_debugmodel (17 layers, dim 256) reading c4_test as text-only 256-token rows, four 256-token micro-batches per rank, one seed checkpoint per batch shape, total grad norm in fp32. Each cell gives the raw value and, beneath it, the change against the reference. Measured at ce67cece4 on main 7349a2282 (the two commits after it change a guard's form and an error's type, not a computation).

One compile lineage per table: the reference runs once cold to fill the inductor and Triton caches and once more on a copy of both, and that second run is the reference; every other cell runs on a copy of the same two caches. The bf16 cold fill is the last row of the first table: on this box it reproduces the warm rerun bitwise, as the fp32 one does, while the dp2 lineage's cold fill picks other kernels and differs from step 2, which is why every cell runs on the warm reference's cache.

None of the following is part of this PR:

the c4 flavors (kimi_k3_debugmodel_c4, _pp_naive, _8stages, _16stages, the last the CI recipe's split), the switches ATTN_RES_NAIVE (after Refactor model configuration and parallelization ownership #4810 the model owns its pipelining, so a switch stands in for the removed pipelining_fn partial), NOSYNC_GA, MB_REVERSE, GN_FP32, and the KDA capability guard lifted for SM90: probe_apply_h100.py
the whole matrix, and the seed checkpoint and cache handling above: run_pp_h100.sh

python probe_apply_h100.py . && export GN_FP32=1   # c4 flavors, the switches, the total grad norm in fp32
COMMON="-m torchtitan.train --module kimi_k3 --debug.seed 42 --debug.deterministic --metrics.log_freq 1 --training.num-tokens-per-train-step 1024 --training.num-tokens-per-microbatch-per-dp-rank 256 --parallelism.data_parallel_shard_degree 1"
torchrun --nproc_per_node=1 $COMMON --config kimi_k3_debugmodel_c4_seed --training.steps 1 --dump-folder seed
cell() { d=$1; n=$2; c=$3; src=$4; shift 4; rm -rf $d; mkdir -p $d; cp -r seed/checkpoint $d/; [ -d cache_$src ] && cp -r cache_$src cache_$d; TORCHINDUCTOR_CACHE_DIR=cache_$d/inductor TRITON_CACHE_DIR=cache_$d/triton torchrun --nproc_per_node=$n $COMMON --config $c --training.steps 100 "$@" --dump-folder $d; }
P2="--parallelism.pipeline_parallel_degree 2 --parallelism.num-pp-microbatches 4"; P4="--parallelism.pipeline_parallel_degree 4 --parallelism.num-pp-microbatches 4"; IL="--parallelism.pipeline_parallel_schedule Interleaved1F1B"
NOSYNC_GA=1 cell ref_cold 1 kimi_k3_debugmodel_c4 none; NOSYNC_GA=1 cell ref 1 kimi_k3_debugmodel_c4 ref_cold   # the second run is the reference
cell dp1 1 kimi_k3_debugmodel_c4 ref; MB_REVERSE=1 cell reversed 1 kimi_k3_debugmodel_c4 ref
cell pp2 2 kimi_k3_debugmodel_c4 ref $P2; cell pp4 4 kimi_k3_debugmodel_c4 ref $P4
cell vp2_cached 2 kimi_k3_debugmodel_c4 ref $P2 $IL; cell vp2_naive 2 kimi_k3_debugmodel_c4_pp_naive ref $P2 $IL   # pp4 x vp2 likewise with $P4
cell pp2vp4_cached 2 kimi_k3_debugmodel_c4_8stages ref $P2 $IL; cell pp4vp4_cached 4 kimi_k3_debugmodel_c4_16stages ref $P4 $IL   # local flavors with the splits spelled out; _naive variants likewise
# dp2 table: --parallelism.data_parallel_shard_degree 2 --training.num-tokens-per-train-step 2048, its own seed checkpoint, reference and cache

cell	loss, step 1	step 10	step 50	step 100	grad norm, step 1	step 10	step 50	step 100
dp1 [1]	`8.003520`	`3.488880`	`2.584620`	`2.511720`	`2.2337`	`1.3804`	`0.6772`	`0.7861`
dp1, stock accumulation [4]	`8.003520` identical	`3.485760` -0.09%	`2.588440` +0.15%	`2.511260` -0.02%	`2.2337` identical	`1.3846` +0.30%	`0.6855` +1.23%	`0.7832` -0.37%
dp1, micro-batches reversed [5]	`8.003520` identical	`3.485890` -0.09%	`2.587830` +0.12%	`2.509630` -0.08%	`2.2338` +0.00%	`1.3743` -0.44%	`0.6857` +1.26%	`0.7884` +0.29%
pp2	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical
pp4	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical
pp2 x vp2, naive [2]	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical
pp2 x vp2, cached	`8.003520` identical	`3.484850` -0.12%	`2.586430` +0.07%	`2.510560` -0.05%	`2.2337` identical	`1.3904` +0.72%	`0.6782` +0.15%	`0.7908` +0.60%
pp2 x vp4, naive [2][3]	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical
pp2 x vp4, cached [3]	`8.003520` identical	`3.484350` -0.13%	`2.585380` +0.03%	`2.509430` -0.09%	`2.2337` identical	`1.3920` +0.84%	`0.6719` -0.78%	`0.7819` -0.53%
pp4 x vp2, naive [2]	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical
pp4 x vp2, cached	`8.003520` identical	`3.487100` -0.05%	`2.584950` +0.01%	`2.508500` -0.13%	`2.2337` identical	`1.3682` -0.88%	`0.6786` +0.21%	`0.7855` -0.08%
pp4 x vp4, naive [2][3]	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical
pp4 x vp4, cached [3]	`8.003520` identical	`3.486570` -0.07%	`2.585890` +0.05%	`2.509130` -0.10%	`2.2339` +0.01%	`1.3861` +0.41%	`0.6705` -0.99%	`0.7889` +0.36%
dp1, second cold compile [6]	`8.003520` identical	`3.488880` identical	`2.584620` identical	`2.511720` identical	`2.2337` identical	`1.3804` identical	`0.6772` identical	`0.7861` identical

[1] reference: the micro-batches accumulate with the gradient sync on the last one, as the pipeline does (NOSYNC_GA); the logged loss is the sum of the four micro-batch losses in one reduction, as the pipeline's last stage logs it
[2] naive transport, the whole block stack on every hop (attn_res_cache=False); "cached" rows use the rank store, the default
[3] 8 and 16 stages with the split spelled out, since no layers_per_stage reaches them for 19 units: the 16 stage split is the CI recipe's, the 8 stage one a local flavor's (probe_apply.py above)
[4] stock accumulation: the gradient sync on every micro-batch (the flavor's default), the same data and cache
[5] the four micro-batches in reverse order, the same data and cache: the floor of a bf16 sum order change
[6] the same code, data and seed compiled a second time from an empty cache: on this box it picks the same kernels and is identical to the reference on all 100 steps
pp2, pp4 and the four naive rows are identical to the reference on all 100 steps in the loss; in the grad norm pp4 and three naive rows on all 100 steps, pp2 and pp2 x vp4 naive on 99, with step 37 printing 1.2084 against 1.2085, one unit in the last digit: under PP core's clip_grad_norm_ squares each rank's norm, all-reduces and takes the root, another summation order than one device's single vector_norm, and the bf16 parameters absorb it
the fp32 end to end comparison the review thread asked for, cache on and off, is in that thread

dp2, 2048 tokens per step, same protocol.

cell	loss, step 1	step 10	step 50	step 100	grad norm, step 1	step 10	step 50	step 100
dp2 [1]	`8.059690`	`3.411970`	`2.565870`	`2.366000`	`2.0111`	`1.4686`	`0.5743`	`0.6493`
dp2 x pp2	`8.059690` identical	`3.411970` identical	`2.565870` identical	`2.366000` identical	`2.0111` identical	`1.4686` identical	`0.5743` identical	`0.6493` identical
dp2 x pp2 x vp2, naive [2]	`8.059690` identical	`3.411970` identical	`2.565870` identical	`2.366000` identical	`2.0111` identical	`1.4686` identical	`0.5743` identical	`0.6493` identical
dp2 x pp2 x vp2, cached	`8.059690` identical	`3.412200` +0.01%	`2.564320` -0.06%	`2.368620` +0.11%	`2.0111` identical	`1.4784` +0.67%	`0.5711` -0.56%	`0.6455` -0.59%

[1] [2] as above; dp2 x pp2 and the naive row are identical to the reference on all 100 steps, loss and grad norm. This lineage's cold fill differs from its warm rerun from step 2 (autotune), so every row runs on the warm reference's cache

The KDA capability guard was widened locally to admit SM 9.0 for these runs; it is not part of this PR.

Test plan

CPU, in the default unit test suite: test_kimi_k3_pp_block_grads.py (the real AttnResPipelineStage under ScheduleInterleaved1F1B on four gloo ranks with integer block gradients, bitwise against cache off and a single device closed form over three steps, bf16 and fp32, a forward only pass between steps), test_kimi_k3_pp_layout.py (the routing tables on uneven splits, the 8 and 16 stage splits, the loop-style guard), test_kimi_k3_pp_stage.py (assembly, routing, the gradient split, the forward only return, the stage rebuild on one gloo rank), test_pipeline_parallel.py (the helper: both ends pinned, the split handed back, the caller's config untouched), test_config_manager.py (the renamed field), test_integration_test_definitions.py (the two B200 cells): pytest tests/unit_tests/cpu -k "kimi_k3 or pipeline or cli or integration_test or config_manager" -q, 107 passed.
B200 suite: kimi_k3_pp4_vp4 (four GPUs, sixteen stages, one layer per stage from layer 5 on, the head alone on the last stage, the second block opening inside stage 2); kimi_k3_mm unchanged.
The tables above: 100 steps per cell on 4 x H100, every cell against a single device reference on the same seed checkpoint and compile lineage.

A `torch.distributed.pipelining` finding

With one layer per stage the last stage holds only the head, whose first op on the block stack is a cat, so autograd hands the stage's input gradients back as views.

PipelineStage._backward_metadata_inference records those strides, _create_grad_recv_info allocates the receive buffer with torch.empty_strided, and c10d rejects it at the first RECV_B with "Tensors for P2P must be non-overlapping and dense".

Every other split passed because a later op consumed the input and autograd accumulated a dense gradient. The subclass returns dense gradients from _compute_input_grads; the library-side fix would be a dense torch.empty receive buffer and .contiguous() before the send.

QIU023 · 2026-08-26T05:47:59Z

cc @shuhuayu @tianyu-l Kimi K3 text side PP PR ready for review, please let me know if the above PR summary might be too verbose or if you have any confusion / question and I will answer ASAP!

tianyu-l

sorry but to be honest, I couldn't understand the code in pipeline_adapter.py. My high level feelings are

too many functions with "install", "patch", "hook", "thread", which feels temporary unblocking rather than "first-principle" infra, but I could be wrong
If skip connections across transformer blocks is becoming popular, we should look at proper ways to build general infra in torch.distributed.pipelining.
numerics feels wrong, see comment #4312 (comment)

I would suggest

Could you write some notes about "how you get there", instead of simply presenting what's ready. E.g. start from scratch -> what's missing, so you added X component; o/w Y is not feasible. This is a complicated piece of infra, so I'd like to understand the decision making.
Maybe if you are available, you could give us a talk and walk over the design? That's totally optional.

tianyu-l · 2026-09-03T08:20:40Z

+# ----- Rank-shared cache across virtual stages ----------------------------- #
+
+
+class RankLocalCache:


@sanketpurandare you mentioned that pytorch PP today already has caching

from my observation: PipelineStage.fwd_cache / bwd_cache hold a stage's own inputs and outputs for its own backward;

But they cannot serve a later stage on the same rank, so the rank store stays; what is shared with them is the chunk id, which the subclass now gets from forward_one_chunk directly.

Please correct me if this assumption is wrong, thanks for both of your raising concerns and answering!

tianyu-l · 2026-09-03T08:44:39Z

    def forward(
        self,
        x_TD: torch.Tensor,
        block_residual_TND: torch.Tensor,


Kimi tech report:

The AttnRes computation is entirely wrapped with checkpointing, so the activation saved for the backward pass at each layer is identical to that of the standard residual architecture

What's the storage life cycle of these block tensors -- e.g. when PP=1, would we keep N tensors of size 1, 2, 3, 4, 5, ..., 93 on dim-1 at the end of forward which grows quadratically?

Quadratic by construction, but in the block count, not the layer count: the stack grows once per 12-layer block, so at the end of forward the graph holds stacks of 1..8 for the 93-layer model, N(N+1)/2 = 36 [T, D] rows, not 1 2 3 ... 93

The larger term is the per-layer AttnRes internals in fp32, O(n) rows per layer; the report's "entirely wrapped with checkpointing" moves exactly that from storage to backward recomputation, so the saved activations match the standard residual net

the report's "entirely wrapped with checkpointing" moves exactly that from storage to backward recomputation, so the saved activations match the standard residual net,

(PS. has implemented as its own follow-up PR on our fork (_apply_attention_residual becomes a checkpoint(use_reentrant=False) entry; without it a no-AC run OOMs at step 2 even at debug scale on a 16 GiB card, with it the same config runs 10 steps with step-1 loss bitwise), kept out of this diff for simplicity)

Quadratic by construction, but in the block count, not the layer count: the stack grows once per 12-layer block, so at the end of forward the graph holds stacks of 1..8 for the 93-layer model, N(N+1)/2 = 36 [T, D] rows, not 1 2 3 ... 93

It seems the answer is that "yes we would save all previous blocks' result for each layer, despite the redundancy" and the overhead is limited. Is this correct?

But I didn't understand the per-layer internals part. Could you be very specific about the activation save / recomputation policy you / Kimi are using? If possible, please work with simple examples to illustrate the idea. E.g. every block has 4 layers, there are two blocks -- in this case, what exactly are you saving in forward, after each layer, and when exactly would things be released? Thanks!

NP, will address tomorrow asap, doing the numerical proof/verification now and log off very soon

Hi Tianyu @tianyu-l , here is the per-stage view, where the cache matters.

Your example, 2 blocks x 4 layers, with pp2 x vp2 (Interleaved1F1B: global stage s runs on rank s % 2) and the split core generates. e is the embedding (stack entry 0), x4 block 1's result (entry 1); block 2's result is consumed by the output aggregation and never enters the stack.

global stage (rank, virtual) layers Block AttnRes stored stack PP Comm, fwd (no cache) PP Comm, fwd (cache) already cached

s0 (r0, v0) emb, 0-1 e - - -

s1 (r1, v0) 2-4 x4 s0 -> [e] s0 -> [e] -

s2 (r0, v1) 5-6 - s1 -> [e, x4] s1 -> [x4] [e]

s3 (r1, v1) 7, head - s2 -> [e, x4] s2 -> [] (empty payload) [e, x4]

Stack entries sent per micro-batch in forward: 5 with the cache off, 2 with it on. The hidden state travels on every hop and is not listed; an empty entry list still posts a 0-byte send / recv.

Backward, s3 -> s0 for each micro-batch:

global stage (rank, virtual) PP Comm, bwd (no cache) PP Comm, bwd (cache) deposits into the rank cache collects from the rank cache

s3 (r1, v1) - - [de, dx4] -

s2 (r0, v1) s3 -> [de, dx4] no op (empty payload) [de] -

s1 (r1, v0) s2 -> [de, dx4] s2 -> [dx4] - s3 stored: [de, dx4]

s0 (r0, v0) s1 -> [de] s1 -> [de] - s2 stored: [de]

Gradient entries sent per micro-batch in backward: 5 with the cache off, 2 with it on. The hidden state's gradient travels on every hop and is not listed; an empty payload needs no gradient, so nothing is sent for it.

One micro-batch of that example with the cache on.

Forward (the hidden state travels on every hop and is not drawn):

sequenceDiagram participant C0 as rank 0 cache participant S0 as s0 (r0) participant S1 as s1 (r1) participant S2 as s2 (r0) participant S3 as s3 (r1) participant C1 as rank 1 cache S0->>C0: put e (committed here) S0->>S1: P2P [e] S1->>C1: put e (received) and x4 (committed here) S1->>S2: P2P [x4] C0-->>S2: read e S2->>C0: put x4 (received) Note over C0: s2 is rank 0's last stage: drop this micro-batch's blocks S2->>S3: P2P [] (0-byte payload) C1-->>S3: read e and x4 Note over C1: s3 is rank 1's last stage: drop this micro-batch's blocks

Loading

Backward (each stage frees its own saved stack and activations when its backward of the micro-batch finishes):

sequenceDiagram participant C0 as rank 0 cache participant S0 as s0 (r0) participant S1 as s1 (r1) participant S2 as s2 (r0) participant S3 as s3 (r1) participant C1 as rank 1 cache S3->>C1: deposit de and dx4 (read from the cache) Note over S2,S3: no block gradient on this hop (the forward payload was empty) S2->>C0: deposit de (read from the cache) S2->>S1: P2P [dx4] C1-->>S1: collect dx4 into the incoming dx4, then collect de after its own backward S1->>S0: P2P [de] C0-->>S0: collect de into the incoming de, then run its own backward Note over C0,C1: every deposit collected, the rank checks none is left

Loading

In conclusion:

Forward: a stage puts each entry it adds or receives into its rank's cache (a detached view of the stack, no copy). The rank drops a micro-batch's entries right after its last stage's forward for that micro-batch (stage 2 on rank 0, stage 3 on rank 1). What backward needs is held by autograd as usual, not by the cache: each stage keeps its own stack (assembled from the payload and the cache) and its saved activations in the pipeline's fwd_cache, and frees them in its own backward of that micro-batch (s3 first, s0 last), the same as with the cache off.

Backward: a stage that read an entry from the cache does not send that entry's gradient anywhere; it deposits it in the rank's cache (stage 3: e and x4; stage 2: e). The stage that put the entry on the rank adds the deposit to its own gradient before sending upstream (stage 1 collects both on rank 1, stage 0 collects e on rank 0). So a deposit lives from the reader's backward to the owner's backward of the same micro-batch, and the rank checks none is left after its first stage's backward.

Inside a stage, what is saved and recomputed is the same as without PP: with the default SelectiveAC each layer is one checkpoint region, and the residual reads are recomputed in that layer's backward.

A hop carries only the entries the receiving rank has not seen yet. The same with 4 blocks x 4 layers, pp4 x vp4 (stage s on rank s % 4; x4, x8, x12 = the results of blocks 1-3):

global stage (rank, virtual) layers Block AttnRes stored stack PP Comm, fwd (no cache) PP Comm, fwd (cache) already cached

s0 (r0, v0) emb, 0 e - - -

s1 (r1, v0) 1-2 - s0 -> [e] s0 -> [e] -

s2 (r2, v0) 3 - s1 -> [e] s1 -> [e] -

s3 (r3, v0) 4 x4 s2 -> [e] s2 -> [e] -

s4 (r0, v1) 5 - s3 -> [e, x4] s3 -> [x4] [e]

s5 (r1, v1) 6 - s4 -> [e, x4] s4 -> [x4] [e]

s6 (r2, v1) 7 - s5 -> [e, x4] s5 -> [x4] [e]

s7 (r3, v1) 8 x8 s6 -> [e, x4] s6 -> [] (empty payload) [e, x4]

s8 (r0, v2) 9 - s7 -> [e, x4, x8] s7 -> [x8] [e, x4]

s9 (r1, v2) 10 - s8 -> [e, x4, x8] s8 -> [x8] [e, x4]

s10 (r2, v2) 11 - s9 -> [e, x4, x8] s9 -> [x8] [e, x4]

s11 (r3, v2) 12 x12 s10 -> [e, x4, x8] s10 -> [] (empty payload) [e, x4, x8]

s12 (r0, v3) 13 - s11 -> [e, x4, x8, x12] s11 -> [x12] [e, x4, x8]

s13 (r1, v3) 14 - s12 -> [e, x4, x8, x12] s12 -> [x12] [e, x4, x8]

s14 (r2, v3) 15 - s13 -> [e, x4, x8, x12] s13 -> [x12] [e, x4, x8]

s15 (r3, v3) head - s14 -> [e, x4, x8, x12] s14 -> [] (empty payload) [e, x4, x8, x12]

Stack entries sent per micro-batch in forward: 39 with the cache off, 12 with it on. The hidden state travels on every hop and is not listed; an empty entry list still posts a 0-byte send / recv.

Backward, s15 -> s0 for each micro-batch:

global stage (rank, virtual) PP Comm, bwd (no cache) PP Comm, bwd (cache) deposits into the rank cache collects from the rank cache

s15 (r3, v3) - - [de, dx4, dx8, dx12] -

s14 (r2, v3) s15 -> [de, dx4, dx8, dx12] no op (empty payload) [de, dx4, dx8] -

s13 (r1, v3) s14 -> [de, dx4, dx8, dx12] s14 -> [dx12] [de, dx4, dx8] -

s12 (r0, v3) s13 -> [de, dx4, dx8, dx12] s13 -> [dx12] [de, dx4, dx8] -

s11 (r3, v2) s12 -> [de, dx4, dx8, dx12] s12 -> [dx12] [de, dx4, dx8] s15 stored: [dx12]

s10 (r2, v2) s11 -> [de, dx4, dx8] no op (empty payload) [de, dx4] s14 stored: [dx8]

s9 (r1, v2) s10 -> [de, dx4, dx8] s10 -> [dx8] [de, dx4] s13 stored: [dx8]

s8 (r0, v2) s9 -> [de, dx4, dx8] s9 -> [dx8] [de, dx4] s12 stored: [dx8]

s7 (r3, v1) s8 -> [de, dx4, dx8] s8 -> [dx8] [de, dx4] s11 stored: [dx8], s15 stored: [dx8]

s6 (r2, v1) s7 -> [de, dx4] no op (empty payload) [de] s10 stored: [dx4], s14 stored: [dx4]

s5 (r1, v1) s6 -> [de, dx4] s6 -> [dx4] [de] s9 stored: [dx4], s13 stored: [dx4]

s4 (r0, v1) s5 -> [de, dx4] s5 -> [dx4] [de] s8 stored: [dx4], s12 stored: [dx4]

s3 (r3, v0) s4 -> [de, dx4] s4 -> [dx4] - s7 stored: [de, dx4], s11 stored: [de, dx4], s15 stored: [de, dx4]

s2 (r2, v0) s3 -> [de] s3 -> [de] - s6 stored: [de], s10 stored: [de], s14 stored: [de]

s1 (r1, v0) s2 -> [de] s2 -> [de] - s5 stored: [de], s9 stored: [de], s13 stored: [de]

s0 (r0, v0) s1 -> [de] s1 -> [de] - s4 stored: [de], s8 stored: [de], s12 stored: [de]

Gradient entries sent per micro-batch in backward: 39 with the cache off, 12 with it on. The hidden state's gradient travels on every hop and is not listed; an empty payload needs no gradient, so nothing is sent for it.

@tianyu-l A quick follow up for optimizing our original naive apply Attention Residual:

What's the storage life cycle of these block tensors -- e.g. when PP=1, would we keep N tensors of size 1, 2, 3, 4, 5, ..., 93 on dim-1 at the end of forward which grows quadratically?

Since the second half of the question now has a change behind it. The stack is as described above: quadratic in the block count rather than the layer count, 1 to 8 for 93 layers at block size 12, so 36 [T, D] rows at the end of forward. The term that dominated is the other one, the per-layer fp32 internals: _apply_attention_residual kept two fp32 [T, N + 1, D] tensors per call alive until backward, which is exactly what made the saved set unlike a standard residual block.

Raised #4780 that replaces that body with an torch.autograd.Function that saves only per-token statistics and recomputes the upcasts in backward. It is filed against main rather than here because the aggregation arrived with #4025 and its three call sites carry no pipeline guard, so this diff does not move.

Debug model saved: so one call at the released hidden size retains 0.2 MiB instead of 1008.2 and its forward peak is 280 MiB instead of 1764, more test data in #4780

sorry, I'm asking that without PP, what is the activation life cycle of blocks and the partial sums, with 2 blocks and each block having 4 layers.

Citing Claude Bot's comment below:

at 2 blocks x 4 layers and no PP, the end of the forward holds three stack columns, one [T, D] partial sum per layer, and, unless AC recomputes them, the per layer fp32 residual internals, which are the cost that matters.

Conclusion:

A column enters the stack at its block's first layer (layer 0 and layer 4) and is read by every later layer and by the head. Under SelectiveAC, the flavor's default, each layer's checkpoint holds the stack it received, so three columns are alive at the end of the forward, each freed once the first layer that read it has run its backward. With AC off no op saves the older stack, so only the current one is alive.

A partial sum lives one layer: the next layer's residual read consumes it and it is freed after that layer's backward, as a plain residual's hidden state is. It never crosses a block boundary.

The extra term is each residual read's fp32 intermediates.

[Kimi K3] Checkpoint the attention-residual computation #4656 checkpoints them, and under selective or full AC the layer's own checkpoint already does;

[Kimi K3] Attention residual: zero initialisation and a recomputing aggregation #4780 keeps them out of the saved set altogether.

QIU023 · 2026-09-03T09:24:24Z

Thank you so much for your comprehensive review Tianyu @tianyu-l ! I will address ASAP and provide a detailed design explanation, and verify the numerical correctness part!

This PR has not yet including the Attention Residual Activation Checkpointing reuse, PP Cache Adapter CPU offloading and the Mooncake cross PP-rank activation transferring for memory LB, I only have some drafted implementation for them on a integration branch but not yet able to cherry-pick them to this PR branch since comprehensive numerical test is needed before it. (And not sure if you are okay with adding Mooncake as external optional dependency for cross PP-rank memory LB, I guess not likely since we have detached from fla, similar for MoonEP case)

Regarding the road of the design, this current impl direction was coming from a few rounds of iteration while trying to implement following the Kimi original author of AttnRes training infra owner way back in April and encountered bunch of training gradient flow / map issues under the core-idea as caching older VP block activations and gradients, I will summarize it to present to you ASAP.

(Referred AttnRes Training Infra Mandarin Blog: https://www.zhihu.com/question/2016993095078684011/answer/2017381145474508331)

QIU023 · 2026-09-04T09:47:27Z

Have done a major refactor about pp_adapter.py for a cleaner design and arch, with addressing @tianyu-l 's comments and questions, numerical matrix is still being ran and will be updated to PR body within 1 day

pending numerical correctness verification, design motivation summarize, haven't know if and when convenient to present a design talk since workday daytime I am not convenient for it but I will try to provide it (perhaps at some night time or weekend ?) to speedup the review and get maintainers familiar with the design.

QIU023 · 2026-09-07T09:34:01Z

Hi Tianyu @tianyu-l , numerical correctness is verified and with some gap explained, most of the comments from Tianyu is addressed, just made a full design idea explanation and summarize of motivation and requested changes idea in torch.distributed.pipelining about cross-layer connection PP support

Please let me know if you have any questions or feedback reading this doc of PP AttnRes design idea/motivation/further dependency plan and I can make the readability better for any of your and other maintainers' need

PP_RUNTIME_DESIGN_NOTE_2026-09-08.en.short.pdf

QIU023 · 2026-09-08T22:49:04Z

Hi Tianyu @tianyu-l, I am really sorry for just realized the above placed md file do not render the descriptive figure I want to present to help explaining, also have saw the ongoing changes in #4486 that AttnRes into this updated file and I can provide some use case for the design there, cc @acisseJZhong

Could you please check the replaced to PDF file:
PP_RUNTIME_DESIGN_NOTE_2026-09-08.en.short.pdf

also please let me know which time you prefer to have the talk about the explanation for PP design idea and I will try my best to work it out

QIU023 · 2026-09-09T22:02:41Z

cc @elfiegg this is the new PP tree after the rebasing of #4025 and had done refactor of PP adapters removing hooks and patch for cross PP-rank caching,

let’s follow up offline with maintainers with the new fixes from you about the super helpful multi node PP fixes added

tianyu-l · 2026-09-10T08:07:33Z

Thanks. I haven't looked into details, but the refactor looks much cleaner, I'll come to details asap.

Without looking into your design / the zhihu blog, I had two options in my mind (after discussing with @sanketpurandare):

[option 1]
Do skip connections, where in initial stages
rank 0 sends rank 0 blocks to rank 1, 2, 3
rank 1 sends rank 1 blocks to rank 2, 3
rank 2 sends rank 2 blocks to rank 3

[option 2]
The alternative is
rank 0 sends rank 0 blocks to rank 1
rank 1 sends rank 0+1 blocks to rank 2
rank 2 sends rank 0+1+2 blocks to rank 3

The main benefit is that we don't need to invent skip connections for now. Message volume stays the same, but the above "rank 2 sending rank 0+1+2 blocks to rank 3" would be [PP degree] times more data enlarging the bubble (hopefully only in the first microbatch). But it's not [num transformer layers] times more.

On the other hand, in addition to the complexity, option 1 "rank 0 sends rank 0 blocks to rank 1, 2, 3" may cause some contention for every microbatch? I'm not sure.

After reading the zhihu blog, it seems Kimi and this PR is doing [option 1], which I think makes the most sense, at least in the short term.

And I noticed that you mentioned [option 2] in your design pdf, which could be a future exploration direction.

btw I don't think we'll go with the design in #4486, at least not for MTP + PP support.

Yes, thanks for your confirming of the [option 1] as ideal way. Actually to further correct, this PR is sending P2P Comm amount of Block Activations in [option 1] but the cross-PP rank comm topo is [option 2] like;

on the real 93-layer shape: 8 blocks, one message = one block, per micro-batch.

shape option 1 option 2 this PR this PR, peak per hop option 2, peak

pp8 x vp1 28 28 28 7 7

pp4 x vp2 18 28 18 3 7

pp2 x vp4 7 28 7 1 7

pp8 x vp4 52 136 52 2 8

pp4 x vp8 24 136 24 1 8

pp16 x vp2 96 136 96 4 8

The volume column for this PR equals option 1's in every shape, and every message is still between adjacent stages: a block is relayed forward, and a hop carries only what the receiving rank does not already hold. A rank keeps what it has seen for its own later virtual stages, which is what makes the relay cost the same as the direct sends -- no block is ever delivered twice to the same rank.

The regular [option 1] (No impl through cache in this PR) will do:
sends from a producer rank to ranks it is not adjacent to, which means P2P outside the schedule's own batches and traffic on the pipeline group that the schedule does not order, reached in different orders by neighbouring ranks under 1F1B.

This PR [Rank comm topo as option 2, but cache most of old Blocks] impl cache avoids it by construction -- the block rides the forward hop that has to happen anyway, so it adds no message and no latency of its own, and the stage issues no collective and no object P2P.

rank 0 sends 0 to rank 1 rank 1 sends 0+1 to rank 2 rank 2 sends 0+1+2 to rank 3 rank 3 sends 1+2+3 to rank 0 (rank 0 already has 0) rank 0 sends 2+3+4 to rank 1 (rank 1 already has 0+1) rank 1 sends 3+4+5 to rank 2 (rank 2 already has 0+1+2) rank 2 sends 4+5+6 to rank 3 (rank 3 already has 0+1+2+3)

(Refer this map again)

Sorry if there is any mis-understanding, but [Option 2] is disabling the entire PP cache and will lead to the whole PPxVP is non-optimized (later PP rank has to send larger amount of Block Attn activations, congesting the cross PP-rank / DC rack-level networking comm since it's a way lower bandwidth. Hence we can keep it but if this PR goes well I hope it can be enabled as default

tianyu-l · 2026-09-10T08:13:30Z

    def forward(
        self,
        x_TD: torch.Tensor,
        block_residual_TND: torch.Tensor,


Quadratic by construction, but in the block count, not the layer count: the stack grows once per 12-layer block, so at the end of forward the graph holds stacks of 1..8 for the 93-layer model, N(N+1)/2 = 36 [T, D] rows, not 1 2 3 ... 93

It seems the answer is that "yes we would save all previous blocks' result for each layer, despite the redundancy" and the overhead is limited. Is this correct?

But I didn't understand the per-layer internals part. Could you be very specific about the activation save / recomputation policy you / Kimi are using? If possible, please work with simple examples to illustrate the idea. E.g. every block has 4 layers, there are two blocks -- in this case, what exactly are you saving in forward, after each layer, and when exactly would things be released? Thanks!

QIU023 · 2026-09-10T09:06:06Z

Thank you so much Tianyu! I will answer your above questions and reply discussions tomorrow night as top-1 priority on this one and waiting for your detailed comments in the pipeline_stage.py

tianyu-l · 2026-09-11T00:37:50Z

+    # 33 layers: two blocks of 12 and the 93-layer model's partial block of 9
+    # (93 = 7 x 12 + 9); 35 units with the embedding and the head, which no
+    # pipeline shape divides, so every split is uneven the way the real one is.
+    "debugmodel": (partial(_debugmodel, num_layers=33), 16384),


33 might be too large for a debugmodel, how many total params does it have?

33 layers = 1.40 B parameters, of which the embedding and the head are 335.5 M (the vocabulary is the real one) and each layer 32.2 M;

in the first round review, you would want to have a non-even splitting model and following the (3KDA+1MLA)*8+1MLA (33 layers, 3 Block AttnRes) pattern, to ensure we can have the pressure PP8xVP4 cache [option 1] test covered, the model need to have at least 32 layers, therefore I introduced this 33 layers debugmodel to both unify and cover all required scenarios

therefore I would suggest to keep original debugmodel (24 layers, evenly splitted) and enable the PP2xVP2 for regular CI testing, and optional for the 33 layer covering uneven splitting and provided to PP8xVP4 cache pressure testing coverage (can be removed if you do not want 2 flavors, but will miss uneven split and higher-pressume PP VP testing); either way works for me! Will change follow your decision next round

the embedding and the head are 335.5 M (the vocabulary is the real one)

shouldn't use full vocab in debug model, please check what other modes do, e.g. qwen

I think we should have small number of layers, 24 still sounds too many. What is the minimum layer number you think would be enough to test VPP > 1? We can also shrink each layers by tuning the other hyperparameters.

Followed Qwen3, vocabulary is 2048 now, and the widths are cut to debug scale: dim 256, 4 heads, latent 128, 8 experts with top 2, a 2 layer vision tower, reduced to 25.8 M parameters

VPP > 1 with every path of the block transport needs two blocks, one opening inside a stage, plus the partial closing block, which the pattern gives at 9 layers. 17 is the least depth in that pattern whose 19 units split into the 16 stages of pp4 x vp4. Would feel this is more well coverage than pp2 x vp2

tianyu-l · 2026-09-11T00:58:01Z

+            kwargs["parallelism"] = dataclasses.replace(
+                parallelism,
+                module_fqns_per_model_part=fqns,
+                pipeline_parallel_layers_per_stage=None,


I think we anyway won't use it, so it sounds unnecessary to set this to None explicitly here https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/pipeline_parallel.py#L89

That said, we probably should error out in post_init of ParallelismConfig when both are given.

Ack, fixed: ParallelismConfig.__post_init__ now raises when both module_fqns_per_model_part and pipeline_parallel_layers_per_stage are given: diff

But dataclasses.replace re-runs __post_init__, and pipeline_with_first_stage_modules keeps layers_per_stage when it writes its split, so the new check would reject the split if the knob is still set, which made us still put this "set to None" here

tianyu-l · 2026-09-11T00:59:02Z

+    # Core spells the head ``output``; this model calls it ``lm_head``. Any
+    # FQN matching no child makes core set that child to None on every stage.
+    fqns = [["lm_head" if n == "output" else n for n in stage] for stage in fqns]
+    tail = [n for n in _KIMI_ATTN_RES_LAST_STAGE_FQNS if hasattr(model, n)]


I think we should extend https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/pipeline_parallel.py#L144 instead of making this model specific

Just found that #4560 has merged and introduced pipeline_with_first_stage_modules(), so have to rebase to build on top of it and make similar changes, for last_stage_modules

please LMK if you prefer to review some post latest-rebase changes or retain old head and rebase after approval

please rebase and extend

Have rebased onto current main and extended as asked: the helper takes last_stage_module_fqns and is pipeline_with_first_last_stage_modules now, and return_split=True hands the split it applied back.

tianyu-l · 2026-09-11T02:12:41Z

+        n_layers=n_layers,
+        layers_per_block=layers_per_block,
+        layer_to_stage=layer_to_stage,
+        cache=attn_res_cache,


numerical correctness is verified and with some gap explained

I'm not convinced that the numerics gap can be this large.

Could you give simple examples of how accum order would change with this flag off / on vs. without PP?

How do you prove it's not caused by bugs?

Thanks for sticking on this criteria and reminding!

have done extensive research and reverification on this to answer these, will provide concrete answers and a report tomorrow ASAP from the AI summarized doc and experiment results to ensure addressing your concern

PS. Kimi infra author mentioned the numerical will be harder to align in the Zhihu blog on PP adapter cache and different rank cached activations / gradients are reduced in different orders comparing to naive impl (above [option 2]), and our this PR was done on the debug model with untrained fixed random init ckpt, might be harder to present a converged difference

Against our bar -- step-1 loss bitwise, step-1 gradients bitwise or located -- the pipeline passes: step 1 matches a single GPU in every cell below, and the one step-1 gradient difference left is located in the last layer's attention backward, before anything crosses a stage boundary.

Moved these runs to H100. KDA runs Attention Gym kernels outside the SM100/SM103 path main supports, with configurations autotuned per process, so how much this debug flavour amplifies ulp-level differences by step 10 depends on the device: the same pp2 cell reads +13.7% at step 10 on RTX 5060 Ti and +3.6% on H100.

Accumulation order, flag off / on vs no PP.

Take one block of the residual stack under pp2 x vp2 (rank 0 runs stages 0 and 2, rank 1 stages 1 and 3). Without PP, autograd adds every layer's read of that block onto one running gradient, top-down. With the cache off, each hop hands that running gradient back and the previous stage keeps adding onto it, so the order is the same as without PP. With the cache on, stages 2 and 3 read the block from their rank's store, sum their own reads from zero and deposit the subtotal, and stages 1 and 0 add it when they collect: the same terms, grouped differently -- bitwise-different in bf16 where terms cancel, identical in float64.

How we know it is not a bug.

Comparing every parameter's step-1 gradient, with the predictions written down before the dumps were read: the cache changes only the parameters that produce a block read from a rank's store, by 2-3 bf16 ulps, and leaves everything downstream bitwise; deleting one gradient deposit, a real bug in that path, moves the same tensors about a hundred times further.

4 x H100 PCIe, one seed checkpoint, 100 steps; 1024 tokens per step because four stages need four 256-token micro-batches; steps stop at 20 because the reference memorises the debug set after that. Percentages against the first row; the last row has no pipeline in it.

The PP2 slot seems to be still causing by #4135 and pytorch #194033 , got comments from Jane and will raise revision tomorrow

cell loss step 1 loss step 10 loss step 20 grad norm step 1 grad norm step 10 grad norm step 20

dp1 12.605700 3.114620 3.373330 18.625 5.4375 3.9844

pp2 12.605700
bitwise 3.227050
+3.61% 3.288290
-2.52% 18.75
+0.67% 5.6875
+4.60% 3.7344
-6.27%

pp2 x vp2, cache on 12.605700
bitwise 3.150940
+1.17% 3.349300
-0.71% 18.625
0% 3.7188
-31.61% 4.0625
+1.96%

pp2 x vp2, cache off 12.605700
bitwise 3.514970
+12.85% 3.281700
-2.72% 18.625
0% 6.0312
+10.92% 3.6719
-7.84%

dp1, accumulation order reversed 12.605700
bitwise 3.247610
+4.27% 3.295370
-2.31% 18.625
0% 6.6562
+22.41% 4.25
+6.67%

dp2, 2048 tokens per step; raw value and, beneath it, the change against dp2; the last row has no pipeline in it.

cell loss step 1 loss step 10 loss step 20 grad norm step 1 grad norm step 10 grad norm step 20

dp2 12.521140 3.221120 2.839810 16.375 7.0625 2.4844

dp2 x pp2 12.521140
same 3.201920
-0.60% 2.846440
+0.23% 16.375
0% 5
-29.20% 2.6562
+6.92%

dp2 x pp2 x vp2, cache on 12.521140
same 3.205680
-0.48% 2.682970
-5.52% 16.375
0% 5.375
-23.89% 2.5156
+1.26%

dp2 x pp2 x vp2, cache off 12.521140
same 3.189240
-0.99% 2.847220
+0.26% 16.375
0% 5.5938
-20.80% 2.375
-4.40%

dp2 x ep2 12.521140
same 3.149310
-2.23% 2.969330
+4.56% 16.375
0% 5.0625
-28.32% 2.9531
+18.87%

Very detailed generated report on numerical analysis on both H100 and 5060Ti

REPLY_4312_NUMERICS_2026-09-11.md

Hi Tianyu @tianyu-l , sorry if above numerical checking message thread might be too much dumping, in conclusion to make the results more cleaner, I aligned to the H100 device and apply the local diff of (active review fix,

above REPLY_4312_NUMERICS_2026-09-11.md file is for reference about the analysis (on both H100 and 5060Ti), the new H100 results on C4 dataset with 100 steps is better and cleaner for your review as below:

Re-run on the same 4 x H100, now on c4_test as text-only rows (at 1024 tokens per step the reference memorised the original 32-sample debug set within the 100 steps), with the total grad norm taken in float32: with bf16 gradients the logged norm, and the clip factor that applies it every step here, depend on how the pipeline splits the parameters (pytorch/pytorch#194033). Without the rank cache every pipeline cell is identical to the reference over 100 steps (the loss on every step, the logged grad norm on every step but one, where it differs in the fourth decimal); the cached cells differ because the cache adds a cached block's gradient contributions in another order.

4 x H100 PCIe, kimi_k3_debugmodel (24 layers), one seed checkpoint, four 256-token micro-batches per rank; the c4 flavor, the accumulation switch, the pp4 x vp4 stage count and the fp32 norm are local probe changes, not part of this PR (c4 patch, stage count, fp32 norm).

Reproduction, on k3_pp_text plus the three probe files linked above:

python gn_fp32_hack.py . && export GN_FP32=1 # total grad norm in fp32 export TORCHINDUCTOR_CACHE_DIR=$PWD/cache/inductor TRITON_CACHE_DIR=$PWD/cache/triton # one compile cache for every cell of a table, warmed by a 1-step run of each configuration first COMMON="-m torchtitan.train --module kimi_k3 --debug.seed 42 --debug.deterministic --training.num-tokens-per-train-step 1024 --training.num-tokens-per-microbatch-per-dp-rank 256 --checkpoint.enable --parallelism.data_parallel_shard_degree 1" torchrun --nproc_per_node=1 $COMMON --config kimi_k3_debugmodel_c4 --training.steps 1 --checkpoint.create_seed_checkpoint --dump-folder seed cell() { d=$1; n=$2; c=$3; shift 3; rm -rf $d; mkdir -p $d; cp -r seed/checkpoint $d/; torchrun --nproc_per_node=$n $COMMON --config $c --training.steps 100 --metrics.log_freq 1 --checkpoint.interval 100000 "$@" --dump-folder $d; } P="--parallelism.pipeline_parallel_degree 2 --parallelism.num-pp-microbatches 4"; IL="--parallelism.pipeline_parallel_schedule Interleaved1F1B"; P4="--parallelism.pipeline_parallel_degree 4 --parallelism.num-pp-microbatches 4 $IL" NOSYNC_GA=1 cell ref 1 kimi_k3_debugmodel_c4 cell pp2 2 kimi_k3_debugmodel_c4 $P; cell vp2_naive 2 kimi_k3_debugmodel_c4_pp_naive $P $IL; cell vp2_cached 2 kimi_k3_debugmodel_c4 $P $IL PP_STAGES_PER_RANK=4 cell pp4vp4_naive 4 kimi_k3_debugmodel_c4_pp_naive $P4; PP_STAGES_PER_RANK=4 cell pp4vp4_cached 4 kimi_k3_debugmodel_c4 $P4 # 2048-token table: --parallelism.data_parallel_shard_degree 2 --training.num-tokens-per-train-step 2048, its own seed checkpoint; NOSYNC_GA=1 for the reference

1024 tokens per step:

cell loss, step 1 step 10 step 20 step 100 grad norm, step 1 step 10 step 20 step 100

dp1 (reference) 12.609980 3.595760 2.998930 2.533960 16.9193 5.0429 2.1888 1.6014

pp2 (all 100 steps identical) 12.609980
identical 3.595760
identical 2.998930
identical 2.533960
identical 16.9193
identical 5.0429
identical 2.1888
identical 1.6014
identical

pp2 x vp2, naive 12.609980
identical 3.595760
identical 2.998930
identical 2.533960
identical 16.9193
identical 5.0429
identical 2.1888
identical 1.6014
identical

pp4 x vp4, naive 12.609980
identical 3.595760
identical 2.998930
identical 2.533960
identical 16.9193
identical 5.0429
identical 2.1888
identical 1.6014
identical

pp2 x vp2, cached 12.609980
identical 3.554240
-1.15% 3.055690
+1.89% 2.571310
+1.47% 16.9177
-0.01% 5.2631
+4.37% 2.9058
+32.76% 1.6073
+0.37%

pp4 x vp4, cached 12.609980
identical 3.309360
-7.96% 3.007760
+0.29% 2.554590
+0.81% 16.9167
-0.02% 3.5408
-29.79% 2.0286
-7.32% 1.5763
-1.57%

2048 tokens per step, dp2:

cell loss, step 1 step 10 step 20 step 100 grad norm, step 1 step 10 step 20 step 100

dp2 (reference) 12.580740 3.576250 2.976110 2.420420 14.4170 12.4659 2.4576 1.0313

dp2 x pp2 (all 100 steps identical) 12.580740
identical 3.576250
identical 2.976110
identical 2.420420
identical 14.4170
identical 12.4659
identical 2.4576
identical 1.0313
identical

dp2 x pp2 x vp2, naive (all 100 steps identical) 12.580740
identical 3.576250
identical 2.976110
identical 2.420420
identical 14.4170
identical 12.4659
identical 2.4576
identical 1.0313
identical

dp2 x pp2 x vp2, cached 12.580740
identical 3.305220
-7.58% 2.949200
-0.90% 2.434300
+0.57% 14.4191
+0.01% 4.5317
-63.65% 2.6572
+8.12% 1.0893
+5.62%

have concluded that the PP cache enabled numerical difference is only coming from backward reduction ordering difference between Block Activations, forward is 100% same between cache enabled and disabled (naive), here is another figure helping illustration,

and adding a new AttnResPipelineStage unit test: to make sure 100% correctness of pipeline_stage.py (numerical difference is only from bf16/fp32 Block Activation reduction order in real usage):

On CPU with 4 gloo ranks,

a linear 16-layer model with 4-layer blocks runs 3 SGD steps through the real AttnResPipelineStage and Interleaved1F1B, split into 16 and 8 stages;

every block-gradient contribution is a small integer, so every sum is exact, and the test checks that with the rank cache on, with it off,

and on a single device, every block gradient and micro-batch loss are bitwise equal at every step and equal their closed form, in bf16 and fp32.

Question 1:
When you compare DP and PP, what global batch size are you using in both? I think to mimic PP grad reduction order, we should compare gradient accumulation vs. PP, while keeping the number of microbatches the same. If you set reshard_after_forward to "always", then I expect both to be bitwise identical, with cache off.
https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/fsdp.py#L126

Question 2:

The PP2 slot seems to be still causing by #4135 and pytorch #194033 , got comments from Jane and will raise revision tomorrow

What dtype are you using? For this study, we at least should use fp32 in backward, with cache on.

Q1. Same global batch for both: 1024 tokens per step, four 256-token sequences, four micro-batches. The dp2 table is 2048 over two ranks.

Not bitwise, and the gap is core's. In fp32, gradient accumulation against pp2 and pp2 x vp2, cache off and reshard_after_forward="always", is identical through step 14 and drifts to -0.11% on the loss at step 100. At step 1 every gradient, updated parameter and micro-batch loss is bitwise; from step 2 only the total grad norm differs, in its last fp32 unit.

The unit comes from clip_grad_norm_: one device takes a single vector_norm over all parameters (utils.py#L594), under PP each rank's norm is squared, all-reduced and rooted (L609-L615), and the clip factor (L617) is live on the first 17 steps. With the per-parameter norms gathered over the pp ranks and reduced in the single-device order, the same cell is bitwise for 100 steps, norm bits included. reshard_after_forward changes nothing either way.

In bf16 the same unit sits in the clip factor on 7 of the first 20 steps and leaves the bf16-valued gradients and the master parameters bitwise (measured at steps 2 and 20), which is why the description's whole-stack rows are identical for 100 steps. Happy to file the clipping order separately.

Q2. The description's tables are bf16 parameters and compute with fp32 reduction. fp32 end to end, cache on and off, same protocol on 4 x H200:

cell loss, step 1 step 10 step 50 step 100 grad norm, step 1 step 10 step 50 step 100

fp32 dp1, reference 8.055300 3.405580 2.578440 2.501960 2.1551 1.2746 0.7113 0.8558

fp32 dp1, stock accumulation 8.055300
identical 3.405580
identical 2.578440
identical 2.501960
identical 2.1551
identical 1.2746
identical 0.7113
identical 0.8558
identical

fp32 pp2 x vp2, naive 8.055300
identical 3.405580
identical 2.576750
-0.07% 2.499220
-0.11% 2.1551
identical 1.2746
identical 0.7125
+0.17% 0.8608
+0.58%

fp32 pp2 x vp2, cached 8.055300
identical 3.405580
identical 2.576820
-0.06% 2.499740
-0.09% 2.1551
identical 1.2746
identical 0.7127
+0.20% 0.8613
+0.64%

fp32 pp4 x vp4, naive 8.055300
identical 3.405580
identical 2.576820
-0.06% 2.499740
-0.09% 2.1551
identical 1.2746
identical 0.7127
+0.20% 0.8613
+0.64%

fp32 pp4 x vp4, cached 8.055300
identical 3.405530
-0.00% 2.574910
-0.14% 2.501610
-0.01% 2.1551
identical 1.2744
-0.02% 0.7067
-0.65% 0.8648
+1.05%

Cache on tracks cache off through step 18 (pp2 x vp2) and step 8 (pp4 x vp4), within 0.1% at step 100. At step 1, 212 of the 459 gradients differ at the rounding level, none missing: the cache changes the order the contributions are added in, nothing else. The cache-off drift from step 15 is the clipping order above, not the transport.

QIU023 · 2026-09-12T10:03:01Z

Have addressed most of the comments and will double check tomorrow, then requesting your next round review, thanks a lot! @tianyu-l

tianyu-l · 2026-09-20T06:09:01Z

+    # 33 layers: two blocks of 12 and the 93-layer model's partial block of 9
+    # (93 = 7 x 12 + 9); 35 units with the embedding and the head, which no
+    # pipeline shape divides, so every split is uneven the way the real one is.
+    "debugmodel": (partial(_debugmodel, num_layers=33), 16384),


the embedding and the head are 335.5 M (the vocabulary is the real one)

shouldn't use full vocab in debug model, please check what other modes do, e.g. qwen

I think we should have small number of layers, 24 still sounds too many. What is the minimum layer number you think would be enough to test VPP > 1? We can also shrink each layers by tuning the other hyperparameters.

tianyu-l · 2026-09-20T06:09:24Z

+    # Core spells the head ``output``; this model calls it ``lm_head``. Any
+    # FQN matching no child makes core set that child to None on every stage.
+    fqns = [["lm_head" if n == "output" else n for n in stage] for stage in fqns]
+    tail = [n for n in _KIMI_ATTN_RES_LAST_STAGE_FQNS if hasattr(model, n)]


please rebase and extend

tianyu-l · 2026-09-20T06:29:14Z

+        n_layers=n_layers,
+        layers_per_block=layers_per_block,
+        layer_to_stage=layer_to_stage,
+        cache=attn_res_cache,


Question 1:
When you compare DP and PP, what global batch size are you using in both? I think to mimic PP grad reduction order, we should compare gradient accumulation vs. PP, while keeping the number of microbatches the same. If you set reshard_after_forward to "always", then I expect both to be bitwise identical, with cache off.
https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/fsdp.py#L126

Question 2:

The PP2 slot seems to be still causing by #4135 and pytorch #194033 , got comments from Jane and will raise revision tomorrow

What dtype are you using? For this study, we at least should use fp32 in backward, with cache on.

tianyu-l · 2026-09-20T06:31:31Z

    def forward(
        self,
        x_TD: torch.Tensor,
        block_residual_TND: torch.Tensor,


sorry, I'm asking that without PP, what is the activation life cycle of blocks and the partial sums, with 2 blocks and each block having 4 layers.

tianyu-l · 2026-09-20T06:39:54Z

@claude review this PR, in particular please answer my questions regarding

[Kimi K3] Pipeline Parallelism support (text LLM side): PP Virtual Stage Cache Adapter for Attention Residual #4312 (comment) memory lifecycle of activations on transformer block inputs, including block_residual_TND, when no PP is used
[Kimi K3] Pipeline Parallelism support (text LLM side): PP Virtual Stage Cache Adapter for Attention Residual #4312 (comment) review the gradient flow in PP implementation and check if there could be bugs

claude · 2026-09-20T06:40:12Z

Claude finished @tianyu-l's task in 23m 24s —— View job

Reviewing PP support for Kimi K3

Q1: Activation memory lifecycle of block inputs / block_residual_TND when no PP is used (2 blocks x 4 layers)
Q2: Review gradient flow in the PP implementation for possible bugs

Q1 - Activation lifecycle with no PP (2 blocks x 4 layers)

Reading model.py, here is exactly what KimiK3Model.forward keeps alive for the backward with block_size = 4, layers 0-7, no PP. The stack (block_residual_TND) is only grown at a block's first layer:

model.py:246  if self.first_layer_in_block:
                  block_residual_TND = torch.cat((block_residual_TND, x_TD.unsqueeze(1)), dim=1)

Forward trace (T tokens, D model dim):

layer	`first_layer_in_block`	`block_residual_TND` after the layer	what `_apply_attention_residual` reads
entry	-	`[T, 0, D]` (empty, `model.py:497`)	-
L0	yes (block 0 opens)	`[T, 1, D]` (embedding)	L0 skips it (`attention_res_proj is None`, `model.py:254`)
L1	no	`[T, 1, D]`	stack + partial `x1`
L2	no	`[T, 1, D]`	stack + partial `x2`
L3	no	`[T, 1, D]`	stack + partial `x3`
L4	yes (block 1 opens)	`[T, 2, D]`	stack of 2
L5-L7	no	`[T, 2, D]`	stack of 2 + partial
output agg	-	reads `[T, 2, D]` (`model.py:511`)	stack of 2

Three distinct classes of retained activation:

The block stack columns. Each torch.cat at a block boundary is a new storage, so the graph retains the [T, 1, D] from block 0 and the [T, 2, D] from block 1 - i.e. columns 1 + 2 = 3, which is N(N+1)/2 with N = 2. This generalizes to 36 columns for the 93-layer model (N = 8, block size 12), not 1 + 2 + ... + 93. This confirms your earlier answer to @tianyu-l: quadratic in the block count, small in absolute terms. These columns are never freed mid-forward - every later layer's attention residual plus the output aggregation reads all earlier blocks, so they live from creation until the backward reaches their region.
Per-layer partial sums (prefix_sum_TD, model.py:274/288). One [T, D] per layer, exactly a standard residual activation, freed as the backward passes each layer.
The _apply_attention_residual fp32 internals - the dominant term. Each call upcasts:
```
model.py:167  values_float = values_TND.float()          # fp32 [T, N(+1), D]
model.py:169  keys_TND = values_float * torch.rsqrt(...)  # fp32 [T, N(+1), D]
```
Both are needed by the matmul/rsqrt backward, so two fp32 [T, N+1, D] tensors per call are held until that layer's backward, and there are ~2 calls/layer (attention residual + ffn residual) plus the aggregation. This - not the stack - is what makes the saved set unlike a plain residual block.

Release / recomputation policy: with the default SelectiveAC each transformer layer is one checkpoint region, so the fp32 internals in (3) are recomputed in that layer's backward rather than held across the whole forward; the stack columns in (1) are the AC region's inputs and are retained. This is consistent with the Kimi report's "entirely wrapped with checkpointing," and PR #4780 (against main) is the targeted fix that turns _apply_attention_residual into an autograd.Function saving only per-token statistics.

So, directly answering the question: with 2 blocks x 4 layers and no PP, at the end of forward you hold 3 stack columns ([T,1,D] + [T,2,D]), one [T,D] partial sum per layer, and - unless AC recomputes them - the fp32 residual internals per layer, which are the real cost.

Q2 - Gradient flow in the PP implementation

I traced the forward/backward routing in pipeline_stage.py against the layout tables. The design is sound and, importantly, it is validated end-to-end - I want to lead with that, then flag the subtle points.

How the gradient is routed. The key move is in assemble_stack (pipeline_stage.py:86):

return stack_TND.detach().requires_grad_(True), order

The assembled stack is a fresh autograd leaf - the connection from the model's stack back to the received delta_TND and to the rank store is deliberately severed, and the gradient is re-routed by hand:

backward_one_chunk reads grad_stack (grad wrt the leaf), then split_stack_grad divides its columns into received blocks (-> grad_delta, in wire order, sent upstream) and stored blocks (-> store.deposit).
_retrieve_recv_grads (before this stage's own backward) folds in the deposits for blocks this stage committed and sent forward, so the model backward sees the full downstream gradient for those columns.
The count invariant deposits_expected (layout.py:78) asserts each block collects exactly one deposit per later same-rank reader.

Why the ordering is correct: Interleaved1F1B runs backward in reverse stage order, so the higher-indexed same-rank readers deposit before the lower-indexed owner collects. The _collect_into count check (pipeline_stage.py:276) would raise if any deposit were missing or extra - a genuine guardrail, not a silent path.

Strongest evidence it's correct: test_kimi_k3_pp_exact_block_grads.py runs the real AttnResPipelineStage through the real ScheduleInterleaved1F1B on 4 gloo ranks, with block gradients constructed as small integers so every sum is exact, and asserts bitwise equality (rtol=0, atol=0) between cache-on, cache-off, and single-device closed form - across 3 optimizer steps, both bf16 and fp32, for three splits including uneven ones and blocks opening mid-stage. This is exactly the right way to separate a routing/gradient bug (which would flip signs / change magnitudes) from floating-point reduction-order differences (which the integer construction eliminates). The cache-on/off numeric gap in the PR body is therefore reduction order only, and this test proves it.

Subtle points worth confirming (not bugs I can demonstrate, but the places I'd double-check):

delta_needs_grad depends on recv-tensor metadata, not on the (detached) graph. In backward_one_chunk:
```
pipeline_stage.py:319  inputs_meta = self._stage_meta.inputs
pipeline_stage.py:320  delta_needs_grad = inputs_meta is not None and inputs_meta[1].requires_grad
```
Because _compute_outputs also assembles via the detached leaf (pipeline_stage.py:348), the module output does not depend on delta_TND during metadata inference, so the framework's own autograd.grad wrt delta would be unused. The gradient is instead sent iff the received delta tensor carries requires_grad. This works (the exact test confirms), but it hinges on _stage_meta.inputs[1] always being the delta and always reflecting the sender's payload requiring grad. Since it is load-bearing and non-obvious, I'd add an assert that inputs_meta[1] corresponds to the delta (e.g. shape rank 3 / block dim) rather than trusting positional index [1], and a one-line comment on why the value is read from the recv metadata rather than the backward grad. Fix this →
_is_last_on_rank() releases the store on the last forward; _is_first_on_rank() asserts no deposits remain after the last backward. These correctly bound the store to in-flight micro-batches. Worth confirming that a schedule that does not order all of a rank's virtual-stage forwards before its backwards (i.e. anything other than the Interleaved family this PR targets) is refused upstream - the layout/collection invariants assume the S = v*P + R congruence. parallelize.py builds on whatever pipeline_parallel_schedule resolves to; a guard that the schedule is an interleaved multi-stage one (or a clear error if a rank holds one stage yet cache=True was requested with a non-congruent schedule) would make the assumption explicit.
_retrieve_recv_grads clones then mutates a view (pipeline_stage.py:267-269): grad_delta = grad_delta.clone() then _collect_into does grad_col_TD.add_(deposit) on grad_delta[:, j]. This is safe (fresh clone), just noting it's the intended in-place accumulation.

Nothing in the flow reads as a correctness bug to me - the detach-and-reroute is a legitimate pattern for skip connections that cross stage boundaries, the deposit/collect counts are checked, and the exact-integer test closes the loop. My only asks are the defensive assert + comment in point 1 and the explicit schedule guard in point 2, both to protect the non-obvious invariants.

Minor, aligned with @tianyu-l's open comments: the 163840 vocab in _debugmodel (__init__.py:486) drives the 335.5M embedding/head @tianyu-l flagged - other debug models (e.g. qwen3) use a small vocab; shrinking it and the layer count would make CI far cheaper without losing VPP>1 coverage. And the file moves he requested (layout.py/pipeline_stage.py -> kimi_k3/pipeline_parallel/) are still open.

QIU023 · 2026-09-20T07:46:37Z

@claude for the first question from Tianyu and your answer, double-check to confirm if using #4780 replacing the _apply_attention_residual() will change the conclusion, (it is expected to), regardless of PP enabling

pytorch-bot · 2026-09-21T06:26:21Z

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

ciflow/rl
ciflow/fake-pg

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.

QIU023 · 2026-09-21T09:31:01Z

Hi Tianyu @tianyu-l , thanks a lot for your comments! Have responded and rebased + changed file structures & extend needed utils as needed, all numerical results and additional fp32 PP cache on/off experiments results attached in above thread, please check, thanks!

(PS. touched the shared utils pipeline_with_first_stage_modules renamed to pipeline_with_first_last_stage_modules for making the module splitting logic common utils, so other relevant model files are touched with naming changes)

shuhuayu

lgtm overall, some more comments. in addition, maybe we can have one plot in folder torchtitan/models/kimi_k3/pipeline_parallel to illustrate the idea of pp for attention residual with cache, similar plots are like https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/MOE_SHARDING.md. cc: @tianyu-l

shuhuayu · 2026-09-23T00:08:16Z

+    return grad_delta, deposits
+
+
+# TODO: decide whether this stage belongs in torch.distributed.pipelining.


my comment would be not at least in the current form since there are many kimi k3 specific features.

ack, removed this comment for now, make sense,

Tianyu's previous thought is only referring to Cross-Layer connection becomes popular and this VPP cache for cross layer connection can be the fundamental infra (not the case now)

shuhuayu · 2026-09-23T00:52:05Z

    return config
+
+
+def kimi_k3_debugmodel_pp4_vp4() -> Trainer.Config:


let's test some composibility. with 8 gpus, we can test fsdp=2, ep=2, tp=2, pp=2. keep that spmd type checking is disabled for pp.

changed recipe to kimi_k3_debugmodel_fsdp2_tp2_ep2_pp2,

please LMK if you are okay or prefer to drop kimi_k3_debugmodel_pp4_vp4 and I will quickly remove this recipe in next commit,

I prefer we keep the debug model to 16 layers and I think pp4vp4 (16 stages) can test more on the caching behavior and numerical difference, pp2vp2 CI might not be complex enough to cover the breaking AttnRes cache change, happy to follow whatever you think is correct here, thanks!

Umm, just realize that perhaps we can set in this composition cell to be PP2 VP4 instead, that would be better coverages for the cache on case,

and in such case the above separated PP4 VP4 cell can be removed, LMK the preferred remaining debug PP setups and will change accordingly

…_model_part The field names the split that pipeline parallelism applies, so it takes the prefix its neighbours carry, and it is exclusive with pipeline_parallel_layers_per_stage: both describe the same split and a run that sets each would have to pick one silently.

The helper already co-located modules with the first stage; a model whose output side has modules outside the decoder needs the same for the last one, so it takes last_stage_module_fqns and a model declares both ends through pipeline_first_stage_module_fqns and pipeline_last_stage_module_fqns. It refuses pipeline_parallel_layers_per_stage instead of overriding it: the split it derives is sized by the schedule's default stage count, so the two cannot both hold. return_split hands the applied split back to a caller that routes along it.

…el owns it The model declares the modules that belong with the first and the last stage and overrides pipeline() to route the attention-residual blocks along the split it gets back. The package holds the stage subclass, whose forward carries the blocks a receiving rank lacks and whose backward reads the delta's gradient need off the receive metadata, and the routing tables, which derive the block count from the layer count and assume the loop-style stage-to-rank assignment the contiguity check enforces. A block's first layer closes the previous block and joins the stack, so a stage that opens inside a block carries a partial sum rather than a closed block, and a stage without the head hands the stack on unchanged.

…abulary The flavor carried the released vocabulary and 24 layers at model widths, which is neither cheap nor a debug shape. 17 is the least depth in the model's own (3 KDA + 1 MLA) pattern whose units split into the 16 stages a pp4 x vp4 cell asks for.

One recipe and one cell, as the llama3 pipeline cells are, with the split spelled out because the debug model's units do not divide into 16 stages. Type checking is off, which the pipeline recipes upstream also do.

…2 x pp2 One eight-GPU cell with the four parallelisms together, as the reviewer asked for: FSDP 2, TP 2 with sequence parallel, EP 2 and PP 2 under 1F1B with four micro-batches. Type checking is off, which the pipeline recipes have in common.

… sharding One note in the pipeline package with one vector figure: a stage on one rank with the rank store, forward and backward, and the same stage with the whole-stack transport, with h, B and delta and their gradients on the arrows and the words kept for a legend. The rest of the note is the two transports side by side and what the stage and the tables do, so a reader has the idea before the code.

shuhuayu · 2026-09-23T09:00:57Z

+        layout, store = self.layout(), self.store()
+        deposit, count = store.collect(mb, b)
+        expected = layout.deposits_expected(b, self.stage_index)
+        if count != expected:


this is expecting every stage must use the cached blocks on it, so it is required to contribute a grad. so there is a splitting constraint that each stage must have a transformer layer, which is satisfied by the titan stage splitting, but not true in general for any stage splitting. non-blocking for this pr.

ack, thanks a lot for reminding here ! Will leave a comment in next revision if any comments need to be addressed, or next PR for K3 after this merged, eg. K3 MoonViT PP DEP

QIU023 requested review from fegin, tianyu-l, wconstab and wwwjn as code owners August 25, 2026 01:46

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 25, 2026

QIU023 marked this pull request as draft August 25, 2026 01:47

QIU023 force-pushed the k3_pp_text branch from 6adf230 to 1326f94 Compare August 25, 2026 07:49

QIU023 marked this pull request as ready for review August 26, 2026 05:44

QIU023 mentioned this pull request Aug 26, 2026

Kimi K3 issue tracking #4272

Open

18 tasks

tianyu-l requested review from acisseJZhong and sanketpurandare August 26, 2026 06:18

QIU023 changed the title ~~Pipeline Parallelism support for Kimi K3 (text LLM side): PP Virtual Stage Cache Adapter for Attention Residual~~ [Kimi K3] Pipeline Parallelism support (text LLM side): PP Virtual Stage Cache Adapter for Attention Residual Aug 31, 2026

QIU023 mentioned this pull request Sep 1, 2026

[Kimi K3] Expert Parallelism support #4314

Merged

tianyu-l requested changes Sep 3, 2026

View reviewed changes

QIU023 force-pushed the k3_pp_text branch from 087c4d1 to 0e7cc5e Compare September 4, 2026 09:30

QIU023 requested a review from tianyu-l September 7, 2026 09:34

tianyu-l reviewed Sep 10, 2026

View reviewed changes

tianyu-l reviewed Sep 11, 2026

View reviewed changes

QIU023 mentioned this pull request Sep 11, 2026

[Kimi K3] Tensor/Sequence Parallelism support #4499

Merged

QIU023 force-pushed the k3_pp_text branch from 75045fe to 66601a7 Compare September 11, 2026 09:18

QIU023 requested a review from tianyu-l September 13, 2026 10:09

QIU023 force-pushed the k3_pp_text branch from 1ca5b0c to de6f295 Compare September 14, 2026 08:16

QIU023 mentioned this pull request Sep 16, 2026

[RFC] Kimi K3 Architecture support for torchtitan based pre-training and post-training (with veRL) #3029

Open

17 tasks

tianyu-l reviewed Sep 20, 2026

View reviewed changes

pytorch-bot Bot added the ciflow/fake-pg Run 1-GPU Fake PG integration tests label Sep 21, 2026

QIU023 force-pushed the k3_pp_text branch from de6f295 to d33f0a1 Compare September 21, 2026 07:55

QIU023 requested review from IvanKobzarev, SherlockNoMad and aditvenk as code owners September 21, 2026 07:55

pytorch-bot Bot added the ciflow/rl label Sep 21, 2026

QIU023 force-pushed the k3_pp_text branch from d33f0a1 to 498dc57 Compare September 21, 2026 08:05

QIU023 requested a review from tianyu-l September 21, 2026 09:34

shuhuayu self-requested a review September 21, 2026 22:46

shuhuayu requested changes Sep 22, 2026

View reviewed changes

shuhuayu requested changes Sep 23, 2026

View reviewed changes

QIU023 added 7 commits September 23, 2026 05:42

kimi_k3: the pipeline cell in the B200 suite is pp4 x vp4

ff9a0ac

One recipe and one cell, as the llama3 pipeline cells are, with the split spelled out because the debug model's units do not divide into 16 stages. Type checking is off, which the pipeline recipes upstream also do.

QIU023 force-pushed the k3_pp_text branch from 4b47d12 to 3f9201e Compare September 23, 2026 08:20

shuhuayu reviewed Sep 23, 2026

View reviewed changes

QIU023 requested a review from shuhuayu September 23, 2026 09:23

		# ----- Rank-shared cache across virtual stages ----------------------------- #


		class RankLocalCache:

global stage (rank, virtual)	layers	Block AttnRes stored stack	PP Comm, fwd (no cache)	PP Comm, fwd (cache)	already cached
s0 (r0, v0)	emb, 0-1	e	-	-	-
s1 (r1, v0)	2-4	x4	s0 -> [e]	s0 -> [e]	-
s2 (r0, v1)	5-6	-	s1 -> [e, x4]	s1 -> [x4]	[e]
s3 (r1, v1)	7, head	-	s2 -> [e, x4]	s2 -> [] (empty payload)	[e, x4]

global stage (rank, virtual)	PP Comm, bwd (no cache)	PP Comm, bwd (cache)	deposits into the rank cache	collects from the rank cache
s3 (r1, v1)	-	-	[de, dx4]	-
s2 (r0, v1)	s3 -> [de, dx4]	no op (empty payload)	[de]	-
s1 (r1, v0)	s2 -> [de, dx4]	s2 -> [dx4]	-	s3 stored: [de, dx4]
s0 (r0, v0)	s1 -> [de]	s1 -> [de]	-	s2 stored: [de]

global stage (rank, virtual)	layers	Block AttnRes stored stack	PP Comm, fwd (no cache)	PP Comm, fwd (cache)	already cached
s0 (r0, v0)	emb, 0	e	-	-	-
s1 (r1, v0)	1-2	-	s0 -> [e]	s0 -> [e]	-
s2 (r2, v0)	3	-	s1 -> [e]	s1 -> [e]	-
s3 (r3, v0)	4	x4	s2 -> [e]	s2 -> [e]	-
s4 (r0, v1)	5	-	s3 -> [e, x4]	s3 -> [x4]	[e]
s5 (r1, v1)	6	-	s4 -> [e, x4]	s4 -> [x4]	[e]
s6 (r2, v1)	7	-	s5 -> [e, x4]	s5 -> [x4]	[e]
s7 (r3, v1)	8	x8	s6 -> [e, x4]	s6 -> [] (empty payload)	[e, x4]
s8 (r0, v2)	9	-	s7 -> [e, x4, x8]	s7 -> [x8]	[e, x4]
s9 (r1, v2)	10	-	s8 -> [e, x4, x8]	s8 -> [x8]	[e, x4]
s10 (r2, v2)	11	-	s9 -> [e, x4, x8]	s9 -> [x8]	[e, x4]
s11 (r3, v2)	12	x12	s10 -> [e, x4, x8]	s10 -> [] (empty payload)	[e, x4, x8]
s12 (r0, v3)	13	-	s11 -> [e, x4, x8, x12]	s11 -> [x12]	[e, x4, x8]
s13 (r1, v3)	14	-	s12 -> [e, x4, x8, x12]	s12 -> [x12]	[e, x4, x8]
s14 (r2, v3)	15	-	s13 -> [e, x4, x8, x12]	s13 -> [x12]	[e, x4, x8]
s15 (r3, v3)	head	-	s14 -> [e, x4, x8, x12]	s14 -> [] (empty payload)	[e, x4, x8, x12]

global stage (rank, virtual)	PP Comm, bwd (no cache)	PP Comm, bwd (cache)	deposits into the rank cache	collects from the rank cache
s15 (r3, v3)	-	-	[de, dx4, dx8, dx12]	-
s14 (r2, v3)	s15 -> [de, dx4, dx8, dx12]	no op (empty payload)	[de, dx4, dx8]	-
s13 (r1, v3)	s14 -> [de, dx4, dx8, dx12]	s14 -> [dx12]	[de, dx4, dx8]	-
s12 (r0, v3)	s13 -> [de, dx4, dx8, dx12]	s13 -> [dx12]	[de, dx4, dx8]	-
s11 (r3, v2)	s12 -> [de, dx4, dx8, dx12]	s12 -> [dx12]	[de, dx4, dx8]	s15 stored: [dx12]
s10 (r2, v2)	s11 -> [de, dx4, dx8]	no op (empty payload)	[de, dx4]	s14 stored: [dx8]
s9 (r1, v2)	s10 -> [de, dx4, dx8]	s10 -> [dx8]	[de, dx4]	s13 stored: [dx8]
s8 (r0, v2)	s9 -> [de, dx4, dx8]	s9 -> [dx8]	[de, dx4]	s12 stored: [dx8]
s7 (r3, v1)	s8 -> [de, dx4, dx8]	s8 -> [dx8]	[de, dx4]	s11 stored: [dx8], s15 stored: [dx8]
s6 (r2, v1)	s7 -> [de, dx4]	no op (empty payload)	[de]	s10 stored: [dx4], s14 stored: [dx4]
s5 (r1, v1)	s6 -> [de, dx4]	s6 -> [dx4]	[de]	s9 stored: [dx4], s13 stored: [dx4]
s4 (r0, v1)	s5 -> [de, dx4]	s5 -> [dx4]	[de]	s8 stored: [dx4], s12 stored: [dx4]
s3 (r3, v0)	s4 -> [de, dx4]	s4 -> [dx4]	-	s7 stored: [de, dx4], s11 stored: [de, dx4], s15 stored: [de, dx4]
s2 (r2, v0)	s3 -> [de]	s3 -> [de]	-	s6 stored: [de], s10 stored: [de], s14 stored: [de]
s1 (r1, v0)	s2 -> [de]	s2 -> [de]	-	s5 stored: [de], s9 stored: [de], s13 stored: [de]
s0 (r0, v0)	s1 -> [de]	s1 -> [de]	-	s4 stored: [de], s8 stored: [de], s12 stored: [de]

shape	option 1	option 2	this PR	this PR, peak per hop	option 2, peak
pp8 x vp1	28	28	28	7	7
pp4 x vp2	18	28	18	3	7
pp2 x vp4	7	28	7	1	7
pp8 x vp4	52	136	52	2	8
pp4 x vp8	24	136	24	1	8
pp16 x vp2	96	136	96	4	8

cell	loss step 1	loss step 10	loss step 20	grad norm step 1	grad norm step 10	grad norm step 20
dp1	`12.605700`	`3.114620`	`3.373330`	`18.625`	`5.4375`	`3.9844`
pp2	`12.605700` bitwise	`3.227050` +3.61%	`3.288290` -2.52%	`18.75` +0.67%	`5.6875` +4.60%	`3.7344` -6.27%
pp2 x vp2, cache on	`12.605700` bitwise	`3.150940` +1.17%	`3.349300` -0.71%	`18.625` 0%	`3.7188` -31.61%	`4.0625` +1.96%
pp2 x vp2, cache off	`12.605700` bitwise	`3.514970` +12.85%	`3.281700` -2.72%	`18.625` 0%	`6.0312` +10.92%	`3.6719` -7.84%
dp1, accumulation order reversed	`12.605700` bitwise	`3.247610` +4.27%	`3.295370` -2.31%	`18.625` 0%	`6.6562` +22.41%	`4.25` +6.67%

cell	loss step 1	loss step 10	loss step 20	grad norm step 1	grad norm step 10	grad norm step 20
dp2	`12.521140`	`3.221120`	`2.839810`	`16.375`	`7.0625`	`2.4844`
dp2 x pp2	`12.521140` same	`3.201920` -0.60%	`2.846440` +0.23%	`16.375` 0%	`5` -29.20%	`2.6562` +6.92%
dp2 x pp2 x vp2, cache on	`12.521140` same	`3.205680` -0.48%	`2.682970` -5.52%	`16.375` 0%	`5.375` -23.89%	`2.5156` +1.26%
dp2 x pp2 x vp2, cache off	`12.521140` same	`3.189240` -0.99%	`2.847220` +0.26%	`16.375` 0%	`5.5938` -20.80%	`2.375` -4.40%
dp2 x ep2	`12.521140` same	`3.149310` -2.23%	`2.969330` +4.56%	`16.375` 0%	`5.0625` -28.32%	`2.9531` +18.87%

cell	loss, step 1	step 10	step 20	step 100	grad norm, step 1	step 10	step 20	step 100
dp1 (reference)	`12.609980`	`3.595760`	`2.998930`	`2.533960`	`16.9193`	`5.0429`	`2.1888`	`1.6014`
pp2 (all 100 steps identical)	`12.609980` identical	`3.595760` identical	`2.998930` identical	`2.533960` identical	`16.9193` identical	`5.0429` identical	`2.1888` identical	`1.6014` identical
pp2 x vp2, naive	`12.609980` identical	`3.595760` identical	`2.998930` identical	`2.533960` identical	`16.9193` identical	`5.0429` identical	`2.1888` identical	`1.6014` identical
pp4 x vp4, naive	`12.609980` identical	`3.595760` identical	`2.998930` identical	`2.533960` identical	`16.9193` identical	`5.0429` identical	`2.1888` identical	`1.6014` identical
pp2 x vp2, cached	`12.609980` identical	`3.554240` -1.15%	`3.055690` +1.89%	`2.571310` +1.47%	`16.9177` -0.01%	`5.2631` +4.37%	`2.9058` +32.76%	`1.6073` +0.37%
pp4 x vp4, cached	`12.609980` identical	`3.309360` -7.96%	`3.007760` +0.29%	`2.554590` +0.81%	`16.9167` -0.02%	`3.5408` -29.79%	`2.0286` -7.32%	`1.5763` -1.57%

cell	loss, step 1	step 10	step 20	step 100	grad norm, step 1	step 10	step 20	step 100
dp2 (reference)	`12.580740`	`3.576250`	`2.976110`	`2.420420`	`14.4170`	`12.4659`	`2.4576`	`1.0313`
dp2 x pp2 (all 100 steps identical)	`12.580740` identical	`3.576250` identical	`2.976110` identical	`2.420420` identical	`14.4170` identical	`12.4659` identical	`2.4576` identical	`1.0313` identical
dp2 x pp2 x vp2, naive (all 100 steps identical)	`12.580740` identical	`3.576250` identical	`2.976110` identical	`2.420420` identical	`14.4170` identical	`12.4659` identical	`2.4576` identical	`1.0313` identical
dp2 x pp2 x vp2, cached	`12.580740` identical	`3.305220` -7.58%	`2.949200` -0.90%	`2.434300` +0.57%	`14.4191` +0.01%	`4.5317` -63.65%	`2.6572` +8.12%	`1.0893` +5.62%

cell	loss, step 1	step 10	step 50	step 100	grad norm, step 1	step 10	step 50	step 100
fp32 dp1, reference	`8.055300`	`3.405580`	`2.578440`	`2.501960`	`2.1551`	`1.2746`	`0.7113`	`0.8558`
fp32 dp1, stock accumulation	`8.055300` identical	`3.405580` identical	`2.578440` identical	`2.501960` identical	`2.1551` identical	`1.2746` identical	`0.7113` identical	`0.8558` identical
fp32 pp2 x vp2, naive	`8.055300` identical	`3.405580` identical	`2.576750` -0.07%	`2.499220` -0.11%	`2.1551` identical	`1.2746` identical	`0.7125` +0.17%	`0.8608` +0.58%
fp32 pp2 x vp2, cached	`8.055300` identical	`3.405580` identical	`2.576820` -0.06%	`2.499740` -0.09%	`2.1551` identical	`1.2746` identical	`0.7127` +0.20%	`0.8613` +0.64%
fp32 pp4 x vp4, naive	`8.055300` identical	`3.405580` identical	`2.576820` -0.06%	`2.499740` -0.09%	`2.1551` identical	`1.2746` identical	`0.7127` +0.20%	`0.8613` +0.64%
fp32 pp4 x vp4, cached	`8.055300` identical	`3.405530` -0.00%	`2.574910` -0.14%	`2.501610` -0.01%	`2.1551` identical	`1.2744` -0.02%	`0.7067` -0.65%	`0.8648` +1.05%

		return grad_delta, deposits


		# TODO: decide whether this stage belongs in torch.distributed.pipelining.

		return config


		def kimi_k3_debugmodel_pp4_vp4() -> Trainer.Config:

Conversation

QIU023 commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Results

Test plan

A torch.distributed.pipelining finding

Uh oh!

QIU023 commented Aug 26, 2026

Uh oh!

tianyu-l left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

QIU023 Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

QIU023 Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

QIU023 commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

QIU023 commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

QIU023 commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

QIU023 commented Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

QIU023 commented Sep 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

QIU023 commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

QIU023 commented Aug 25, 2026 •

edited

Loading

A `torch.distributed.pipelining` finding

tianyu-l left a comment •

edited

Loading

QIU023 Sep 13, 2026 •

edited

Loading

QIU023 Sep 21, 2026 •

edited

Loading

QIU023 commented Sep 3, 2026 •

edited

Loading

QIU023 commented Sep 4, 2026 •

edited

Loading

QIU023 commented Sep 7, 2026 •

edited

Loading

QIU023 commented Sep 8, 2026 •

edited

Loading

QIU023 commented Sep 10, 2026 •

edited

Loading

QIU023 Sep 12, 2026 •

edited

Loading

QIU023 Sep 11, 2026 •

edited

Loading

QIU023 Sep 12, 2026 •

edited

Loading

QIU023 Sep 12, 2026 •

edited

Loading

QIU023 Sep 13, 2026 •

edited

Loading

QIU023 Sep 14, 2026 •

edited

Loading