[ZenAlexa]

why

ChunkedLossWrapper runs several head-only backwards before decoder backward. main restores head sync unconditionally on the final chunk, overwriting the caller's FSDP and HSDP accumulation policies. an inner root callback can also reduce earlier decoder gradients before the current microbatch arrives.

pipeline schedules choose their backward policy after the loss runs. the wrapper now follows the schedule's manual-finalization owner from the first head pass, so GPipe and 1F1B both defer every head microbatch until schedule finalization.

what changed

• keep runtime policy and grouped-forward handling in torchtitan/distributed/fsdp_lifecycle.py, shared by loss and validation. the loss loop uses before_last_forward and after_forward without private FSDP reads.
• save and restore each group's reduction, replica all-reduce and reshard settings and the root's is_last_backward, including when a chunk raises.
• keep decoder gradients, prefetched weights and root finalization out of the inner head-only backwards. outside manual finalization, retain pending-gradient detection so a shared norm/head group coalesces its accumulated gradients into one reduction.
• defer head reduction unconditionally under manual pipeline finalization, including the first microbatch. parameter residency stays with the schedule. outside pipelines, preserve the caller's early-reshard policy and the unreduced-gradient lifetime.
• complete partial grouped forwards for training with a frozen decoder and for evaluation cleanup. register root iteration reset after dynamic pipeline metadata inference on success and failure, preserving gradients between normal accumulated schedule steps.
• replace the private-state mock with real FSDP lifecycle observations, including tuple multi-output unshard and reshard behavior.

GraphTrainer's production path uses SimpleFSDP and keeps its separate reduction owner.

current integration and validation

implemented the lifecycle correction in 049301bfb83c and integrated main@196b4a5fcfb3 through ordinary merge 636e16da4e78, with complete tree 7482898122bd. the latest main integration leaves all eight partial-backward patch files unchanged. pipeline manual finalization now unconditionally owns head reduction, including the first microbatch, and the wrapper leaves parameter residency with the schedule. the last-forward hook is named before_last_forward.

on macOS arm64, Python 3.13.18, PyTorch 2.16.0.dev20261026 and four Gloo ranks, the current lifecycle scenario passed 56 combinations plus one real grouped, frozen-decoder, two-output scenario. the matrix covers FSDP and HSDP, GPipe and 1F1B, deferred finalization, frozen decoders, tied weights, caller-policy restoration, reset/retry and losses, gradients and updated parameters against a full-logits reference. each finalized pipeline accumulation performs two reduce-scatters for a trainable decoder or one for a frozen decoder. an accumulation without finalization performs zero. metadata inference is counted separately because its gradients are discarded.

the two-output scenario runs two chunks of each output through one real FSDP group, observes one explicit unshard and one actual reshard, and compares its weighted loss and gradients against full logits. the former private-state mock and its duplicate checks are removed.

the same reduction-count scenario fails on the previous e0029100edbb revision, where GPipe performs five reductions in the checked accumulation. the final combined run passed the lifecycle entry and 89 selected loss and pipeline regressions, reporting 90 passed with one LoRA CPU FlexAttention case deselected. changed-file pre-commit checks, including Pyrefly, passed.

GLOO_SOCKET_IFNAME=lo0 pytest tests/unit_tests/cpu/components/test_chunked_loss_fsdp.py -q
pytest tests/unit_tests/cpu/components/test_loss.py tests/unit_tests/cpu/distributed/test_pipeline_parallel.py -k 'not lora' -q

numerical and hardware merge gate

pipeline head gradients now use the same reduction owner as the rest of their stage. this deliberately changes floating-point accumulation order from schedule-dependent partial head reductions to the schedule's final reduction. the existing tolerances and golden files remain unchanged.

the current eight-H100 attempt reached no workload and provides no GPU result. please run the current GPU lifecycle and model workflows, including llama3_fsdp+tp+pp, deepseek_v3_fsdp+cp+pp+ep and gpt_oss_pp+fsdp+cp+ep+sacop, then regenerate the affected A10G and MI350X goldens from those exact backends. same-environment loss and grad-norm comparisons, last-stage peak memory and throughput, and representative convergence remain required before merge. historical device results in this description retain their original revisions.

upstream lifecycle trackers

• pytorch/pytorch#203176 tracks scoped partial backwards and completion of repeated standalone grouped forwards, including frozen inputs and evaluation cleanup.
• pytorch/pytorch#203177 tracks value-based policy snapshots and exact restoration, including independent HSDP replica all-reduce settings.
• pytorch/pytorch#203178 tracks FSDP iteration reset after pipeline metadata inference.

all three include self-contained two-rank CPU reproductions on PyTorch 2.16.0.dev20261026 and references to the current upstream implementation. the related pipeline policy RFC owns broader configurable schedule reduction and residency policy. this PR implements the existing schedule's finalization contract.

project GPU results on e0029100edbb

the November 2 multi-GPU job executed and passed the prior 40-combination lifecycle scenario. the full job reported 46 passed, 11 skipped, 358 deselected and one xfailed on the eight-A10G runner with PyTorch 2.16.0.dev20261101+cu132. checkout dcc041a merged e0029100edbb into d86c77d609b4. this is evidence for the earlier revision, before unconditional pipeline-head deferral and the added multi-output scenario. the model job on that revision again failed the strict imported goldens for DeepSeek-V3 and GPT-OSS, beginning at grad-norm steps six and three respectively.

earlier CPU lifecycle evidence

at 81b3eec872d2, the HSDP/pipeline correction passed complete loss, gradient and updated-parameter checks against the full-logits reference, preserving caller policy. a separate four-rank metadata-inference failure scenario passed GPipe and 1F1B with tied weights, second-chunk failure, cleanup, retry, delayed accumulation and train → eval → train. cleanup occurred at each metadata-inference boundary and remained absent between normal accumulation groups. those results retain that revision.

project GPU results on 142f59599

• the multi-GPU job ran and passed the previous 24-combination lifecycle scenario. the enclosing job reported 45 passed, 11 skipped, 268 deselected and one xfailed.
• the real-PG features job passed the two-rank deferred-gradient-reduction execution, including ten training steps and CUDA Graph capture.
• the model job completed the GPT-OSS and DeepSeek-V3 pipeline training runs and failed their strict imported-golden comparisons. these results belong to 142f59599, before the latest merge and HSDP pipeline correction.

project GPU results on 65b11734

• the multi-GPU job actually ran the four-rank lifecycle scenario and passed its previous 16 combinations, including HSDP, failure recovery and reduce_scatter_single counts. runtime: eight-A10G runner, Python 3.12.15, PyTorch 2.16.0.dev20261026+cu132, NCCL 2.32.3, driver 580.178.04. the checked-out merge tree matches 65b11734.
• the real-PG features job passed fsdp_deferred_gradient_reduction on two ranks and fsdp_pp_deferred_gradient_reduction on four ranks, both for ten steps with recorded CUDA graphs. these are execution checks, without a baseline/head bitwise comparison.
• fsdp_per_group_cuda_graph and shared_eval_cuda_graph passed in the single-GPU fake-PG job.

earlier native evidence

the two-H100 comparison on 1bbb64b44 used driver 580.95.05, PyTorch 2.16.0.dev20261018+cu130, CUDA 13.0, NCCL 2.32.3, seed 42 and deterministic algorithms. sixteen reduced-model configurations across two revisions completed ten steps. ordinary FSDP loss, gradients, grad norm and parameters were bitwise equal; deferred eager reduce-scatters dropped from seven to three separate owners or two grouped owners. both revisions peaked at 676.144 MiB on the large-vocabulary default workload. native lifecycle and reset/retry checks also passed.

those scripts explicitly set is_last_backward(final) in both modes. their numbers apply to that setting and revision. the checked-in scenario follows the actual default is_last_backward=True. the earlier separate four-GPU allocation expired in its ten-minute queue budget without a test process or result and was removed.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[ZenAlexa]

@tianyu-l could you approve the pending Lint, CPU Unit Test, and GPU Unit Tests workflows for this revision? The two-H100 deterministic comparison and deferred-reduction, CUDA Graph, and recovery results are in the description.

[tianyu-l]

@claude explain the issue to me in simpler ways -- use example if it's helpful

also review the PR

[claude]

Claude finished @tianyu-l's task in 2m 59s —— View job

---

Part 1 -- The problem, explained simply

Background: what deferred accumulation is trying to do. Say one optimizer step uses 3 microbatches (mb0, mb1, mb2). With defer_fsdp_gradient_reduction, TrainingEngine (training_engine.py:772-776) tells the FSDP root:

	requires_gradient_sync	reshard_after_backward	is_last_backward
mb0, mb1	False	False	False
mb2	True	True	True

So each param group should keep its unsharded gradients locally through mb0 and mb1, then do one reduce-scatter per group during mb2's backward. With HSDP, set_requires_all_reduce(is_last) works the same way for the replica all-reduce.

What ChunkedLossWrapper does inside one microbatch. It splits the sequence into N chunks and runs a separate chunk_loss.backward() for each chunk, only through lm_head, because the hidden states are detached. After the loop, the decoder backward runs once, as a separate autograd pass.

That gives N+1 backward passes per microbatch, and FSDP treats each one as a full backward.

Bug A: the wrapper overwrites the caller's settings. On main, the wrapper does this:

lm_head.set_requires_gradient_sync(False)     # chunks 0..N-2
...
lm_head.set_requires_gradient_sync(True)      # final chunk, unconditionally
...
lm_head.set_reshard_after_backward(True)      # after the loop, unconditionally

Example: on mb0 the engine says "don't sync", but the wrapper turns sync on for lm_head anyway. So lm_head reduce-scatters on every microbatch, and deferral never applies to it. set_requires_gradient_sync(True) also turns all_reduce_grads back on, which undoes the engine's HSDP set_requires_all_reduce(False). These are mostly wasted collectives. The math is still right, since the sum of reductions equals the reduction of the sum.

Bug B: an inner chunk backward reduces the decoder too early. When an autograd pass touches any FSDP module, FSDP's root post-backward callback runs over every group in the root context, not just lm_head. On mb2 (sync=True), the decoder groups are still holding unsharded grads from mb0 and mb1. So:

mb2, chunk-0 backward (lm_head only)
  -> root final callback visits decoder groups
  -> they have pending mb0+mb1 grads and sync=True
  -> reduce-scatter now                         <- reduction #1
mb2, decoder backward
  -> mb2 grads arrive -> reduce-scatter again   <- reduction #2

Each decoder group gets reduced twice instead of once, plus an early is_last_backward finalize. That matches the "7 reduce-scatters -> 3" numbers in the description.

Bug C: a grouped fully_shard([norm, lm_head]) with a frozen decoder. FSDP fires the group's post-forward, which registers the backward hooks, only after all modules in the group have run forward. The chunk loop only calls lm_head, so the group never finishes its forward. With a frozen decoder, the inputs don't require grad, so no input-side hook gets registered either. As a result, nothing triggers the group's backward lifecycle.

The PR's fix:

1. Save each group's reduce/reshard policy and restore it on exit, including on exceptions, instead of hard-coding True.
2. If any group already has pending grads (or HSDP partial state), keep sync off on all groups and set is_last_backward=False during the chunks. The outer decoder backward's root callback then reduces everything, including lm_head, exactly once.
3. For case C, manually call FSDP's _post_forward so the output hook gets registered.

---

Part 2 -- Review

The diagnosis looks right to me, and I agree with the goal. My concerns are mostly about how the fix gets there.

1. Heavy dependence on FSDP2 private internals. This is the main concern. _chunked_loss_fsdp (loss.py:505-571) and the frozen-head branch (loss.py:~775) read and write these private attributes:

• _get_fsdp_state, _lazy_init, _state_ctx.is_last_backward, _state_ctx.all_states, _auto_reshard_after_forward
• on param groups: _fsdp_param_groups, post_forward_mesh_info, reduce_grads, all_reduce_grads, _post_backward_pending, _partial_reduce_output
• state._post_forward(...) and _modules_to_run_forward.clear()

That is a lot of private state for a core component (Core Principle 1: this kind of complexity belongs upstream). Any FSDP2 refactor of these names or semantics would break it silently. The CPU tests won't catch that, because _RecordingFSDPLinear (test_loss.py:801-819) mocks the same private fields with a SimpleNamespace. Suggestions:

• Upstream needs are really two narrow public APIs: (a) get/set policies as values, so they can be saved and restored, plus a "has pending unreduced grads" query; (b) a way to scope the root post-backward callback to one state, or a way to "complete a partial grouped forward". Could we open a PyTorch issue/PR and land this here with a TODO(<link>)?
• At minimum, add a fail-fast check (e.g. hasattr on the private fields, with a clear error) so a PyTorch bump fails loudly instead of quietly doing the wrong thing.

2. Default HSDP behavior changes, but was only checked on CPU. In the HSDP default (non-deferred) mode, set_requires_all_reduce(False) on mb0..mbN-2 leaves _partial_reduce_output set. That makes defer_to_decoder True on mb1+ even without defer_fsdp_gradient_reduction. So the "default path is bitwise identical to main" claim applies to plain FSDP. For HSDP, the default path now uses the new deferral path on CUDA, and that hasn't been run natively (the description says native 4-GPU HSDP is unverified). Please run HSDP (dp_replicate>1) loss/grad_norm bitwise comparisons before merge.

3. Reshard before reduction on the last microbatch with deferral (loss.py:563-571). On mb2 with defer_to_decoder, the caller's reshard_after_backward is True. So group.reshard() runs on the lm_head groups at wrapper exit, before the outer callback's post_backward has reduced their unsharded grads. This relies on FSDP keeping unsharded_param.grad / unsharded_accumulated_grad alive and readable after to_sharded(). It apparently works on the tested nightly, but it's subtle and undocumented. Either skip the early reshard when defer_to_decoder (the outer post_backward already reshards), or add a comment explaining why it's safe.

4. The frozen-decoder grouped-head workaround (loss.py:~775-786).

• Calling _post_forward by hand and clearing _modules_to_run_forward means faking FSDP's forward state machine. This seems like an FSDP limitation (a partially-run grouped forward gets no hooks), so it should be an upstream fix with a TODO.
• It runs every chunk with a tuple of logits from all outputs. Please confirm that registering the output hook N times doesn't trigger N pre-backward unshards or N callback queues.
• Is case C a bug on main today, i.e. are lm_head grads never reduced for head-only training with a grouped norm/head? If so, it deserves its own small PR and a regression test. It's separate from the policy-restoration fix.

5. Test coverage doesn't cover the new behavior. The test diff only adapts the existing mock. There's no checked-in test for:

• caller policy restored after the wrapper (sync, reshard_after_backward, all_reduce, is_last_backward), including when an exception is raised mid-loop
• defer_to_decoder being chosen when pending grads exist, and not chosen otherwise
• the frozen-decoder grouped-head path

The description lists a lot of good ad-hoc verification (Gloo 2/4-rank, 2xH100), but none of it is in the repo. CLAUDE.md asks for CPU unit tests plus a GPU/multi-rank test for parallelism changes. The ciflow/fake-pg label is set; could some of the Gloo scenarios become fake-PG or multi-process tests? Something like "count reduce-scatters on the last microbatch with defer_fsdp_gradient_reduction + chunked loss" would lock in the 7->3 fix.

6. Audit sibling callsites. torchtitan/experiments/graph_trainer/chunked_loss.py:102,124 uses the same pattern: it restores with set_requires_gradient_sync(True, recurse=False) unconditionally, so it overrides the caller's policy and HSDP all_reduce_grads. GraphTrainer currently rejects FSDP2 deferred reduction (test_trace_module.py:163), so this may be moot there, but please confirm that HSDP gradient accumulation isn't affected.

7. Nits.

• _chunked_loss_fsdp yields either None or a callable, and None is overloaded to mean both "not FSDP" and "deferred". A small enum or two explicit fields would read more clearly at loss.py:751-756.
• policies stores 5-tuples that get unpacked as _, _, _ in three places. A small dataclass or NamedTuple would help.
• The docstring sentence "Deferred accumulation retains the unsharded weights and gradients for the next microbatch" is accurate but vague about who reduces lm_head in the deferred case (the outer decoder backward's root callback). Worth stating that explicitly.
• The isinstance(lm_head, FSDPModule) check is now done twice: in the context manager and in the per-chunk frozen-head branch.

Before merge, in addition to the author's request: run fsdp_deferred_gradient_reduction, fsdp_pp_deferred_gradient_reduction (PP's stage backward interacts with is_last_backward differently), fsdp_per_group_cuda_graph, shared_eval_cuda_graph, and an HSDP config with chunked loss.

Overall: the bugs are real and the numerical evidence is strong. My hesitation is the amount of private FSDP2 state being manipulated in core loss.py. I'd prefer narrowing this down: split the frozen-grouped-head fix (C) into its own PR, and pursue upstream public hooks for (1) and (4), tracked with a TODO.

I only read the code; I didn't run any tests (no torch or GPUs in this environment).

[ZenAlexa]

The wrapper now preserves the caller's policies across the inner chunk backwards, so a three-microbatch step can keep its gradients together until decoder backward. Added real four-rank FSDP/HSDP CPU and GPU regressions for grouped and frozen heads, failures with pending gradients, and reset/retry; all 16 CPU combinations pass. The public sync setter also changes HSDP's separate all-reduce flag, so the policy snapshot remains in one private-state boundary. Kept early reshard and documented its ownership: the autograd leaf retains gradients after full weight storage is released; current native coverage and the two model-CI failures are recorded in the description.

[tianyu-l]

@claude take another look

• verify correctness
• see if the readability and maintainability can be improved -- I feel there's room
• give suggestions on what should go to upstream FSDP API

[claude]

Claude finished @tianyu-l's task in 7m 14s —— View job

---

Re-review: chunked-loss FSDP policy preservation

• Gather context (diff vs 4103c72, FSDP2 call paths, engine/PP/fsdp.py callers)
• Verify correctness
• Readability / maintainability suggestions
• Upstream FSDP API suggestions

I only read the code and traced FSDP2's root post-backward callback, post_backward, finalize_backward and the grouped-forward hooks by hand. torch isn't installed in this runner, so I couldn't check the private attribute names against the pinned nightly or run the tests.

---

1. Correctness

Correct as written

• Bug B (decoder reduced too early) is real, and the defer path fixes it. In each inner chunk_loss.backward(), the root final callback calls post_backward() on every group in all_states that isn't already POST_BACKWARD. On the last microbatch, that reduce-scatters the decoder's pending mb0+mb1 gradients during chunk 0, and then again in decoder backward. In the defer path, every group has reduce_grads=False and reshard_after_backward=False, and is_last_backward=False. Those callbacks then only accumulate gradients, and decoder backward does a single reduction.
• Early reshard with pending gradients (my earlier item 3) is safe. FSDP's own reduce_grads=False branch of post_backward calls reshard() before to_accumulated_grad_if_needed(). to_sharded() only frees all-gather storage, and _unsharded_param.grad survives. So the wrapper relies on the same contract FSDP relies on internally, and the new comment at loss.py:571-572 is accurate. Withdrawn.
• Bug C (frozen decoder, grouped head) is real on main. Only chunk 0's logits finish the grouped forward. Chunks 1..N-1 get neither an output pre-backward hook nor an input post-backward hook, because frozen inputs carry no hook. Their lm_head gradients sit in unsharded_param.grad with nothing to reduce them. Calling _post_forward each chunk fixes this. Implicit prefetch indices stay balanced, because each chunk adds one post-forward index and its backward pops one.
• Exception path. The finally block restores everything the setup touched. The test checks the policies both after a mid-loop failure and after reset_iter_state().

Needs attention before merge

(a) PP: what the wrapper preserves is the previous microbatch's policy. Under PP, ChunkedLossWrapper runs inside the last stage's forward step (loss computation). PyTorch pipelining sets the FSDP policy in backward_maybe_with_nosync, i.e. after the loss. On main, the wrapper force-reduces lm_head on every microbatch, so this didn't matter. With this PR, lm_head's reduction follows whatever the previous stage backward or the previous step's finalization left behind. The pending-gradient detection can also switch to the defer path partway through the schedule.

Mathematically it should still be correct, because PP's finalization reduces whatever is pending. But it changes how lm_head gradients are grouped into reductions, which changes floating-point summation order. My guess is that this explains the DeepSeek-V3 / GPT-OSS PP loss/grad-norm failures in the model job. I haven't confirmed it, and it needs a same-environment main-vs-PR rerun. If the comparison there is exact-match, the expected-numerics change needs to be stated and re-baselined. If it's tolerance-based, it's a real bug. Either way, it shouldn't land unexplained.

(b) Production uses only the grouped layout; the separate layout is test-only. torchtitan/distributed/fsdp.py:339-358 always does fully_shard([norm, lm_head]), or [tok_embeddings, norm, lm_head] with weight tying. The test matrix covers grouped, which is good. The tied 3-module group isn't covered, though: its group pre-forward fires on tok_embeddings, whose input is integer token ids, so no input post-backward hook gets registered. Please add a grouped_tied variant to the matrix, or at least run the weight-tied model through fsdp_deferred_gradient_reduction.

(c) A pre-existing perf issue that the defer path fixes only some of the time. With the default is_last_backward=True, chunk 0's root callback runs _finalize_backward() on all states. As I recall upstream (not checked against your nightly), that clears every decoder group's _post_forward_indices and the shared post_forward_order. Decoder backward would then lose implicit backward prefetch for every layer. The ep_degree > 1 path uses explicit prefetch (fsdp.py:510-520) and is unaffected. The defer path avoids this because it sets is_last_backward=False. The non-defer path, which is the default plain-FSDP step and the frozen-decoder case, still has it. Worth confirming with a trace: either way, it is a strong argument for §2.1.

(d) Test robustness. The test counts reductions by patching dist.reduce_scatter_single. Upstream FSDP2 has historically called dist.reduce_scatter_tensor. If CI's nightly uses a different entry point, the count assertion fails loudly, which is fine. Please confirm it runs in the CPU and GPU CI images, not just locally.

---

2. Readability and maintainability

There is room here. In rough order of impact:

2.1 Drop pending-gradient detection by always isolating the head pass (when decoder_requires_grad).
defer_to_decoder exists because the inner callbacks touch non-head groups. Mask those groups unconditionally instead:

# decoder_requires_grad: decoder backward follows in this microbatch and owns
# the root finalization. Inner head-only passes must not reduce, reshard or
# finalize any non-head group.
for group in non_head_groups:
    group.reduce_grads = False
    group.reshard_after_backward = False
state_ctx.is_last_backward = False
# head groups: reduce_grads=False until the final chunk, then the caller's policy (unchanged)

• In the default plain-FSDP step, non-head groups have no gradients during chunks, so masking them changes nothing numerically. lm_head still reduces at the final chunk, overlapped with decoder backward, exactly as on main. The bitwise-identical default is preserved.
• In deferred mode on the last microbatch, decoder groups can't reduce early, so bug B is fixed without the _post_backward_pending / _partial_reduce_output reads.
• It also fixes (c) for every path, and stops a prefetched last block from being resharded and then re-gathered between chunk 0 and decoder backward.
• It deletes the None | Callable yield and the mid-loop restore_gradient_sync() call. There is a single path when the decoder is trainable and a single path when it is frozen. The frozen path leaves the other groups and is_last_backward alone, since no decoder backward follows.
• Cost: with a grouped [norm, lm_head] in deferred mode, the last microbatch does one extra small reduce-scatter (norm's mb2 gradient, in decoder backward), so the test would expect 3 instead of 2. I think that's worth it to drop the two private reads that encode FSDP's notion of "pending".

Fix this →

2.2 Put all private FSDP access behind one object, and move it to torchtitan/distributed/.
The PR description says the private accesses stay inside the chunked FSDP boundary. But the frozen-head _post_forward / _modules_to_run_forward workaround sits in __call__ (loss.py:778-788), which re-checks isinstance(lm_head, FSDPModule). It also duplicates post_validate_fsdp_for_cuda_graph (distributed/fsdp.py:57-59), which has a matching TODO: Use public FSDP APIs. Suggestion:

# torchtitan/distributed/fsdp.py
def finish_grouped_forward(module: FSDPModule, output: Any) -> Any: ...  # shared with post_validate_fsdp_for_cuda_graph

@contextmanager
def isolated_head_backward(head, *, run_backward, decoder_requires_grad) -> Iterator[HeadBackwardHooks]:
    ...

class HeadBackwardHooks:          # no-op implementation when head is not an FSDPModule
    def after_forward(self, output): ...   # frozen grouped head -> finish_grouped_forward
    def before_last_chunk(self): ...       # restore caller sync on head groups

Then loss.py is FSDP-agnostic. It needs no FSDPModule import or private access, and the chunk loop reads hooks.before_last_chunk() / logits = hooks.after_forward(logits). This follows the "model-agnostic parallelism helpers -> torchtitan/distributed/" rule, and the GraphTrainer variant could reuse the hooks interface later.

2.3 Smaller items

• Replace the 5-tuples (unpacked as _, _, _ three times) with a @dataclass(slots=True) class _GroupPolicy that has snapshot(group) and restore(group) methods.
• loss.py:566-569: name the condition, e.g. release_head_weights = reshard_after_backward or not run_backward # forward-only has no backward to reshard, and precompute head_groups = set(state._fsdp_param_groups) instead of a list membership test inside the loop.
• The class docstring (loss.py:603-607) should say who reduces lm_head in each case: the final chunk when the decoder is frozen or nothing is pending; otherwise decoder backward via the group's post-backward hook or the root callback. "Deferred accumulation retains the unsharded weights and gradients" is still vague.
• Most of the loss.py diff is re-indenting the chunk body. Pulling the per-chunk body into a _run_chunk(...) method would shrink both this diff and __call__.
• Tests:
  • test_accumulation_and_recovery is a single 150-line method with closures, using raise RuntimeError instead of assertions. Split it into helpers (_build, _accumulate, _inject_failure, _check_against_reference), and use torch.testing.assert_close / self.assertEqual so failures show diffs.
  • _RecordingFSDPLinear in test_loss.py now rebuilds FSDP's private layout with SimpleNamespace. Those tests would pass even if FSDP renamed the fields. Now that the real-FSDP test exists, I'd trim the mock tests back to the event-order checks (unshard once, N forwards) and drop the policy assertions from them.

---

3. What belongs in the upstream FSDP API

Listed by how much private state each removes from torchtitan:

1. A scoped "partial backward" for the root callback. This is the main one. For example: with fsdp_root.partial_backward(): or backward(..., fsdp_partial=True). Inside the scope, the root final callback only runs post_backward for states whose hooks actually fired, skips _finalize_backward, and leaves post_forward_order / _post_forward_indices alone. Chunked loss, PP's split-backward / zero-bubble, and any multi-pass backward need these semantics. It would remove the is_last_backward overwrite and the non-head group masking, and it would fix (c) upstream.
2. Value-based policy snapshot and restore.

   policy = module.get_fsdp_policy(recurse=False)   # reduce_grads, all_reduce_grads, reshard_after_backward, reshard_after_forward (incl. int / mesh info)
   with module.fsdp_policy(requires_gradient_sync=False, reshard_after_forward=False, reshard_after_backward=False, recurse=False): ...

   Also decouple set_requires_gradient_sync from all_reduce_grads, or add a keep_all_reduce=True argument. Today set_requires_gradient_sync(True) silently undoes set_requires_all_reduce(False), which is the original HSDP part of bug A.
3. Finishing a partial grouped forward. Either FSDPModule.finish_forward(output), or make a re-entrant call to one member of a fully_shard([a, b]) group register the output pre-backward hook on every call. This would replace both the frozen-head workaround here and post_validate_fsdp_for_cuda_graph.
4. A documented contract that reshard() keeps unreduced gradients. FSDP already relies on this internally. Documenting it turns loss.py:571-572 into a contract instead of an observation.
5. has_pending_gradients() / has_pending_partial_reduction(). Only needed if (1) doesn't land and §2.1 isn't adopted.

Each private access should get a TODO(<upstream issue link>) pointing at the item above that replaces it. Please open (1)–(3) as PyTorch issues, since they are the exit path for this code.

---

Summary

• Correctness: the fixes for B and C are correct, and the early-reshard concern is resolved.
• Blocking:
  • The PP loss/grad-norm failures need an explanation; (a) is my hypothesis.
  • The weight-tied 3-module group has no coverage.
  • The real GPU/HSDP runs listed in the description are still outstanding.
• Maintainability: §2.1 removes the pending-gradient detection and the overloaded yield while keeping the default path bitwise-identical. §2.2 actually puts all private FSDP access in one place, in torchtitan/distributed/.
• Upstream: a scoped partial backward, value-based policy get/set (with the HSDP all-reduce decoupled), and grouped-forward completion.

[tianyu-l]

should this be in upstream pytorch pipelining code or should it be here?

[ZenAlexa]

this belongs in upstream PipelineStage cleanup.I kept the registration here cuz preserving deferred reduction leaves metadata-probe state outside param.grad, and the first real step needs that state cleared.

main forces head sync on in the last chunk, which clears that state early, so the cleanup needs to stay with this fix

[ZenAlexa]

tracked the cleanup in pytorch/pytorch#203178 with a standalone reproduction. the local registration keeps discarded metadata-probe gradients out of the first training step until PipelineStage owns that reset.

[tianyu-l]

should this be a separate file, or should it be included in fsdp.py, or is it worth creating a folder for fsdp/

[ZenAlexa]

I think should keep it separate since this file owns runtime state and fsdp.py owns sharding setup

both loss and validation use it, and keeping the private accesses here gives the upstream replacement one place to remove

added the module docstring and moved the grouped-forward TODO to the helper that owns it

[tianyu-l]

@claude reivew my latest comments

[claude]

Claude finished @tianyu-l's task in 1m 59s —— View job

---

Replies to your three inline comments

• Read the current diff of fsdp_lifecycle.py, fsdp.py, pipeline_parallel.py and loss.py
• Traced the callers and the existing register_post_metadata_inference_cleanup users
• Answers below

I couldn't import torch in this runner (permission prompt), so I haven't checked the upstream pipelining and FSDP source on the pinned nightly. The points about upstream internals below come from reading the code and from memory.

---

1. pipeline_parallel.py:517-523: should the metadata-inference cleanup be upstream or here?

It belongs upstream in the long run. Keep it here as a stopgap with a TODO.

• The hook API is already upstream. register_post_metadata_inference_cleanup is a method on the PyTorch schedule; torchtitan doesn't define it. Dist-MoE already uses it (models/common/dist_moe/runtime.py:396-402). So the only open question is who registers the FSDP cleanup.
• The bug comes from PipelineStage and FSDPModule interacting, not from torchtitan. Metadata inference runs a real backward. The stage then clears the sharded .grad, but it doesn't touch FSDP's per-iteration state: unsharded or accumulated grads kept under reduce_grads=False, and pending post-backward state. Upstream stage.py already special-cases FSDPModule (backward_maybe_with_nosync, and the manual reduce/finalize at the end of a step). The cleanup should live in that same file, next to the existing grad clearing, e.g. if isinstance(self.submod, FSDPModule): self.submod.reset_iter_state(). Then every PipelineStage + FSDP2 user gets the fix, not only torchtitan.
• Until then, here is the right place. _build_pipeline_schedule is the one spot where the schedule and stages meet, and GraphTrainer (SimpleFSDP) is unaffected because the isinstance check filters it out. Suggestions:
  • Add a TODO(<pytorch issue>) stating that PipelineStage should reset FSDPModule iteration state after metadata inference.
  • Reword the comment so it says why the stage's own cleanup isn't enough, e.g. "The stage clears param.grad, but FSDP keeps unreduced unsharded grads and pending post-backward state outside param.grad when reduction is deferred."
• Scope question for the author: can this leak happen on main (plain FSDP+PP with defer_fsdp_gradient_reduction)? Or only with this PR's defer_head_reduction under manual_backward_finalization?
  • If it happens on main: it's an independent correctness fix and should be its own small PR with its own test. That also keeps this PR's PP numerics story clean.
  • If it only happens with this PR: say so in the comment. Otherwise a reader can't tell why chunked loss requires a PP change.

---

2. fsdp_lifecycle.py: separate file, merge into fsdp.py, or a fsdp/ package?

Recommendation: keep it a separate file and don't create a package yet. Make the file's purpose explicit.

• Two different kinds of code. fsdp.py is parallelize-time code: it applies fully_shard, picks meshes and reshard policies, and imports models.common.linear and models.common.moe. fsdp_lifecycle.py is runtime code that manipulates FSDP2 private state, which is the part you want to shrink and eventually delete. Keeping it in its own file means:
  • every private-API access in this PR sits in one ~130-line file, and loss.py no longer touches FSDP internals (good, matches the earlier §2.2 ask);
  • the upstream exit path is one file to delete, not scattered hunks in a 518-line fsdp.py;
  • it has no dependency on torchtitan.models. The dependency direction is already right: fsdp.py imports complete_fsdp_grouped_forward from it.
• Why not a fsdp/ package: two files don't justify it. Turning torchtitan/distributed/fsdp.py into a package means churning or re-exporting every torchtitan.distributed.fsdp import across models and experiments. The other private-FSDP users (fsdp.py:410-416, rl/generator/vllm/model_wrapper.py, experiments/torchft/manager.py) are unrelated and shouldn't be pulled in by this PR. Revisit if a third runtime FSDP-state helper shows up in core.
• Small follow-ups in the file:
  • Add a module docstring, e.g. "Runtime FSDP2 lifecycle helpers that depend on private FSDP state; each function names the upstream API that would replace it."
  • Move the TODO: Use public FSDP APIs for grouped eval cleanup from fsdp.py:52 onto complete_fsdp_grouped_forward, since that's now where the private access lives.
  • Give each function a TODO(<link>) that maps to the upstream items: (1) scoped partial backward, (2) value-based policy snapshot/restore with HSDP all-reduce decoupled, (3) grouped-forward completion.

---

3. fsdp_lifecycle.py:41: is chunked_loss_fsdp really tied to chunked loss?

The mechanism isn't tied to chunked loss; only the vocabulary is. What the context manager actually does, for one FSDP module:

1. keeps it unsharded across several forward calls;
2. runs several partial backwards that stop at detached inputs;
3. stops those inner backwards from reducing, resharding or finalizing any other group in the root (is_last_backward=False, non-head groups masked);
4. hands the reduction back to the caller's policy on the last pass, unless pending grads mean it should coalesce into the outer backward;
5. restores every policy on exit.

That is exactly upstream item (1), "scoped partial backward". The same pattern applies to GraphTrainer's chunked loss, chunked MTP heads, or any detach-then-backward-a-submodule trick. Naming it after the eventual upstream concept lets a reader see what the private-state code stands in for, and makes the swap mechanical later.

Suggested renames:

current	suggested
chunked_loss_fsdp	fsdp_partial_backwards (or isolated_fsdp_partial_backwards)
lm_head param	module
decoder_requires_grad	outer_backward_follows. This is the real meaning: another backward through the root follows and owns finalization.
_ChunkedLossHooks	PartialBackwardHooks
before_last_chunk	before_last_backward
defer_head_reduction	defer_to_outer_backward

If you'd rather keep it chunked-loss-specific, that's defensible too, since the before_last_chunk contract is tailored to it. In that case the current name is honest, but the docstring should say so explicitly: "Assumes module is the only FSDP module run in the inner passes and that all inner passes precede a single outer backward."

A related readability nit from the earlier review that's still open: the context manager yields None when the head isn't FSDP, so loss.py checks if fsdp is not None twice inside the chunk loop (loss.py:736, loss.py:756). Yielding a no-op PartialBackwardHooks instead removes both branches, and the loop body then reads the same with or without FSDP.

---

Summary:

1. The PP cleanup should eventually live in upstream PipelineStage; keep it here with a TODO, and confirm whether it's actually a separate fix that's reachable on main.
2. Keep fsdp_lifecycle.py as a separate file (no package), with a module docstring and upstream TODOs.
3. The logic is generic, so I'd rename it after the "scoped partial backward" concept rather than chunked loss.

I only posted review comments; I made no code changes.

[tianyu-l]

@ZenAlexa what claude says makes sense to me. Could you help address the final nits?

[tianyu-l]

@claude review again

also check if the CI failure is introduced by this PR

[claude]

Claude finished @tianyu-l's task in 3m 33s —— View job

---

Request changes -- this PR causes the CI failure: the DeepSeek-V3 and GPT-OSS PP goldens fail here and pass on main (#5664). Review

[claude]

Verdict: Request changes -- This PR causes the CI failure: the two PP model goldens fail on this head and pass on main. That numerics change is still unexplained, and the golden files haven't been regenerated.

What changed since the last review (64646a2). The commit is mostly mechanical:

• renames (fsdp_partial_backwards, outer_backward_follows, _PartialBackwardHooks, before_last_backward);
• a no-op hooks object for non-FSDP heads, which removes the two if fsdp is not None branches in loss.py;
• a module docstring, and the grouped-forward TODO moved onto complete_fsdp_grouped_forward;
• a reworded PP cleanup comment.

None of these change behavior.

CI failure: caused by this PR. The failing job is 8 GPU Integration (Real PG, models), run on head 64646a2. Two tests fail the exact-match golden check:

test	first mismatch	example
deepseek_v3_fsdp+cp+pp+ep	grad_norm at step 6, loss at step 7	loss 4.0435534 vs golden 4.0435486
gpt_oss_pp+fsdp+cp+ep+sacop	grad_norm at step 3, loss at step 4

Every other step is bitwise equal, and llama3_fsdp+tp+pp passes in the same job. The evidence that this PR is the cause:

• On #5664, the latest commit on main, the same job passes both tests with the same golden files ("All losses are equal ... Assertion passed!"). That job fails only on an unrelated deepseek_v4_dist_muon_fsdp+ep config error ('TrainingConfig' object has no attribute 'disable_cuda_graphs').
• The same two tests failed the same way on this PR's earlier 142f59599.
• Both failing configs use Interleaved1F1B with 8 microbatches and use ChunkedLossWrapper (it is the debug model default). That is the path where this PR changes how lm_head gradients are grouped into reductions; see the inline comment on fsdp_lifecycle.py:86-108.

The size and onset of the drift (~1e-6 relative, appearing after a few steps) are consistent with a change in summation order, not a gradient bug. Unverified: that it is purely an ordering effect. A same-environment main vs. head run on these two configs would confirm it.

Must fix

1. fsdp_lifecycle.py:86-108 -- Under PP, the policy snapshot captures a value left over from the previous microbatch, and that leftover changes how head gradients are grouped into reductions, so the PP goldens fail. Either declare it an intended numerics change and regenerate deepseek_v3_cp_pp.txt / gpt_oss_pp.txt (and every affected backend), or keep main's per-microbatch head reduction under PP (see inline).
2. fsdp_lifecycle.py:65-66, pipeline_parallel.py:519-520 -- These TODOs link to comments on this PR instead of upstream issues. File the three PyTorch issues and link them (see inline).

Still open from earlier reviews

• §2.1 from the 10-27 review: always isolate non-head groups when an outer backward follows, instead of reading _post_backward_pending / _partial_reduce_output. It wasn't adopted and there's no reply. That would remove the two private reads, and it is also the natural place to make the PP behavior deliberate (must-fix 1).
• test_loss.py's _RecordingFSDPLinear still mocks private FSDP fields with SimpleNamespace. Now that the real-FSDP lifecycle test exists, drop the policy assertions from the mock.

Resolved

• Coverage for the weight-tied 3-module group: added to both the lifecycle and pipeline matrices.
• The multi-GPU lifecycle test runs at head: Multi-GPU Unit Tests is green.
• Naming, the no-op hooks object and the module docstring, as requested in the 10-28 review.
• Scope of the PP cleanup: the author confirmed it is needed only because of this PR, and the comment now says so.

Should consider

• fsdp_lifecycle.py:103 -- before_last_backward is invoked before the last forward (inline nit).

Not verified: that the PP drift is purely a summation-order effect. A maintainer still needs to review, since this PR touches core loss and PP code.

[claude]

This is where the PP golden failures come from. Under PP, the loss runs inside the last stage's forward step, before the schedule sets the backward policy for that microbatch. So the "caller policy" snapshotted at line 75 is whatever the previous stage backward, or the previous step's finalization, left behind. That value then decides whether defer_reduction holds and whether the head groups reduce at the final chunk.

On main, the head reduce-scatters on every microbatch. With this PR, the number of head reductions per step depends on the schedule, which matches your own CPU 1F1B measurement of 5 -> 3. The result is still mathematically correct, but the fp summation order changes, and Interleaved1F1B with 8 microbatches makes that visible from step 6 onward.

Please choose one of the following:

• treat it as an intended numerics change: say so in the description, and regenerate deepseek_v3_cp_pp.txt and gpt_oss_pp.txt (plus any other affected backend goldens);
• or keep main's per-microbatch head reduction when manual_backward_finalization / a PP schedule owns the step. That leaves PP numerics unchanged and confines the change to the non-PP accumulation path.

Either way, under PP the "preserve caller policy" here is preserving a leftover value rather than a policy anyone chose for this microbatch. The docstring should say that.

[ZenAlexa]

manual finalization now owns head reduction from the first microbatch for every schedule, and parameter residency follows that same owner. the accumulation-order change is explicit in the description. backend-specific golden regeneration and the current memory, throughput and model checks remain open for the project runners.

[tianyu-l]

@claude review again

[claude]

Claude finished @tianyu-l's task in 1m 17s —— View job

---

Request changes: the only new commit is a merge of main, so the PP golden decision and the upstream-issue TODOs are still open. Review

[claude]

Verdict: Request changes -- The only new commit since the 10-30 review is a merge of main. It leaves the lifecycle code untouched, so both must-fix items from that review are still open. The PP numerics change is still undecided, and the golden files are unchanged.

What changed since the last review (64646a2 -> e002910). There is one commit, e002910, an ordinary merge of main@d86c77d (#5661, TP cross-entropy).

• fsdp_lifecycle.py, pipeline_parallel.py, fsdp.py and the tests are byte-identical to 64646a2.
• loss.py against main is the same hunk as before (chunk loop wrapped in fsdp_partial_backwards), and the merge resolved cleanly.
• I found no new problems in the merge.

CI. Lint, CPU Unit Test and GPU Unit Tests at e002910 are waiting for workflow approval (action_required). The 8-GPU model job hasn't run on this head. The latest model result is still the one from 64646a2: deepseek_v3_fsdp+cp+pp+ep and gpt_oss_pp+fsdp+cp+ep+sacop fail their goldens, and both pass on main.

Must fix (still open)

1. fsdp_lifecycle.py:86-108 -- PP golden failures. The description now confirms the mechanism on CPU: 1F1B goes from 5 to 3 reduce-scatters per step, with a gradient difference of up to 2.98e-8. That makes this an expected numerics change under PP, not a hypothesis. The same-environment GPU rerun would quantify it but won't remove the need to choose one of these:

   • (a) declare it intended and regenerate deepseek_v3_cp_pp.txt / gpt_oss_pp.txt, plus every other affected backend;
   • (b) keep main's per-microbatch head reduction when a PP schedule owns finalization.

   Under PP, the snapshotted "caller policy" is whatever the previous microbatch left behind, so (b) loses nothing intended. @tianyu-l, this one is your call.

2. fsdp_lifecycle.py:65-66, pipeline_parallel.py:519-520 -- The TODOs still link to comments on this PR. The description's "upstream FSDP API boundaries" section lists three proposals but no filed issues. Please file them in pytorch/pytorch and link them:

   • scoped partial backward;
   • value-based policy snapshot/restore, with set_requires_gradient_sync decoupled from all_reduce_grads;
   • PipelineStage resetting FSDPModule iteration state after metadata inference.

Still open, no author reply

• §2.1 from the 10-27 review: mask non-head groups unconditionally when an outer backward follows, instead of reading _post_backward_pending / _partial_reduce_output. Option (b) above fits naturally there.
• test_loss.py:974-1007 -- _RecordingFSDPLinear still rebuilds FSDP's private layout with SimpleNamespace and asserts reduce_grads on it. The real-FSDP test in test_chunked_loss_fsdp.py already covers policy restoration, so trim the mock back to the unshard and forward counts.
• fsdp_lifecycle.py:103 -- before_last_backward runs before the last forward, so the name is a nit.

A maintainer still needs to review, since this touches core loss and PP code.

Not verified: the GPU size of the PP drift at e002910. That needs the model job, which is waiting for approval.

[tianyu-l]

@claude explain the PP numerics change with easy-to-understand examples before vs. after

1F1B goes from 5 to 3 reduce-scatters per step

sounds that we should do it?

[claude]

Claude finished @tianyu-l's task in 3m 38s —— View job

---

Short answer: reducing lm_head once per step under PP is worth doing, but this PR only gets there by accident. The "5 -> 3" comes from whatever FSDP policy the previous microbatch left behind, so the number of head reductions depends on the schedule (GPipe: no change). I'd keep main's behavior under PP in this PR and make "head reduces once per step under PP" its own small PR, with memory numbers and regenerated goldens.

Why PP is different

Under PP, the loss runs inside the last stage's forward action. That covers ChunkedLossWrapper and its head-only chunk backwards. The FSDP policy for a microbatch's backward is set later: the stage backward turns sync off, and PP turns it back on only when it finalizes at the end of the step. So the "caller's policy" the wrapper snapshots (fsdp_lifecycle.py:75-85) is whatever the previous action left behind.

Example. One stage holds decoder plus the grouped (norm, lm_head), with 1F1B and 3 microbatches. The action order is F0 B0 F1 B1 F2 B2, then finalize. This is the setup behind the 5 -> 3 count, the same shape as test_chunked_loss_fsdp.py:128-156.

action	policy the wrapper sees	main: head group	PR: head group
F0	sync on (left by last step's finalize)	reduce-scatter mb0	reduce-scatter mb0
B0	PP turns sync off
F1	sync off (left by B0)	forces sync on, reduce-scatter mb1	keeps it off, accumulates locally
F2	sync off (left by B1)	forces sync on, reduce-scatter mb2	accumulates locally
finalize	PP turns sync on	head group (norm's mb2 grad) + decoder	head group (head mb1+mb2, norm) + decoder
reduce-scatters per step		5	3

Why the bits change. The math is the same, but the order of additions differs. Let RS be the reduce-scatter (sum over ranks) and g_k the head gradient of microbatch k:

main: shard = RS(g0) + RS(g1) + RS(g2)   # sum over ranks, then over microbatches
PR:   shard = RS(g0) + RS(g1 + g2)       # add mb1+mb2 locally first, then over ranks

Float addition isn't associative. In fp32, (1 + 1e8) - 1e8 = 0, but 1 + (1e8 - 1e8) = 1. Each step's difference is about 1 ulp (the CPU run in the description saw 2.98e-8), and the optimizer compounds it. That fits the goldens matching for a few steps and then diverging: step 3 for gpt_oss, step 6 for deepseek_v3. Unverified: that the GPU drift comes only from this.

Should we do it?

The goal, yes. On main, lm_head is the only parameter in a PP stage that reduces every microbatch. Everything else waits for PP finalization, and lm_head is usually the largest single weight (Llama3-8B: 128256 x 4096, about 525M params). But the PR doesn't implement that goal deliberately:

• The count depends on the schedule. GPipe runs every forward before any backward, so every forward sees "sync on" and the head still reduces per microbatch. That is why the description's GPipe run is bitwise equal to main. 1F1B reduces at F0 and again at finalize. Interleaved1F1B falls in between, depending on how many last-stage forwards run before its first backward.
• The policy it "preserves" was never chosen. Under PP it is the previous action's leftover, so no comment or test can state the intended behavior.
• The memory cost isn't measured. Deferring keeps the head's unsharded gradient alive across microbatches on the last stage. It also keeps the unsharded weight: the leftover reshard_after_backward=False skips the reshard at fsdp_lifecycle.py:140-146. Main releases both every microbatch. For Llama3-8B that is roughly 1 GiB of bf16 weight plus 1-2 GiB of gradient per rank (unverified on GPU). It also removes the per-microbatch head all-gather, which is a throughput gain worth measuring.

Recommendation

1. This PR: keep main's per-microbatch head reduction when PP owns finalization (state_ctx.manual_backward_finalization). The one exception is the HSDP case with deferred replica all-reduce, which the PR has to defer anyway. PP goldens stay bitwise equal, and the PR stays about non-PP policy preservation. By the author's own reply on pipeline_parallel.py:517, the metadata-inference cleanup is only needed because of this PP deferral, so that hunk could probably move to the follow-up too. Unverified: that no other path needs it.
2. Follow-up PR: make the deferral deliberate. Under manual finalization, always defer the head. Roughly, drop the any(not policy.all_reduce_grads ...) clause at fsdp_lifecycle.py:88-94, so the head reduces exactly once per step for every schedule. Include last-stage peak memory and throughput on the 8-GPU PP configs, and regenerate deepseek_v3_cp_pp.txt / gpt_oss_pp.txt (and any other affected backends) with that explanation.

If you'd rather land it in one PR, the minimum is option 2 here, deterministic rather than leftover-dependent, plus the memory numbers and regenerated goldens. Either way it's your call. The PR's current form is the only one I'd rule out.

Assumptions behind the example

• These come from reading the code and from memory of upstream PipelineStage: the stage backward calls set_requires_gradient_sync(False) / set_reshard_after_backward(False) / set_is_last_backward(False), and finalization turns sync back on. I couldn't import torch in this runner to check this against the pinned nightly.
• The 5/3 breakdown is my reconstruction. It matches the description's totals for single-stage 1F1B with 3 microbatches (no warmup forwards). That comparison script isn't checked in.
• On main, F1's forced head reduce-scatter also picks up norm's pending mb0 gradient, because norm is in the same group. I left that out of the table for simplicity. It doesn't change the counts.

[tianyu-l]

@claude

What still puzzles me is "reduce-scatter mb0" on F0. My impression has been that, without chunked loss wrapper, every FSDP module will only incur one reduce-scatter across all the PP microbatches, after the last backward happens. Correct me if I'm wrong.

It seems it's not the case with chunked loss wrapper, even with changes in this PR. Regardless of whether or not we fix in this PR (I think we should), is there a way that we defer this mb0 reduce-scatter on lm_head as well?

[claude]

Claude finished @tianyu-l's task in 2m 38s —— View job

---

You're right that without the wrapper each FSDP group reduce-scatters once per PP step. The mb0 reduce-scatter on lm_head happens because the wrapper runs lm_head's backward inside the last stage's forward action, and at F0 the stage hasn't switched FSDP to no-sync yet. Yes, it can be deferred: when PP owns finalization, always defer the head (a one-clause change in fsdp_lifecycle.py:86-101), and I'd do it in this PR.

Why there's normally one reduce-scatter per step

PyTorch's PipelineStage switches its FSDP root to no-sync at the start of every backward action (reduce_grads=False, reshard_after_backward=False, is_last_backward=False). The step's finalization then turns sync back on and runs post_backward on every group once. So each group reduces once per schedule step. With defer_fsdp_gradient_reduction, that becomes once per optimizer step, because finalize_gradients is False on earlier accumulation groups (training_engine.py:782-784).

Why lm_head breaks that pattern under chunked loss

The loss runs inside forward_one_chunk, and ChunkedLossWrapper calls chunk_loss.backward() there (loss.py:819-821). So lm_head's backward happens during F0, before B0 has set no-sync. At that point, FSDP still holds what the previous step's finalization left behind: sync on and reshard on. On step 1 it holds FSDP's defaults, which are the same.

action	lm_head policy at that moment	main	this PR
F0	sync on (left by previous finalize)	forces sync on: RS	keeps sync on: RS
B0	stage sets no-sync
F1, F2	no-sync (left by B0, B1)	forces sync on: RS each	keeps no-sync: accumulates
finalize	stage sets sync on	RS head group + decoder	RS head group + decoder

At F0, the PR's defer_reduction (fsdp_lifecycle.py:86-101) comes out False for plain FSDP, for two reasons:

• the manual_backward_finalization clause only fires when HSDP replica all-reduce is off;
• nothing is pending yet at the start of the step.

So before_last_backward restores reduce_grads=True, and the final chunk reduces mb0. Under GPipe, every forward runs before any backward, so every forward sees sync on. That's why GPipe gives the same result as main.

How to defer it

The stage already sets set_manual_backward_finalization(True) in forward_one_chunk, before the loss runs. You can see the same call mirrored at models/kimi_k3/pipeline_parallel/stage.py:178-180. So the wrapper can already tell at F0 that PP owns finalization. Defer unconditionally in that case:

defer_reduction = state_ctx.manual_backward_finalization or (
    outer_backward_follows
    and any(
        policy.group._post_backward_pending
        or policy.group._partial_reduce_output is not None
        for policy in policies
    )
)

Details: what this changes, and what to update with it

• The head reduces exactly once per step on every schedule (GPipe, 1F1B, Interleaved1F1B). The step's finalization reduces it together with the other groups, through the same post_backward path. The PR's HSDP + PP case already relies on that path, since it defers under manual finalization when all-reduce is off. The new clause includes that HSDP clause, so the HSDP clause can be deleted.

• The "leftover policy" problem goes away. Under PP, the head follows the stage's backward policy for every microbatch instead of whatever the previous action left. The docstring line "Manual finalization retains full gradients for deferred replica reduction" can become "Under manual (PP) finalization, the head's reduction is always deferred to the schedule's finalization."

• Counts. For the single-stage 3-microbatch example from the previous comment, reduce-scatters per step go from 5 (main) and 3 (this PR) to 2: one for the head group, one for the decoder. This is my reconstruction, so the expected counts in the PP lifecycle tests need updating.

• Reshard. Right now the early reshard at fsdp_lifecycle.py:140-146 follows the snapshot, so the head weight is resharded after F0 but kept unsharded after F1 onward. Pick one rule under PP:

  • always reshard: same weight memory as main, plus one head all-gather per microbatch;
  • never reshard: matches the stage's no-sync policy and drops those all-gathers.

  Either way, the unsharded gradient stays alive across the step. That memory cost is the price of deferring, and every other group on the stage already pays it.

• Numerics. This is an intended change in summation order: RS(g0+...+gN) instead of per-microbatch RS(g_k) sums. Regenerate every PP golden whose config uses ChunkedLossWrapper, with this explanation. Besides deepseek_v3_cp_pp.txt and gpt_oss_pp.txt, that probably includes llama3_fsdp+tp+pp, which passes today. For the PR description, report last-stage peak memory and throughput on the 8-GPU PP configs.

• The metadata-inference cleanup (pipeline_parallel.py:517-523) is still needed, because the deferred head grads from the probe backward must be discarded.

• Upstream angle. The underlying gap is that PipelineStage applies its no-sync policy only when a backward action starts, even though a loss computed in forward can run backward. If upstream applied it once at step start, or in forward_one_chunk next to set_manual_backward_finalization, the wrapper wouldn't need this special case. That fits the PP item in the upstream issue list.

Not verified: I couldn't import torch here, so I haven't checked the upstream PipelineStage no-sync and finalization behavior or the reduce-scatter counts against the pinned nightly. Both come from reading the code and from memory of upstream.

[tianyu-l]

@ZenAlexa could you fix the remaining gap in PP, so that whether chunked loss do grad sync aligns with the PP+FSDP setup?

[ZenAlexa]

yep, the first head pass now follows the schedule's manual-finalization mode too, so every head microbatch waits for the schedule's reduction and parameter residency stays with the schedule. this removes the dependence on the previous backward's flags for both GPipe and 1F1B. the resulting summation-order change is intentional, and the remaining model checks and backend-specific golden regeneration are listed in the description.

non-head groups are already masked for every inner pass when decoder backward follows, and that keeps their gradients and prefetch state with the outer backward. the pending-gradient check now applies only outside manual finalization, where it keeps the shared norm/head group's accumulated gradients together for one reduction. restoring head sync on every last chunk would split that group's reduction again, so the check stays within this single lifecycle owner.

[tianyu-l]

@claude review again

[claude]

Claude finished @tianyu-l's task in 2m 40s —— View job

---

Request changes: the code now defers the head under PP as requested; merge still needs regenerated PP goldens and last-stage memory and throughput numbers from a GPU run. Review

[claude]

Verdict: Request changes -- The lifecycle code now does what was asked: under PP, the head's reduction always waits for the schedule's finalization. What still blocks merge is evidence, not code. The PP goldens this intentionally breaks haven't been regenerated, and there are no last-stage memory or throughput numbers yet. Both need a GPU run that a maintainer has to trigger.

What changed since the last review (e002910 -> 636e16d). The real changes are in 049301b, which is a merge commit that also carries edits. 636e16d is a plain merge of main and leaves the PR files untouched.

• fsdp_lifecycle.py:88-97: when the schedule finalizes backward manually (manual_backward_finalization), defer_reduction is now unconditionally True. The HSDP-only clause is gone.
• fsdp_lifecycle.py:136-143: under manual finalization the wrapper no longer reshards the head after a backward. Residency follows the schedule's no-sync policy.
• The before_last_backward hook is renamed to before_last_forward. TODOs now link pytorch/pytorch#203176, #203177 and #203178.
• The SimpleNamespace mock tests in test_loss.py are deleted and replaced by a real-FSDP multi-output scenario. The PP matrix adds a plain-FSDP axis and asserts reduce-scatter counts at finalization.

Example. One stage holds decoder plus the grouped (norm, lm_head), runs 1F1B, and takes 3 microbatches (counts are reduce-scatters per step):

            F0        F1        F2        finalize      total
main        head RS   head RS   head RS   head+decoder  5
e002910     head RS   accum     accum     head+decoder  3   (GPipe: same as main)
636e16d     accum     accum     accum     head+decoder  2   (every schedule)

The expected_reductions check in test_chunked_loss_fsdp.py:191-199 asserts exactly this count (2, or 1 with a frozen decoder) for both GPipe and 1F1B.

Resolved since the last review

• PP head deferral now depends on the schedule, not on the policy the previous action left behind (must-fix 1, code side).
• The TODOs link upstream issues instead of comments on this PR (must-fix 2).
• The _RecordingFSDPLinear private-state mock is removed.
• The before_last_backward naming nit is fixed.
• §2.1 (pending-gradient reads): the author explained that outside PP, the check keeps a shared [norm, lm_head] group's accumulated gradients in one reduction, and there was no pushback. Accepted.

Must fix

1. Regenerate the PP goldens in this PR. The summation-order change now covers every PP + FSDP config that uses ChunkedLossWrapper, not just the two that failed before. llama3_fsdp+tp+pp passed on e002910 only because its schedule happened to reduce per microbatch. Expect it to drift now too (unverified). Per CLAUDE.md principle 6, every affected backend (A10G and MI350X) needs regenerated files, with the explanation already in the description. A maintainer needs to trigger the 8-GPU model job at 636e16d.
2. Report last-stage peak memory and throughput on an 8-GPU PP config before and after.
   • Cost: the head's unsharded weight and accumulated gradient now stay alive from F0 until finalization. With finalize_gradients=False, they stay alive across schedule steps. Main resharded and reduced the head every microbatch.
   • Gain: the per-microbatch head all-gather and reduce-scatter go away.
   • This is the intended trade, and it matches what PP's no-sync already does to every other group on the stage. But it is a memory change on the main PP path, so it needs numbers.

Should consider

• test_chunked_loss_fsdp.py still reports failures with raise RuntimeError(...) (for example at lines 195-199). self.assertEqual or torch.testing.assert_close would show diffs. This was raised earlier and is still open.
• The PR description has grown into a revision-by-revision log, with evidence from e002910, 81b3eec, 142f595, 65b1173 and 1bbb64b. Trimming it to the current head's behavior, the before/after table above and the pending GPU gates would make it easier to verify.

A maintainer still needs to review, since this touches core loss and PP code.

Not verified:

• that the upstream PipelineStage sets reshard_after_backward=False on non-final backwards (the memory claim in must-fix 2 rests on this);
• whether llama3_fsdp+tp+pp now drifts;
• the CPU, GPU and Lint runs at 636e16d, which were still queued or running at review time.