fix lint

Author: meichangsu1

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

@@ -870,7 +870,7 @@ class SequenceParallelConfig:
     enabled: bool = True
     ulysses_size: Optional[int] = None
     gather_logits: bool = True
-    loss_reduction: str = "mean"
+    loss_reduction: str = 'mean'
     compensate_fsdp_avg: bool = False
 
 
@@ -975,17 +975,16 @@ def reduce_loss(self, loss: torch.Tensor, labels: Optional[torch.Tensor], ignore
         if labels is None or sequence_parallel._sp_group is None:
             return loss
         # Compute global loss via autograd-aware all-reduce.
-        reduction = str(self.sp_config.get("loss_reduction", "mean")).lower()
-        if reduction == "none":
-            raise ValueError(
-                "SequenceParallelStrategy.reduce_loss only supports reduction='sum' or 'mean'. "
-                "Please aggregate per-token losses before calling reduce_loss."
-            )
-        compensate_fsdp_avg = bool(self.sp_config.get("compensate_fsdp_avg", False))
+        reduction = str(self.sp_config.get('loss_reduction', 'mean')).lower()
+        if reduction == 'none':
+            raise ValueError("SequenceParallelStrategy.reduce_loss only supports reduction='sum' or 'mean'. "
+                             'Please aggregate per-token losses before calling reduce_loss.')
+        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)
 
         class _ReduceSequenceParallelLoss(torch.autograd.Function):
+
             @staticmethod
             def forward(ctx, local_mean: torch.Tensor, num_valid_tokens: torch.Tensor) -> torch.Tensor:
                 local_tokens = num_valid_tokens.detach().clone()
@@ -1011,6 +1010,7 @@ def backward(ctx, grad_output: torch.Tensor):
                 return grad_local_mean, None
 
         class _ReduceSequenceParallelSum(torch.autograd.Function):
+
             @staticmethod
             def forward(ctx, local_sum: torch.Tensor) -> torch.Tensor:
                 ctx.sum_metric_scale = sum_metric_scale
@@ -1026,7 +1026,7 @@ def backward(ctx, grad_output: torch.Tensor):
                 # logging/metric alignment under outer collect='mean'.
                 return grad_output
 
-        if reduction == "sum":
+        if reduction == 'sum':
             return _ReduceSequenceParallelSum.apply(loss)
 
         # Default to mean reduction: `loss` is local mean.

src/twinkle/model/transformers/transformers.py

@@ -252,7 +252,7 @@ def _ensure_sp_strategy(self) -> None:
         # accelerate DDP/FSDP paths), compensate SP loss backward to keep gradient scale.
         if isinstance(self.strategy, (NativeFSDPStrategy, AccelerateStrategy)) and self.device_mesh is not None:
             if (self.device_mesh.ulysses_size or 1) > 1 and (self.device_mesh.data_world_size or 1) > 1:
-                sp_config["compensate_fsdp_avg"] = True
+                sp_config['compensate_fsdp_avg'] = True
         self.sp_strategy = SequenceParallelStrategy(
             self.device_mesh,
             sp_config,
@@ -440,9 +440,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:
-            reduction = getattr(loss_instance, "reduction", None)
+            reduction = getattr(loss_instance, 'reduction', None)
             if reduction is not None:
-                self.sp_strategy.sp_config["loss_reduction"] = str(reduction)
+                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