[FSDP][Loss] Switch gradient reduction to SUM

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.

Author: HAOCHENYE

tests/model/test_reduce_sum_grad.py

@@ -6,7 +6,6 @@
 import torch.nn.functional as F
 from cyclopts import Parameter
 from torch.distributed.device_mesh import DeviceMesh
-from torch.distributed.nn.functional import all_reduce
 
 from xtuner.v1.utils.device import get_device
 
@@ -282,10 +281,10 @@ def forward(
 
         extra_info["local_base_loss"] = loss.detach().clone()
 
-        # Step 2.c in the loss calculation: reduce the loss over all ranks using all_reduce with autograd support
-        if dist.is_initialized():
-            loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD)
-
+        # Under reduce-sum gradients the loss stays as this rank's local component (local token sum
+        # over the global token denominator). Cross-rank aggregation happens on the gradients via the
+        # FSDP SUM reduce-scatter, so no autograd WORLD all_reduce is injected here. The global loss
+        # scalar for logging is restored separately by the detached display pipeline (§5.4).
         return loss, (logits, extra_info)
 
 

xtuner/v1/loss/ce_loss.py

@@ -155,22 +155,18 @@ def finalize(
         alpha = self.loss_cfg.balancing_loss_alpha
 
         if self.loss_cfg.balancing_loss_global_average and dist.is_initialized():
-            group = dist.group.WORLD
-            assert group is not None
             tokens_global = tokens_per_expert_global.sum(-1)
             seqlen_global = tokens_global // num_experts_per_tok
             scale_global = n_routed_experts / tokens_global
 
-            routing_weights_sum_global = all_reduce_autograd(local_gating_sum, "sum", group)
-            routing_weights_mean_global = routing_weights_sum_global / seqlen_global.unsqueeze(-1)
-            loss_vec = scale_global * (tokens_per_expert_global * routing_weights_mean_global).sum(-1)
-
-            # Detached local component: same global denominators, but this rank's own gating sum in
-            # place of the cross-rank sum. Because `all_reduce_autograd` is a plain SUM and every
-            # other factor here (scale_global, seqlen_global, tokens_per_expert_global) is detached
-            # and global, summing `local_vec` over ranks reproduces `loss_vec` exactly.
-            routing_weights_mean_local = local_gating_sum.detach() / seqlen_global.unsqueeze(-1)
-            local_vec = scale_global * (tokens_per_expert_global * routing_weights_mean_local).sum(-1)
+            # Under reduce-sum gradients the loss stays as this rank's local component: use this
+            # rank's own gating sum with the global denominators, without any cross-rank all_reduce.
+            # Cross-rank aggregation happens on the gradients (FSDP / scale_and_reduce_grad SUM);
+            # since every other factor here is detached and global, summing the loss over ranks
+            # reproduces the global balancing loss.
+            routing_weights_mean = local_gating_sum / seqlen_global.unsqueeze(-1)
+            loss_vec = scale_global * (tokens_per_expert_global * routing_weights_mean).sum(-1)
+            local_vec = loss_vec.detach()
         else:
             valid_tokens = max(non_pad_token, 1)
             scale_global = n_routed_experts / (valid_tokens * num_experts_per_tok)
@@ -298,20 +294,18 @@ def accumulate(
         denom_local = max(num_tokens_local, 1)
         base = torch.logsumexp(router_logits, dim=-1).square().sum() / denom_local
 
-        local_loss = base
         loss = base
         if self.loss_cfg.z_loss_global_average and num_tokens_global is not None:
-            # Local component: this rank's share of the global-average z-loss, without the
-            # `× world_size` factor. The backward path keeps `× world_size` (removed in the
-            # reduce-sum switch); summing `local_loss` over ranks reproduces the global z-loss.
+            # Under reduce-sum gradients the injected z-loss stays as this rank's local component
+            # (its share of the global-average z-loss, WITHOUT any `× world_size`). Cross-rank
+            # aggregation happens on the gradients via the FSDP SUM reduce-scatter; summing this
+            # local component over ranks reproduces the global z-loss.
             denom_global = torch.clamp(num_tokens_global, min=1)
-            local_loss = base * num_tokens_local / denom_global
-            loss = local_loss * world_size
+            loss = base * num_tokens_local / denom_global
 
-        local_loss = local_loss * self.loss_cfg.z_loss_alpha / self._batch_size
         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())
         return loss
 
     def finalize(self) -> tuple[torch.Tensor, torch.Tensor]:

xtuner/v1/loss/moe_loss.py

@@ -676,6 +676,10 @@ def fully_shard(
             reshard_after_forward=fsdp_config.reshard_after_forward,
             offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None,
         )
+        # Reduce-scatter gradients with pure SUM (no divide). Combined with the loss forwards no
+        # longer injecting x world_size, each param's gradient is the sum of per-rank local-component
+        # gradients, i.e. the global loss gradient. Covers nested/child FSDP modules via self.modules().
+        self.set_gradient_reduce_sum()
         return self
 
     def _fully_shard(
@@ -1388,6 +1392,21 @@ def post_micro_batch_forward(self, batch_outputs: Sequence[ModelOutputs]) -> Bat
         if "reduced_base_loss" in reduced_other_losses:
             reduced_other_losses["reduced_llm_loss"] = reduced_other_losses.pop("reduced_base_loss")
 
+        # 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:
+            for name in output.model_fields:
+                field_value = getattr(output, name, None)
+                if "loss" in name and isinstance(field_value, torch.Tensor):
+                    expected = "reduced_llm_loss" if name == "loss" else f"reduced_{name}"
+                    if expected not in reduced_other_losses:
+                        raise RuntimeError(
+                            f"Loss field '{name}' has no display curve ('{expected}' missing from reduced "
+                            f"logs): register its detached local component on extra_info (see "
+                            f"`_store_local_display_losses` / CE's `local_base_loss`)."
+                        )
+
         ret = BatchForwardInfo(
             logs_info=reduced_other_losses,
             extra_info=train_engine_extra_info,

xtuner/v1/model/base.py

@@ -138,6 +138,12 @@ def fully_shard(
         self.language_model.set_modules_to_forward_prefetch([self.language_model.layers["0"]])  # type: ignore
 
         self._to_empty_meta()
+        # Reduce-scatter gradients with pure SUM for every sharded submodule. The vision tower,
+        # projector, and this compose root are sharded by their own fully_shard overrides / the root
+        # wrap above, none of which set reduce-sum; a single root-level pass over self.modules()
+        # covers them all (and is idempotent for the language model, already set). Without this the
+        # vision/projector grads silently fall back to FSDP AVG and lose a 1/fsdp_size factor.
+        self.set_gradient_reduce_sum()
         return self
 
     def from_hf(self, hf_path: str | Path, strict=True):

xtuner/v1/model/compose/base.py

@@ -98,6 +98,10 @@ def fully_shard(
         self.language_model.set_modules_to_forward_prefetch([self.language_model.layers["0"]])  # type: ignore
 
         self._to_empty_meta()
+        # Reduce-scatter gradients with pure SUM for every sharded submodule (vision tower, projector,
+        # compose root); their own fully_shard overrides do not set reduce-sum. One root-level pass
+        # over self.modules() covers them all (idempotent for the already-set language model).
+        self.set_gradient_reduce_sum()
         return self
 
     def extract_feature(self, pixel_values):

xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py

@@ -300,6 +300,10 @@ def fully_shard(
         # Make sure it works properly when using fsdp
         if self.config.tie_word_embeddings:
             self.lm_head.weight = self.embed_tokens.weight
+        # Reduce-scatter gradients with pure SUM (no divide) for every sharded submodule; combined
+        # with the loss forwards no longer injecting x world_size, this yields the global loss
+        # gradient. See BaseModel.set_gradient_reduce_sum.
+        self.set_gradient_reduce_sum()
         return self
 
     # TODO: 支持 tp

xtuner/v1/model/dense/dense.py

@@ -1193,6 +1193,10 @@ def fully_shard(
                 module.forward = types.MethodType(self.patched_emb_forward, module)  # type: ignore
 
         self._to_empty_meta()
+        # Reduce-scatter gradients with pure SUM (no divide) for every sharded submodule; the
+        # expert / replicated grads not covered by reduce-scatter are handled without division in
+        # scale_and_reduce_grad. See BaseModel.set_gradient_reduce_sum.
+        self.set_gradient_reduce_sum()
         return self
 
     @property
@@ -1222,10 +1226,10 @@ def scale_and_reduce_grad(self):
             if param.grad is None:
                 continue
 
-            # Expert parameters live on a unique EP rank, so no cross-rank reduction
-            # is needed — just rescale by `ep_size` to keep the effective average.
+            # Expert parameters live on a unique EP rank; their FSDP sharding is only over the
+            # experts_fsdp sub-dim, already SUM-reduced by reduce-scatter. No cross-rank reduction
+            # and no rescaling: under reduce-sum the local-component gradient is what we keep.
             if ep_enabled and ".experts" in name:
-                param.grad.div_(self.ep_mesh.size())  # type: ignore
                 continue
 
             if not isinstance(param, DTensor):
@@ -1252,8 +1256,8 @@ def scale_and_reduce_grad(self):
                 flat_mesh = param.device_mesh[replicate_dim_names[0]]
 
             grad = param.grad.to_local() if isinstance(param.grad, DTensor) else param.grad
-            # Pre-scale locally so the SUM all_reduce below yields the mean across replicas.
-            grad.div_(flat_mesh.size())  # type: ignore
+            # Replicated params get no reduce-scatter; SUM their per-rank local-component grads
+            # across the replicate group with NO pre-divide, matching the reduce-sum invariant.
             grads_by_group.setdefault(flat_mesh.get_group(), []).append(grad)  # type: ignore
 
         # One coalesced all_reduce per process group covers all replicated grads.