[sanketpurandare]

Prepared with AI assistance.

Loss normalization belongs to stages that evaluate the objective. Require
denominators on those owners, and discard unused denominators on other ranks
before device conversion or optimizer-step mutation. Preserve existing RL,
TorchFT and direct engine callers instead of requiring an unrelated API change.
Routing denominators remain globally reduced on every MoE stage.

Keep token-seen and auxiliary metric collectives on all ranks; reduce and log
loss/count metrics only on owners. Non-owners retain the existing explicit
loss sentinel. No throughput improvement is claimed for the smaller payload.

Verification

Owning-commit CPU checks: 116 passed, 2 warnings in 17.02s.

Strict relevant Sphinx builds and all non-type repository hooks passed.
Three unchanged main type findings remain attributed, not reported as a clean
full type run. External reviewer and CI acceptance are separate, still pending.

Complete-stack integration (not this PR in isolation): real four-GPU PP2/DP2/
EP2 and PP4/DP1/EP1 ran ten updates with exact model/optimizer tensor parity,
optimizer graphs and sequential validation. A fresh native PP2/PP4 checkpoint
matrix matched resumed state exactly. These are debug-model contract checks,
not 671B throughput or convergence results.

[tianyu-l]

I think we should do either of

• only handle loss related logic when owns_loss=True (need to be careful to not hang in collective due to non-spmd behavior)
• or continue to let every rank handle loss -- I used to see some codebase sending final loss to all PP ranks.

We shouldn't use the name global_token_counts but fills it implicitly with local values.

[sanketpurandare]

Loss counts are now absent on non-loss stages, rather than local values under a global name. Loss-owner DP/CP peers make the same reduction decision, while routing denominators still reduce on every MoE stage.

[tianyu-l]

I guess the purpose of this PR is to save this all-reduce on non last-stage PP ranks, thus improving throughput? Not clear how much it'd bring from this tiny collective.

[sanketpurandare]

This skips loss-only payload and validation reduction work on non-owners, not the routing collective. I removed the implied throughput claim; I haven't measured a TPS improvement attributable to this tiny collective alone.

[tianyu-l]

@claude review

[claude]

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

---

Verdict: minor changes requested. The design follows option 1 from the earlier review (owns_loss-only loss logic) and the training collectives stay SPMD-safe, but the validator's zero-token error now fires on loss owners only, and the torchft copy of train_step still has the old naming pattern. Three inline comments: validator, engine, MTP.md.

Details

What I checked and found safe

• Collective consistency: the dp and loss meshes are both per PP stage (parallelism_context.py:339,387), so every rank in a collective makes the same owns_loss decision. The payload is smaller on non-owners, but within each mesh it is the same everywhere. Moving the ntokens_seen dist_sum ahead of the loss reductions keeps the same order on all peers.
• CUDA graphs: None is a pytree node with no leaves, so CUDAGraphInputSpec sees a stable structure on non-owners (always None) and on owners (always a tensor).
• Loss kwargs on non-last stages: only loss functions read global_loss_token_counts (loss.py, MTPLoss, dapo, ChunkedLossWrapper), and the PP schedule never calls them on non-last stages.
• GraphTrainer GraphPP asserts that every rank holds both the first and last stage, so its torch.Tensor-typed overrides never receive None.
• Logging: non-owners now log -1.0 for train and validation loss. Before, they logged dist_sum(-1 sentinel), which scaled with the DP/CP size, so this is cleaner.

Issues

1. Validator zero-token check (inline): under PP only loss owners raise this error. Non-owners go back to training and block on the next PP or FSDP communication until the job is torn down.
2. Callsite audit (inline on engine): experiments/torchft/trainer.py:390-416 still passes local loss counts named as global on non-owners, and the engine silently drops them. Either update the FT copy and remove the override, or explain why the compatibility shim should stay. Separately, not pp_enabled or engine.pp_has_last_stage can be just engine.pp_has_last_stage in both the trainer and the engine.
3. MTP.md (nit): the new text is trainer and collective detail inside a model doc, and the line wrap is off.

Tests

• The CPU tests are heavily mocked but cover the owner/non-owner × DP × should_log combinations. The engine test case (pp_enabled=False, last_stage=False) is a state that can't occur, because the engine forces pp_has_last_stage=True without PP.
• I did not run the tests or lint.

On the value question from the earlier review: the PR now states it claims no throughput gain, so the case for merging rests on semantic clarity. With +304/-70 lines, mostly tests, that's a maintainer call.

[claude]

This failure path is no longer SPMD under PP. On the base branch every PP stage built total_global_valid_tokens from the same per-DP-rank data, so all ranks raised this ValueError together. Now only loss owners raise it. Non-owners log -1.0, return to train(), and block on the next PP P2P or FSDP collective until the job is torn down (NCCL timeout, or the torchrun agent killing the group). Under PP, rank 0 is a non-owner, and that is the rank whose logs people usually read, so the actual error is easy to miss.

This is an edge case (fully masked validation labels), but it is the "careful to not hang in collective due to non-SPMD behavior" concern from the earlier review. Some options:

• keep the per-step DP count reduction on every rank, which is a scalar and cheap, and do the zero check everywhere; or
• accept the asymmetry and add a comment saying non-owners can't detect this case.

[claude]

Two things here.

1. engine.pp_has_last_stage is already True when PP is disabled (training_engine.py:408), so not self.parallelism_context.pp_enabled or self.pp_has_last_stage reduces to self.pp_has_last_stage. This file already uses the bare form at L743/L880/L891, and so does rl/trainer/trainer.py:303. The same applies to trainer.py:365. Only the validator needs the long form, because its pp_has_last_stage is None without PP.

2. Overwriting a caller's tensor with None on non-owners means the engine still accepts the old pattern this PR is removing: local counts passed under a global name. experiments/torchft/trainer.py:390-416 still has that exact block (and computes global_main_loss_token_count from it), so the naming concern from the earlier review is still there in the duplicated FT train step. I'd either update the FT copy to match core and drop the override, keeping only the owner-side ValueError, or explain why this compatibility shim should stay. Discarding the input saves only a tiny H2D copy of a value that is never read on non-last stages.

[claude]

nit: this is trainer and collective bookkeeping ("reduces its loss payload, not its collective count", the -1 logging sentinel) inside the MTP model doc. That belongs in the PR description, or a comment in train_step. Also, Each prediction is now left on its own short line. Something like this would do:

Suggested change

	vectors over the optimizer step. Routing counts reduce on every MoE stage;
	loss counts reduce only on loss-owning pipeline stages. Non-owners pass no
	loss denominator (`None`) to the engine and retain a `-1` logging sentinel;
	local counts are not substituted for a global loss denominator. The routing
	count collective still runs on every MoE stage during training; this reduces
	its loss payload, not its collective count. Validation skips loss-only
	reductions on non-owners. Each prediction
	vectors over the optimizer step. Routing counts are reduced on every MoE
	pipeline stage; loss counts are reduced only on stages that evaluate the loss.
	Each prediction