lint

Author: tastelikefeet

src/twinkle/dataset/base.py

@@ -86,9 +86,9 @@ def encode(self, add_generation_prompt: bool = False, **kwargs):
         encode_fn = partial(self.template.batch_encode, add_generation_prompt=add_generation_prompt)
         with processing_lock('dataset'):
             # use a default lock because encode is to all datasets
-            self.dataset = self.dataset.map(encode_fn,
-                                            **kwargs).filter(lambda batch: [True] * len(next(iter(batch.values()))) if 'input_ids' not in batch else [len(x) > 0 for x in batch['input_ids']],
-                                                             **kwargs)
+            self.dataset = self.dataset.map(encode_fn, **kwargs).filter(
+                lambda batch: [True] * len(next(iter(batch.values())))
+                if 'input_ids' not in batch else [len(x) > 0 for x in batch['input_ids']], **kwargs)
 
     @remote_function()
     def check(self, **kwargs):

src/twinkle/infra/collectors.py

@@ -1,12 +1,9 @@
 import numpy as np
-from typing import TYPE_CHECKING, Any, Dict, List
+from typing import Any, Dict, List
 
 from twinkle import DeviceMesh
 from twinkle.utils import pad_and_stack_tensors
 
-if TYPE_CHECKING:
-    import torch
-
 
 def collect_tensor_dict(outputs: List[Dict[str, Any]], device_mesh: DeviceMesh) -> Dict[str, Any]:
     import torch

src/twinkle/loss/dpo.py

@@ -9,8 +9,8 @@
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from twinkle.data_format import LossOutput
-from twinkle.utils.torch_utils import selective_log_softmax
 from twinkle.loss.base import Loss
+from twinkle.utils.torch_utils import selective_log_softmax
 
 if TYPE_CHECKING:
     import torch
@@ -176,14 +176,10 @@ def _align_logps(
                 # Truncate right (keep left part) - may happen in Ray result merging
                 return logps[:, :target_seq_len]
             else:
-                raise ValueError(
-                    f'ref_logps seq_len ({src_seq_len}) < target seq_len ({target_seq_len}). '
-                    f'This should not happen when both models process the same batch.'
-                )
+                raise ValueError(f'ref_logps seq_len ({src_seq_len}) < target seq_len ({target_seq_len}). '
+                                 f'This should not happen when both models process the same batch.')
 
-        raise ValueError(
-            f'Cannot align ref_logps shape {logps.shape} to target shape {target_shape}'
-        )
+        raise ValueError(f'Cannot align ref_logps shape {logps.shape} to target shape {target_shape}')
 
     def _compute_dpo_loss(
         self,
@@ -227,7 +223,7 @@ def _compute_dpo_loss(
         elif self.loss_type == 'ipo':
             # IPO (Identity Preference Optimization) loss
             # Reference: "A General Theoretical Paradigm to Understand Learning from Human Feedback"
-            losses = (logits - 1 / (2 * self.beta)) ** 2
+            losses = (logits - 1 / (2 * self.beta))**2
         elif self.loss_type == 'kto_pair':
             # KTO pair loss (simplified version)
             chosen_logratios_scaled = self.beta * chosen_logratios
@@ -236,7 +232,7 @@ def _compute_dpo_loss(
             rejected_losses = F.sigmoid(rejected_logratios_scaled)
             losses = chosen_losses + rejected_losses
         else:
-            raise ValueError(f"Unknown loss_type: {self.loss_type}")
+            raise ValueError(f'Unknown loss_type: {self.loss_type}')
 
         # Apply label smoothing if specified
         if self.label_smoothing > 0:
@@ -292,7 +288,7 @@ def __call__(
             labels = labels.unsqueeze(0)
 
         batch_size = labels.shape[0]
-        assert batch_size % 2 == 0, "Batch size must be even (chosen + rejected pairs)"
+        assert batch_size % 2 == 0, 'Batch size must be even (chosen + rejected pairs)'
 
         # Get log probabilities from outputs
         logps = self._get_logps_from_outputs(outputs, labels)
@@ -314,9 +310,7 @@ def __call__(
             reference_rejected_logps = ref_rejected_logps.to(device=device, dtype=dtype)
         elif ref_logps is not None:
             # Per-token reference log probs provided, need to align and sum
-            ref_logps_aligned = self._align_logps(
-                ref_logps, labels.shape, device, dtype
-            )
+            ref_logps_aligned = self._align_logps(ref_logps, labels.shape, device, dtype)
             ref_chosen, ref_rejected = self._split_chosen_rejected(ref_logps_aligned)
             reference_chosen_logps = self._compute_sequence_logps(ref_chosen, chosen_labels)
             reference_rejected_logps = self._compute_sequence_logps(ref_rejected, rejected_labels)
@@ -392,7 +386,7 @@ def __call__(
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)
 
-        assert labels.shape[0] % 2 == 0, "Batch size must be even (chosen + rejected pairs)"
+        assert labels.shape[0] % 2 == 0, 'Batch size must be even (chosen + rejected pairs)'
 
         # Get log probabilities
         logps = self._get_logps_from_outputs(outputs, labels)
@@ -455,7 +449,7 @@ def __call__(
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)
 
-        assert labels.shape[0] % 2 == 0, "Batch size must be even"
+        assert labels.shape[0] % 2 == 0, 'Batch size must be even'
 
         # Get log probabilities
         logps = self._get_logps_from_outputs(outputs, labels)
@@ -521,7 +515,7 @@ def __call__(
         if labels.dim() == 1:
             labels = labels.unsqueeze(0)
 
-        assert labels.shape[0] % 2 == 0, "Batch size must be even"
+        assert labels.shape[0] % 2 == 0, 'Batch size must be even'
 
         # Get log probabilities
         logps = self._get_logps_from_outputs(outputs, labels)
@@ -540,8 +534,8 @@ def __call__(
         # Odds ratio: log(odds_chosen / odds_rejected)
         # log_odds = log(p/(1-p)) = log(p) - log(1-p)
         # Use numerically stable computation
-        prob_chosen = torch.exp(chosen_avg_logps).clamp(min=1e-7, max=1-1e-7)
-        prob_rejected = torch.exp(rejected_avg_logps).clamp(min=1e-7, max=1-1e-7)
+        prob_chosen = torch.exp(chosen_avg_logps).clamp(min=1e-7, max=1 - 1e-7)
+        prob_rejected = torch.exp(rejected_avg_logps).clamp(min=1e-7, max=1 - 1e-7)
         log_odds_chosen = torch.log(prob_chosen) - torch.log(1 - prob_chosen)
         log_odds_rejected = torch.log(prob_rejected) - torch.log(1 - prob_rejected)
 

src/twinkle/metric/dpo.py

@@ -57,10 +57,8 @@ def _align_logps(self, logps, target_shape, device, dtype):
         if src_len == target_len:
             return logps
         elif src_len < target_len:
-            raise ValueError(
-                f'ref_logps seq_len ({src_len}) < target seq_len ({target_len}). '
-                f'This should not happen when both models process the same batch.'
-            )
+            raise ValueError(f'ref_logps seq_len ({src_len}) < target seq_len ({target_len}). '
+                             f'This should not happen when both models process the same batch.')
         else:
             return logps[:, :target_len]
 
@@ -84,7 +82,7 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M
         logps = outputs.get('logps')
         if logps is None or len(logps) == 0:
             return
-        
+
         if isinstance(logps, list) and logps:
             logps = pad_and_stack_tensors(logps)
 
@@ -128,9 +126,7 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M
             ref_logps = ref_outputs.get('logps')
             if ref_logps is not None:
                 # Align ref_logps to match labels shape (handles different seq lengths)
-                ref_logps = self._align_logps(
-                    ref_logps, labels.shape, labels.device, logps.dtype
-                )
+                ref_logps = self._align_logps(ref_logps, labels.shape, labels.device, logps.dtype)
 
                 ref_seq_logps = self._compute_sequence_logps(ref_logps, labels)
                 ref_chosen_logps, ref_rejected_logps = self._split_chosen_rejected(ref_seq_logps)

src/twinkle/model/megatron/megatron.py

@@ -25,12 +25,12 @@
 from twinkle import DeviceMesh, Platform, remote_class, remote_function, requires, torch_util
 from twinkle.checkpoint_engine.mixin import CheckpointEngineMixin
 from twinkle.data_format import InputFeature, ModelOutput, Trajectory
-from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus
 from twinkle.hub import HubOperation
 from twinkle.infra import collect_tensor_dict
 from twinkle.loss import CrossEntropyLoss, Loss
 from twinkle.metric import LossMetric, Metric, TrainMetric
 from twinkle.model.base import TwinkleModel
+from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus
 from twinkle.patch import Patch, apply_patch
 from twinkle.processor import InputProcessor
 from twinkle.template import Template
@@ -435,7 +435,7 @@ def forward_backward(self,
                 seq_length = original_seq_length + (divisor - original_seq_length % divisor)
             else:
                 seq_length = original_seq_length
-        
+
         num_microbatches = len(inputs)
         loss_extra_kwargs_per_mb = []
         if num_microbatches <= 1:
@@ -463,10 +463,12 @@ def post_loss_function(output_tensor, inputs, logps):
             if not counts:
                 # Later will gather this value, so it becomes:
                 # 1. SUM loss: gather_sum(local_num_tokens) = global_num_tokens
-                # 2. PER TOKEN MEAN loss: gather_sum(1 * gradient_accumulation_steps ) = gradient_accumulation_steps * world_size
+                # 2. PER TOKEN MEAN loss: gather_sum(1 * gradient_accumulation_steps )
+                #       = gradient_accumulation_steps * world_size
                 # Then, grad will divided by this value:
                 # 1. SUM loss: (global_sum_grad) / (global_num_tokens) = global_sum_grad/global_num_tokens
-                # 2. PER TOKEN MEAN loss: (gather_sum(per_token_grad * gradient_accumulation_steps)) / (gradient_accumulation_steps  * world_size ) = avg_per_token_grad
+                # 2. PER TOKEN MEAN loss: (gather_sum(per_token_grad * gradient_accumulation_steps))
+                #       / (gradient_accumulation_steps  * world_size ) = avg_per_token_grad
                 counts = torch.tensor(1, device=losses.device)
             return self.strategy.reduce_loss(losses, counts, output_tensor, logps)
 

src/twinkle/model/optimizer_group.py

@@ -1,9 +1,8 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional
-
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
+from typing import Any, Dict, List, Optional
 
 from twinkle import DeviceMesh
 from twinkle.data_format import InputFeature, ModelOutput
@@ -83,4 +82,3 @@ def calculate_metrics(self, is_training):
         status.inputs = None
         status.outputs = None
         return results
-

src/twinkle/model/transformers/multi_lora_transformers.py

@@ -8,9 +8,9 @@
 from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union
 
 from twinkle import DeviceMesh, remote_class, remote_function, template
-from twinkle.infra import collect_tensor_dict
 from twinkle.data_format import InputFeature, Trajectory
 from twinkle.hub import HubOperation
+from twinkle.infra import collect_tensor_dict
 from twinkle.loss import Loss
 from twinkle.metric import Metric
 from twinkle.processor import InputProcessor

src/twinkle/model/transformers/transformers.py

@@ -474,13 +474,16 @@ def calculate_loss(self, **kwargs):
             counts = torch.tensor(1, device=loss_value.device)
         # Later will gather this value, so it becomes:
         # 1. SUM loss: gather_sum(local_num_tokens / dp_world_size) = global_num_tokens / dp_world_size
-        # 2. PER TOKEN MEAN loss: gather_sum(1 * gradient_accumulation_steps / dp_world_size ) = gradient_accumulation_steps
+        # 2. PER TOKEN MEAN loss: gather_sum(1 * gradient_accumulation_steps / dp_world_size )
+        #   = gradient_accumulation_steps
         # Then, grad will divided by this value:
         # 1. SUM loss: gather_mean(local_sum_grad) / (global_num_tokens / dp_world_size)
         #              = (global_sum_grad / dp_world_size) / (global_num_tokens / dp_world_size)
         #              = global_sum_grad/global_num_tokens
-        # 2. PER TOKEN MEAN loss: gather_mean(per_token_grad * gradient_accumulation_steps) / gradient_accumulation_steps
-        #                         = (global_per_token_grad * gradient_accumulation_steps / dp_world_size ) / gradient_accumulation_steps
+        # 2. PER TOKEN MEAN loss: gather_mean(per_token_grad * gradient_accumulation_steps)
+        #                               / gradient_accumulation_steps
+        #                         = (global_per_token_grad * gradient_accumulation_steps / dp_world_size )
+        #                               / gradient_accumulation_steps
         #                         = global_per_token_grad / dp_world_size = avg_per_token_grad
         counts = counts / self.device_mesh.data_world_size
         optimizer_config = self.optimizer_group[adapter_name]

src/twinkle/template/base.py

@@ -320,7 +320,7 @@ def _encode_messages(self, trajectory: Trajectory, add_generation_prompt: bool =
             input_ids = self.tokenizer.encode(text)
             encoded = {}
             labels = deepcopy(input_ids)
-        
+
         input_feature = InputFeature(
             input_ids=np.array(input_ids),
             labels=np.array(labels),
@@ -398,9 +398,7 @@ def batch_encode(
                 for key in trajectories:
                     if key in traj_keys:
                         # Encode this trajectory list
-                        result[key] = self.batch_encode(
-                            trajectories[key], add_generation_prompt=add_generation_prompt
-                        )
+                        result[key] = self.batch_encode(trajectories[key], add_generation_prompt=add_generation_prompt)
                     else:
                         # Keep non-trajectory columns as-is
                         result[key] = trajectories[key]

src/twinkle/utils/__init__.py

@@ -10,7 +10,8 @@
 from .parallel import processing_lock
 from .platforms import GPU, NPU, Platform, ensure_hccl_socket_env, ensure_npu_backend
 from .safetensors import LazyTensor, SafetensorLazyLoader, StreamingSafetensorSaver
-from .torch_utils import pad_sequence_to_length, selective_log_softmax, stateless_init_process_group, to_device, pad_and_stack_tensors
+from .torch_utils import (pad_and_stack_tensors, pad_sequence_to_length, selective_log_softmax,
+                          stateless_init_process_group, to_device)
 from .transformers_utils import find_all_linears, find_layers, get_modules_to_not_convert
 from .unsafe import check_unsafe, trust_remote_code
 from .utils import copy_files_by_pattern, deep_getattr