[Refactor] Drop non-global aux-loss averaging mode#1969
Closed
HAOCHENYE wants to merge 2 commits into
Closed
Conversation
Reduce every parameter's gradient with pure SUM instead of mean, so each param receives the sum of per-rank local-component gradients (the global-loss gradient) with no compensating divides. This merges the previously-split helper, logging, switch, world_size-cleanup and fp32-reduce changes into one atomic, self-consistent commit (the gradient flip and the three x world_size removals must land together). Gradients: - BaseModel.set_gradient_reduce_sum() pairs set_gradient_divide_factor(1.0) with set_force_sum_reduction_for_comms(True); wired into every fully_shard (BaseModel / MoE / Dense / compose base / InternS1). - scale_and_reduce_grad: drop expert div_(ep_size) and the replicated pre-divide; BaseModel now SUM-all_reduces (no divide) Replicate-placement / fp32 ignored-param grads (fixing Dense + compose fp32 params that got no reduction). MoE keeps its override. - Loss: drop the three x world_size injections (CE WORLD all_reduce, balancing all_reduce_autograd, z-loss x world_size) and the now-dead world_size plumbing (ZLossContext / AuxLossContext.accumulate args, _AllReduce/all_reduce_autograd). - Fix the domino-EP MTP crash (stale world_size=z_world_size). Display (each rank shows its OWN loss, no all_reduce): - Each loss context exposes calibrate() -> detached per-rank display value (CE: per-token mean = local_loss x global_denom/local_grad_tokens; z-loss: this rank's raw z-loss mean; balancing: this rank's balancing loss from local stats, detached/display-only). The model forward stores them on output.calibrated_losses; post_micro_batch_forward (overridable) and train_step sum them per-rank with NO cross-rank all_reduce. Summed over the batch each rank logs its own loss, equal to the global loss at world size 1. - Remove the local_* record machinery; finalize returns only backward outputs; AuxLossContext.finalize returns a TypedDict. Verified on torch 2.10: EP1/EP2 token-mean parity (median 0.9998), balancing and z-loss element-wise old-vs-new gate A/B (rtol 1e-4), replicate+shard param reduce-sum, Dense fp32 ignored-param SUM, compose reduce-sum hook, MTP domino path, and per-rank display (distinct data -> per-rank values differ; identical data -> each rank == global). pre-commit clean.
Remove the `balancing_loss_global_average` / `z_loss_global_average` config switches and the non-global averaging branch they selected. Under reduce-sum gradients the global-average form is the only correct reduction, so the switches only offered an incorrect mode. The single-process (`dist.is_initialized()` is False) path is kept as the world-size-1 case for reference / eval: it is numerically identical to the distributed branch at world size 1.
This was referenced Jul 20, 2026
HAOCHENYE
force-pushed
the
reduce-sum-core
branch
from
July 22, 2026 09:20
5f4f804 to
3ce8fea
Compare
Collaborator
Author
|
Folded into #1968 (single squashed commit). The non-global aux-loss mode removal is now part of that PR. Closing. |
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.
Stacked on #1968.
Removes the
balancing_loss_global_average/z_loss_global_averageconfig switches and their non-global-average branch. That mode computes a full per-rank loss with local denominators, which is incompatible with reduce-sum — its replicated router/gate gradient would be inflated byworld_sizeunder SUM reduction. Thedist.is_initialized()-False single-process path is kept as the world-size-1 case (numerically identical to the global branch at W=1).Breaking config change: removes two config fields (
balancing_loss_global_average,z_loss_global_average).