[FSDP][Loss] Reduce-sum gradient reduction with per-batch aux-loss hub

xtuner/v1/engine/train_engine.py

@@ -35,7 +35,6 @@
     BatchForwardInfo,
     DataBatchInfo,
     ModelItem,
-    ModelOutputs,
     XTunerBaseModelConfig,
 )
 from xtuner.v1.patch.xtuner_storage import XtunerCacheWriter, _get_async_dcp_save_timeout
@@ -51,7 +50,7 @@
 
 
 class TrainStepInfo(DataBatchInfo, BatchForwardInfo):
-    total_loss: float
+    pass
 
 
 logger = get_logger()
@@ -216,7 +215,6 @@ 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)
 
         for i in range(0, len(data_batches), intra_layer_micro_batch):
             ProberList.set_micro_batch_iter(micro_batch_iter)
@@ -238,14 +236,15 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo:
 
             micro_batch_results.append(output)
 
-            loss = self._get_total_loss(output)
-            loss.backward()
-            total_loss += loss.detach()
+            # The model folds every loss term (LM/policy + aux + MTP) into output.loss; the engine
+            # just backprops it. Per-term display values ride on output.loss_log.
+            assert output.loss is not None, "model forward must return a loss in training mode"
+            output.loss.backward()
             # 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(**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)
@@ -558,17 +557,3 @@ def _maybe_precompute_float8_dynamic_scale_for_fsdp(self):
         for model in self.model.modules():
             if isinstance(model, BaseModel) and model.float8_handler is not None:
                 model.float8_handler.precompute_float8_dynamic_scale_for_fsdp(model)
-
-    def _get_total_loss(self, model_outputs: ModelOutputs) -> torch.Tensor:
-        # TODO: This logic should be moved into the model layer. The model should be responsible
-        # for aggregating all losses (CE loss, balancing loss, z loss, etc.) and returning a
-        # single total_loss. The engine should only call model.forward() and use the returned
-        # total_loss directly, rather than iterating through fields to sum losses here.
-        # This would provide better separation of concerns and make the loss computation logic
-        # more explicit and maintainable.
-        loss = torch.tensor(0.0, device=DEVICE)
-        for key in model_outputs.model_fields:
-            value = getattr(model_outputs, key)
-            if "loss" in key and isinstance(value, torch.Tensor):
-                loss += value
-        return loss

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 /
+    # rank_denom) 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)
 
@@ -171,17 +175,28 @@ def build_batches(  # type: ignore[override]
             loss_ctx.loss_kwargs.loss_weight = loss_weight
             loss_weight_list.append(loss_weight)
 
-        # Compute the denominator used in the global calibration of the loss
+        # Denominators for calibration. rank_denominator is this rank's local weight sum; clone before
+        # the in-place all_reduce so global_denominator does not overwrite it (same tensor otherwise).
         rank_denominator = sum(loss_weight.sum() for loss_weight in loss_weight_list)
         rank_denominator = cast(torch.Tensor, rank_denominator)
-        global_denominator = rank_denominator
+        global_denominator = rank_denominator.clone()
         if dist.is_initialized():
             dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM)
 
+        # Coefficient to turn each rank's backward loss (local weighted sum / global_denominator) back
+        # into this rank's own reduction-correct mean for display: multiply by global_denominator,
+        # divide by this rank's own weight sum. The weight sum is the right per-rank denominator for
+        # every reduction -- token (weight==1, so it equals the grad-token count), sample, and square
+        # alike -- whereas a raw grad-token count is only correct for token reduction. Summing
+        # `calibrate()` over the micro-batches yields this rank's mean with no cross-rank all_reduce,
+        # equal to the global loss when world size is 1 (global_denominator == rank_denominator).
+        display_coeff = (global_denominator / rank_denominator.clamp_min(1e-12)).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