loss debug

Author: meichangsu1

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

@@ -1,4 +1,5 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import os
 from functools import partial
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -921,6 +922,43 @@ def __init__(
         self._tokenizer_id = tokenizer_id
         self._tokenizer = None
         self._initialized = False
+        debug_flag = os.getenv("TWINKLE_DEBUG_SP_LOSS", "").strip().lower()
+        self._debug_sp_loss = debug_flag not in ("", "0", "false", "off", "no")
+        try:
+            self._debug_sp_loss_max_steps = max(1, int(os.getenv("TWINKLE_DEBUG_SP_LOSS_STEPS", "8")))
+        except ValueError:
+            self._debug_sp_loss_max_steps = 8
+        self._debug_sp_loss_seen = 0
+
+    def _maybe_debug_reduce_loss(
+        self,
+        reduction: str,
+        loss_in: torch.Tensor,
+        loss_out: torch.Tensor,
+        num_valid_tokens: Optional[torch.Tensor],
+        compensate_factor: float,
+    ) -> None:
+        if not self._debug_sp_loss or self._debug_sp_loss_seen >= self._debug_sp_loss_max_steps:
+            return
+        rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+        sp_rank = (
+            dist.get_rank(sequence_parallel._sp_group)
+            if sequence_parallel._sp_group is not None and dist.is_initialized()
+            else 0
+        )
+        token_str = "None"
+        if num_valid_tokens is not None:
+            token_str = str(int(num_valid_tokens.detach().item()))
+        in_val = float(loss_in.detach().item())
+        out_val = float(loss_out.detach().item())
+        print(
+            "[SP-LOSS-DEBUG] "
+            f"rank={rank} sp_rank={sp_rank} reduction={reduction} "
+            f"loss_in={in_val:.6f} loss_out={out_val:.6f} "
+            f"local_valid_tokens={token_str} compensate_factor={compensate_factor:.4f}",
+            flush=True,
+        )
+        self._debug_sp_loss_seen += 1
 
     def _get_tokenizer(self) -> Optional[PreTrainedTokenizer]:
         if self._tokenizer is not None:
@@ -1003,13 +1041,15 @@ 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)
+        loss_in = loss.detach()
 
         class _ReduceSequenceParallelLoss(torch.autograd.Function):
             @staticmethod
-            def forward(ctx, local_sum: torch.Tensor, num_valid_tokens: torch.Tensor) -> torch.Tensor:
-                if num_valid_tokens.item() == 0:
-                    local_sum = torch.nan_to_num(local_sum)
+            def forward(ctx, local_mean: torch.Tensor, num_valid_tokens: torch.Tensor) -> torch.Tensor:
                 local_tokens = num_valid_tokens.detach().clone()
+                local_sum = local_mean * local_tokens
+                if local_tokens.item() == 0:
+                    local_sum = torch.nan_to_num(local_sum)
                 global_sum = local_sum.detach().clone()
                 dist.all_reduce(global_sum, group=sequence_parallel._sp_group)
                 global_tokens = num_valid_tokens.detach().clone()
@@ -1023,9 +1063,10 @@ def forward(ctx, local_sum: torch.Tensor, num_valid_tokens: torch.Tensor) -> tor
             def backward(ctx, grad_output: torch.Tensor):
                 local_tokens, global_tokens = ctx.saved_tensors
                 if global_tokens.item() == 0:
-                    return grad_output, None
-                grad_local_sum = grad_output * (local_tokens / global_tokens) * compensate_factor
-                return grad_local_sum, None
+                    return torch.zeros_like(grad_output), None
+                # d(global_mean)/d(local_mean) = local_tokens / global_tokens.
+                grad_local_mean = grad_output * (local_tokens / global_tokens) * compensate_factor
+                return grad_local_mean, None
 
         class _ReduceSequenceParallelSum(torch.autograd.Function):
             @staticmethod
@@ -1039,12 +1080,30 @@ def backward(ctx, grad_output: torch.Tensor):
                 return grad_output
 
         if reduction == "sum":
-            return _ReduceSequenceParallelSum.apply(loss)
+            out = _ReduceSequenceParallelSum.apply(loss)
+            num_valid_tokens = None
+            if self._debug_sp_loss:
+                num_valid_tokens = (labels != ignore_index).sum().to(loss.device).detach()
+            self._maybe_debug_reduce_loss(
+                reduction,
+                loss_in=loss_in,
+                loss_out=out.detach(),
+                num_valid_tokens=num_valid_tokens,
+                compensate_factor=compensate_factor,
+            )
+            return out
 
-        # Default to mean reduction: assume `loss` is local mean, convert to local sum.
+        # Default to mean reduction: `loss` is local mean.
         num_valid_tokens = (labels != ignore_index).sum().to(loss.device)
-        local_sum = loss * num_valid_tokens
-        return _ReduceSequenceParallelLoss.apply(local_sum, num_valid_tokens)
+        out = _ReduceSequenceParallelLoss.apply(loss, num_valid_tokens)
+        self._maybe_debug_reduce_loss(
+            reduction,
+            loss_in=loss_in,
+            loss_out=out.detach(),
+            num_valid_tokens=num_valid_tokens.detach(),
+            compensate_factor=compensate_factor,
+        )
+        return out
 
     def wrap_model(self, model, optimizer=None):
         self.initialize()

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='sum'),
+            LossMetric(self._device_mesh, self._dp_group, loss_reduction='mean'),
             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='sum'),
+            LossMetric(self._device_mesh, self._dp_group, loss_reduction='mean'),
             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='sum'),
+            loss_instance=CrossEntropyLoss(reduction='mean'),
             template=Template(self.tokenizer_id),
             processor=InputProcessor(self.device_mesh),
             _device_mesh=self.device_mesh,