loss metric fix

Author: meichangsu1

src/twinkle/model/transformers/strategy/sequence_parallel.py

@@ -1041,6 +1041,7 @@ def reduce_loss(self, loss: torch.Tensor, labels: Optional[torch.Tensor], ignore
             )
         compensate_fsdp_avg = bool(self.sp_config.get("compensate_fsdp_avg", False))
         compensate_factor = float(self.ulysses_size if compensate_fsdp_avg else 1.0)
+        sum_metric_scale = float(self.ulysses_size)
         loss_in = loss.detach()
 
         class _ReduceSequenceParallelLoss(torch.autograd.Function):
@@ -1071,13 +1072,17 @@ def backward(ctx, grad_output: torch.Tensor):
         class _ReduceSequenceParallelSum(torch.autograd.Function):
             @staticmethod
             def forward(ctx, local_sum: torch.Tensor) -> torch.Tensor:
+                ctx.sum_metric_scale = sum_metric_scale
                 global_sum = local_sum.detach().clone()
                 dist.all_reduce(global_sum, group=sequence_parallel._sp_group)
-                return global_sum
+                # Keep logging/metric value aligned with non-SP sum semantics under
+                # outer collect='mean' by removing one SP replication factor.
+                return global_sum / ctx.sum_metric_scale
 
             @staticmethod
             def backward(ctx, grad_output: torch.Tensor):
-                return grad_output
+                # Preserve original gradient scale (before forward-side metric scaling).
+                return grad_output * ctx.sum_metric_scale
 
         if reduction == "sum":
             out = _ReduceSequenceParallelSum.apply(loss)

src/twinkle/model/transformers/transformers.py

@@ -80,13 +80,13 @@ def __post_init__(self):
 
     def _build_metrics(self):
         self.train_metrics = [
-            LossMetric(self._device_mesh, self._dp_group, loss_reduction='mean'),
+            LossMetric(self._device_mesh, self._dp_group, loss_reduction='sum'),
             Accuracy(self._device_mesh, self._dp_group),
             TrainMetric(self._device_mesh, self._dp_group),
         ]
 
         self.eval_metrics = [
-            LossMetric(self._device_mesh, self._dp_group, loss_reduction='mean'),
+            LossMetric(self._device_mesh, self._dp_group, loss_reduction='sum'),
             Accuracy(self._device_mesh, self._dp_group),
             TrainMetric(self._device_mesh, self._dp_group),
         ]
@@ -317,7 +317,7 @@ def _ensure_optimizer_dp_groups(self):
 
     def _construct_default_optimizer_group(self):
         return OptimizerGroup(
-            loss_instance=CrossEntropyLoss(reduction='mean'),
+            loss_instance=CrossEntropyLoss(reduction='sum'),
             template=Template(self.tokenizer_id),
             processor=InputProcessor(self.device_mesh),
             _device_mesh=self.device_mesh,
@@ -431,10 +431,9 @@ def calculate_loss(self, **kwargs):
         optimizer_config = self.optimizer_group[adapter_name]
         optimizer_config.num_tokens += counts.item()
         if self.sp_strategy is not None and 'labels' in inputs:
-            if "loss_reduction" not in self.sp_strategy.sp_config:
-                reduction = getattr(loss_instance, "reduction", None)
-                if reduction is not None:
-                    self.sp_strategy.sp_config["loss_reduction"] = str(reduction)
+            reduction = getattr(loss_instance, "reduction", None)
+            if reduction is not None:
+                self.sp_strategy.sp_config["loss_reduction"] = str(reduction)
             loss_value = self.sp_strategy.reduce_loss(loss_value, inputs['labels'])
         optimizer_config.loss_value += loss_value
         outputs['loss'] = optimizer_config.loss_value