[FSDP][Loss] Switch gradient reduction to SUM#1963
Draft
HAOCHENYE wants to merge 1 commit into
Draft
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
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