[Refactor] Drop non-global aux-loss averaging mode

xtuner/v1/engine/train_engine.py

@@ -216,7 +216,10 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo:
         micro_batch_results = []
 
         data_batch_info = self.model.pre_micro_batch_forward(data_batches)
-        total_loss = torch.tensor(0.0, device=DEVICE)
+        # Display total loss: sum of every loss term's per-rank display value (from each loss ctx's
+        # calibrate(), stored on output.calibrated_losses) across micro-batches -- NO cross-rank
+        # all_reduce, so each rank reports its own loss (equal to the global loss at world size 1).
+        display_total_loss = torch.tensor(0.0, device=DEVICE)
 
         for i in range(0, len(data_batches), intra_layer_micro_batch):
             ProberList.set_micro_batch_iter(micro_batch_iter)
@@ -240,12 +243,13 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo:
 
             loss = self._get_total_loss(output)
             loss.backward()
-            total_loss += loss.detach()
+            for value in (output.calibrated_losses or {}).values():
+                display_total_loss += value.detach().float()
             # call dump_forward_records after backward to record the recomputed activations
             ProberList.after_micro_iter_forward()
 
         batch_forward_info = self.model.post_micro_batch_forward(micro_batch_results)
-        return TrainStepInfo(total_loss=total_loss.item(), **data_batch_info, **batch_forward_info)
+        return TrainStepInfo(total_loss=display_total_loss.item(), **data_batch_info, **batch_forward_info)
 
     def from_hf(self, hf_path: str | Path, strict: bool = False):
         self.model.from_hf(hf_path=hf_path, strict=strict)

xtuner/v1/loss/aux_loss.py

@@ -3,10 +3,22 @@
 from pydantic import BaseModel, ConfigDict
 from torch import distributed as dist
 from torch.distributed._functional_collectives import all_reduce
+from typing_extensions import TypedDict
 
 from xtuner.v1.loss.moe_loss import BalancingLossContext, ZLossContext
 
 
+class AuxLossFinalizeOutput(TypedDict):
+    """Backward-relevant outputs of :meth:`AuxLossContext.finalize`.
+
+    Per-rank display values are not returned here; each loss context exposes them via ``calibrate()``.
+    """
+
+    balancing_loss: torch.Tensor | None
+    z_loss: torch.Tensor | None
+    tokens_per_expert_global: torch.Tensor
+
+
 class AuxLossScaler(torch.autograd.Function):
     """Inject an auxiliary loss into the main forward graph as a passthrough.
 
@@ -91,7 +103,6 @@ def accumulate(
         z_ctx: list[ZLossContext] | ZLossContext | None = None,
         num_tokens_local: int = 0,
         num_tokens_global: torch.Tensor | None = None,
-        world_size: int = 1,
     ) -> torch.Tensor:
         """Accumulate routing statistics for one layer and inject z-loss into
         the main graph.
@@ -111,9 +122,8 @@ def accumulate(
             num_tokens_local (int): Non-padding token count on this rank for the current forward.
                 Required when any z-loss context is provided.
             num_tokens_global (torch.Tensor | None): All-reduced non-padding token count across
-                ranks (int64 scalar). Pass ``None`` when ``z_loss_global_average`` is off or no
-                process group is initialized.
-            world_size (int): World size that produced ``num_tokens_global``.
+                ranks (int64 scalar). Pass ``None`` when no process group is initialized
+                (single-process reference / eval).
 
         Returns:
             torch.Tensor: ``hidden_states`` augmented with the per-layer z-loss autograd hook.
@@ -139,7 +149,6 @@ def accumulate(
                 router_logits=selected_router_logits,
                 num_tokens_local=num_tokens_local,
                 num_tokens_global=num_tokens_global,
-                world_size=world_size,
             )
             hidden_states = AuxLossScaler.apply(hidden_states, z_loss_l)
 
@@ -151,15 +160,20 @@ def finalize(
         balancing_ctx: list[BalancingLossContext] | BalancingLossContext | None,
         z_ctx: list[ZLossContext] | ZLossContext | None,
         non_pad_token: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
-        """Finalize split auxiliary losses and expert counts from runtime
-        state."""
+    ) -> "AuxLossFinalizeOutput":
+        """Finalize auxiliary losses and expert counts from runtime state.
+
+        Returns:
+            AuxLossFinalizeOutput: ``balancing_loss`` / ``z_loss`` carry the backward graph and the
+            globally reduced ``tokens_per_expert_global``. Per-rank display values are produced by
+            each context's ``calibrate()``, not returned here.
+        """
         tokens_per_expert_local, tokens_per_expert_global = self._cal_tokens_per_expert()
 
         balancing_loss: torch.Tensor | None = None
         balancing_list = _as_list(balancing_ctx)
         if balancing_list:
-            partials = [
+            losses = [
                 ctx.finalize(
                     tokens_per_expert_local=tokens_per_expert_local,
                     tokens_per_expert_global=tokens_per_expert_global,
@@ -169,15 +183,19 @@ def finalize(
                 )
                 for ctx in balancing_list
             ]
-            balancing_loss = partials[0] if len(partials) == 1 else torch.stack(partials).sum(dim=0)
+            balancing_loss = losses[0] if len(losses) == 1 else torch.stack(losses).sum(dim=0)
 
         z_loss: torch.Tensor | None = None
         z_list = _as_list(z_ctx)
         if z_list:
-            partials = [ctx.finalize() for ctx in z_list]
-            z_loss = partials[0] if len(partials) == 1 else torch.stack(partials).sum(dim=0)
+            z_losses = [ctx.finalize() for ctx in z_list]
+            z_loss = z_losses[0] if len(z_losses) == 1 else torch.stack(z_losses).sum(dim=0)
 
-        return balancing_loss, z_loss, tokens_per_expert_global
+        return AuxLossFinalizeOutput(
+            balancing_loss=balancing_loss,
+            z_loss=z_loss,
+            tokens_per_expert_global=tokens_per_expert_global,
+        )
 
     def _cal_tokens_per_expert(self) -> tuple[torch.Tensor, torch.Tensor]:
         """Stack per-layer expert counts and produce both local and globally

xtuner/v1/loss/base_loss_ctx.py

@@ -22,20 +22,21 @@
 # 1. Compute the global loss mask sum among dp, sp and grad accumulation:
 #    global_loss_mask_sum = all_reduce(sum([loss_mask.sum() for loss_mask in loss_masks_grad_acc]), op=dist.ReduceOp.SUM, group=world)
 #                           = (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13)
-# 2. Compute the iter loss, take rank0 iter0 as an example:
+# 2. Compute the iter loss (this rank's local component), take rank0 iter0 as an example:
 #    a. loss_{rank0iter0} = (l00 * w00 * m00 + l01 * w01 * m01)
 #    b. loss_{rank0iter0} = loss_{rank0iter0} / global_loss_mask_sum
 #                         = (l00 * w00 * m00 + l01 * w01 * m01) /
 #                           (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13)
-#    c. loss_{rank0iter0} = all_reduce_autograd(loss_{rank0iter0}, op=dist.ReduceOp.SUM, group=world)
-#                         = (l00 * w00 * m00 + l01 * w01 * m01 + l02 * w02 * m02 + l03 * w03 * m03) /
-#                           (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13)
-# 3. Compute the step loss:
+#       Under reduce-sum gradients this stays a local component: it is NOT all-reduced across ranks.
+#       Each rank keeps its own numerator over the shared (detached) global denominator; the
+#       cross-rank aggregation happens on the gradients (FSDP / scale_and_reduce_grad SUM), whose
+#       sum over ranks equals the global-batch loss gradient.
+# 3. Compute the step loss (this rank's local component summed over grad-acc iters):
 #    step_loss = loss_{rank0iter0} + loss_{rank0iter1}
-#              = (l00 * w00 * m00 + l01 * w01 * m01 + l02 * w02 * m02 + l03 * w03 * m03 +
-#                 l10 * w10 * m10 + l11 * w11 * m11 + l12 * w12 * m12 + l13 * w13 * m13) /
+#              = (l00 * w00 * m00 + l01 * w01 * m01 + l10 * w10 * m10 + l11 * w11 * m11) /
 #                (m00 + m01 + m02 + m03 + m10 + m11 + m12 + m13)
-#    It's equivalent to loss calculation in sp1, dp1 and grad acc 1.
+#    Its SUM over all ranks equals the sp1/dp1/grad-acc-1 loss; that global value is restored for
+#    display by a separate detached SUM all_reduce (see the reduce-sum design's logging section).
 
 
 class BaseLossKwargs(BaseModel):

xtuner/v1/loss/ce_loss.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
 
@@ -106,6 +105,11 @@ class LMHeadLossContext(BaseLossContext):
     loss_cfg: CELossConfig
     loss_kwargs: CELossKwargs
 
+    # Display state for calibrate(): coefficient stashed by build_batches (global_denom /
+    # step_grad_tokens) and this rank's detached backward loss stashed by forward().
+    _display_coeff: torch.Tensor
+    _backward_loss: torch.Tensor
+
     def __init__(self, loss_cfg: CELossConfig, loss_kwargs: CELossKwargs):
         super().__init__(loss_cfg, loss_kwargs)
 
@@ -178,10 +182,21 @@ def build_batches(  # type: ignore[override]
         if dist.is_initialized():
             dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM)
 
+        # Coefficient to turn each rank's local backward loss (local token sum / global denominator)
+        # back into this rank's per-token mean for display: multiply by global_denominator, divide by
+        # this rank's own grad-token count across the whole step. Summing `calibrate()` over the
+        # micro-batches then yields this rank's per-token-mean loss with no cross-rank all_reduce, and
+        # equals the global loss when world size is 1 (global_denominator == step_grad_tokens).
+        step_grad_tokens = torch.zeros((), device=global_denominator.device)
+        for loss_ctx in loss_ctx_list:
+            step_grad_tokens += (loss_ctx.loss_kwargs.shifted_labels != loss_cfg.ignore_idx).sum()
+        display_coeff = (global_denominator / step_grad_tokens.clamp_min(1.0)).detach()
+
         for loss_ctx in loss_ctx_list:
             loss_ctx._batch_size = len(loss_ctx_list)
             assert loss_ctx.loss_kwargs.loss_weight is not None
             loss_ctx.loss_kwargs.loss_weight /= global_denominator + 1e-12
+            loss_ctx._display_coeff = display_coeff
         return loss_ctx_list
 
     def loss_fn(
@@ -280,14 +295,18 @@ def forward(
         if not isinstance(extra_info, ModelForwardExtraLogInfo):
             extra_info = ModelForwardExtraLogInfo(extra_info)
 
-        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 per-rank
+        # display value is produced by `calibrate()`; stash the detached loss for it.
+        self._backward_loss = loss.detach()
         return loss, (logits, extra_info)
 
+    def calibrate(self) -> torch.Tensor:
+        """This rank's per-token-mean CE for display (detached, no cross-rank
+        all_reduce)."""
+        return self._backward_loss * self._display_coeff
+
 
 # Deprecated: Use LMHeadLossContext instead. Will be removed in version 1.1.0
 CELossContext = LMHeadLossContext