[FSDP][Loss] Switch gradient reduction to SUM#1963
Closed
HAOCHENYE wants to merge 1 commit into
Closed
Conversation
Flip gradient reduction from mean to pure SUM, so each parameter receives the sum of per-rank local-component gradients, i.e. the global loss gradient, with no compensating divides. This is the single atomic semantic change; the three cancelling x world_size factors and the reduce divides are removed together. - BaseModel/MoE/Dense.fully_shard call set_gradient_reduce_sum() at the end. - MoE.scale_and_reduce_grad drops the expert div_(ep_size) and the replicated div_(flat_mesh.size()); only the coalesced SUM all_reduce remains. - CE: drop the WORLD autograd all_reduce; the loss stays this rank's local component (display global value restored by the C2 detached pipeline). - Balancing: use local_gating_sum directly instead of all_reduce_autograd; the global detached statistics (tokens_global/seqlen_global/scale_global) are kept unchanged. - Z-loss: drop the x world_size in the global-average branch. Verified on torch 2.10 (bf16 force-sum): distributed full gradient reproduces a single-process full-batch token-mean CE reference at EP=1 and EP=2 (norm-ratio median 0.9998, EP-invariant); global display loss unchanged through the flip; balancing+z backward stays finite. Regression tests cover the bf16 reduce-sum mechanism, token-mean parity with grad-acc=2, aux-loss finite gradients, and an isolated fp32 balancing-only gate-gradient A/B (new local+SUM vs old all_reduce_autograd+AVG) asserting element-wise equality.
HAOCHENYE
force-pushed
the
reduce-sum-02-split-logging
branch
from
July 17, 2026 08:11
a13bca8 to
316b1cc
Compare
HAOCHENYE
force-pushed
the
reduce-sum-03-switch-sum
branch
from
July 17, 2026 08:11
51c5491 to
590c9e7
Compare
HAOCHENYE
commented
Jul 20, 2026
| loss = loss * self.loss_cfg.z_loss_alpha / self._batch_size | ||
| self._update_running(loss.detach()) | ||
| self._update_local_running(local_loss.detach()) | ||
| self._update_local_running(loss.detach()) |
Collaborator
Author
There was a problem hiding this comment.
If the changes in PR 1962 are necessary, it is recommended to merge PR 1962 and PR 1963 into a single one, as it makes no sense for a function defined in PR 1962 to be called in PR 1963 — it feels fragmented
| # Safety net: every `*loss` tensor field an output exposes must have produced a display curve | ||
| # from a registered local component. Otherwise a newly added loss term would silently vanish | ||
| # from the logged curves while still driving backward, hiding the regression. | ||
| for output in batch_outputs: |
Collaborator
Author
There was a problem hiding this comment.
Encapsulate it into an internal function; subclasses can override it if needed
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 (3/5) — the single atomic semantic flip. FSDP reduce-scatters with pure SUM;
scale_and_reduce_graddrops the expertdiv_(ep_size)and replicateddiv_(flat_mesh.size()); CE drops the WORLD autograd all_reduce; balancing useslocal_gating_sumwith global detached stats; z-loss dropsx world_size; compose/VLMfully_shardalso wired to SUM. Carries all regression tests (bf16 mechanism, token-mean parity EP1/EP2, balancing & z-loss new-vs-old A/B, compose hook).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