remove debug log

Author: meichangsu1

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

@@ -1,5 +1,4 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-import os
 from functools import partial
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -801,7 +800,8 @@ def pad_and_split_inputs(self,
                 # - In next-token-aligned labels, this appears at labels[b-1]
                 boundary_starts = (real_position_ids == 0)
                 prev = torch.zeros_like(boundary_starts, dtype=torch.bool)
-                prev[..., 1:] = boundary_starts[..., :-1]
+                # Mask token b-1 when boundary starts at b.
+                prev[..., :-1] = boundary_starts[..., 1:]
                 labels = labels.clone()
                 labels[prev] = -100
                 # Also avoid any potential wrap-around supervision at the end of the concatenated stream.
@@ -922,43 +922,6 @@ 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:
@@ -1042,7 +1005,6 @@ 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):
             @staticmethod
@@ -1081,34 +1043,16 @@ def forward(ctx, local_sum: torch.Tensor) -> torch.Tensor:
 
             @staticmethod
             def backward(ctx, grad_output: torch.Tensor):
-                # Preserve original gradient scale (before forward-side metric scaling).
-                return grad_output * ctx.sum_metric_scale
+                # Keep training gradient scale unchanged; forward-side scaling is for
+                # logging/metric alignment under outer collect='mean'.
+                return grad_output
 
         if reduction == "sum":
-            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
+            return _ReduceSequenceParallelSum.apply(loss)
 
         # Default to mean reduction: `loss` is local mean.
         num_valid_tokens = (labels != ignore_index).sum().to(loss.device)
-        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
+        return _ReduceSequenceParallelLoss.apply(loss, num_valid_tokens)
 
     def wrap_model(self, model, optimizer=None):
         self.initialize()