[Loss] Split logging loss from backward loss#1962
Draft
HAOCHENYE wants to merge 1 commit into
Draft
Conversation
Restore every displayed loss scalar from a detached per-rank LOCAL component reduced across ranks with a SUM all_reduce, instead of mean-reducing the (globally reduced) backward loss tensors. Because each term's cross-rank SUM equals the global loss, the logged curves are unchanged while backward still uses the current global autograd path; this decouples display from backward so the reduce-sum switch can make backward per-rank-local without touching logs. - BalancingLossContext.finalize / ZLossContext now also produce a detached local component (z-loss's without the x world_size factor); MoE forward records them on extra_info as local_*loss (CE already exposes local_base_loss, MTP mirrored). - BaseModel.post_micro_batch_forward and train_step total_loss reduce these local components with detached SUM; train_step total_loss is now the global value. Behavior-neutral: EP=2, balancing+z on, grad-acc=2, symmetric data reproduces C1's total_loss / reduced_llm_loss / reduced_balancing_loss / reduced_z_loss bit-for-bit.
HAOCHENYE
force-pushed
the
reduce-sum-01-fsdp-helper
branch
from
July 17, 2026 08:11
c216448 to
184b6ca
Compare
HAOCHENYE
force-pushed
the
reduce-sum-02-split-logging
branch
from
July 17, 2026 08:11
a13bca8 to
316b1cc
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR (2/5) — behavior-neutral. Routes displayed/logged losses through detached per-rank local components reduced by cross-rank SUM; the backward loss is still globally reduced. Isolates the display pipeline so the next PR flips gradients without touching logging.
Full stack (merge bottom-up, under #1959):
reduce-sum-01-fsdp-helper— [FSDP] Add reduce-sum gradient reduction helperreduce-sum-02-split-logging— [Loss] Split logging loss from backward lossreduce-sum-03-switch-sum— [FSDP][Loss] Switch gradient reduction to SUMreduce-sum-04-remove-world-size— [Refactor] Remove unused world_size plumbingreduce-sum-05-drop-nonglobal— [Refactor] Drop non-global aux-loss averaging mode