[Loss] Split logging loss from backward loss#1962
Closed
HAOCHENYE wants to merge 1 commit into
Closed
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
HAOCHENYE
commented
Jul 20, 2026
| loss = self._get_total_loss(output) | ||
| loss.backward() | ||
| total_loss += loss.detach() | ||
| local_display_loss += self.model.local_display_loss(output) |
Collaborator
Author
There was a problem hiding this comment.
loss reduction, otherwise still do it in post_micro_batch_forward?
| local_z_loss = z_locals[0] if len(z_locals) == 1 else torch.stack(z_locals).sum(dim=0) | ||
|
|
||
| return balancing_loss, z_loss, tokens_per_expert_global | ||
| return balancing_loss, z_loss, tokens_per_expert_global, local_balancing_loss, local_z_loss |
Collaborator
Author
There was a problem hiding this comment.
Return a TypedDict, as the tuple return type is not convenient for extension
Comment on lines
+250
to
+251
| if dist.is_initialized(): | ||
| dist.all_reduce(local_display_loss, op=dist.ReduceOp.SUM) |
Collaborator
Author
There was a problem hiding this comment.
This part of the logic is placed in post_micro_batch_forward, and I hope each rank can display its own loss without performing an all_reduce sum. Could you try multiplying this coefficient back instead of doing a global sum?
Comment on lines
+182
to
+184
| loss = loss_vec.sum() * alpha / self._batch_size | ||
| local_loss = local_vec.sum() * alpha / self._batch_size | ||
| return loss, local_loss.detach() |
Collaborator
Author
There was a problem hiding this comment.
Why return two losses? If a subsequent module needs to record it, just detach it there. Why implement it this way? The interfaces between modules are all very strange now
| non_pad_token=non_pad_token, | ||
| ) | ||
| balancing_loss, z_loss, tokens_per_expert_global = split_aux_output | ||
| balancing_loss, z_loss, tokens_per_expert_global, local_balancing_loss, local_z_loss = split_aux_output |
Collaborator
Author
There was a problem hiding this comment.
If local_loss can be omitted, please simplify the interface
HAOCHENYE
commented
Jul 20, 2026
Collaborator
Author
There was a problem hiding this comment.
I don't see the need to make so many changes just for printing loss
Collaborator
Author
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