chore: import upstream snapshot with attribution

swift/trainers/__init__.py

@@ -0,0 +1,40 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+from . import patcher
+
+if TYPE_CHECKING:
+    from .arguments import Seq2SeqTrainingArguments, TrainArgumentsMixin, TrainingArguments
+    from .embedding_trainer import EmbeddingTrainer
+    from .mixin import DataLoaderMixin, SwiftMixin
+    from .reranker_trainer import RerankerTrainer
+    from .seq2seq_trainer import Seq2SeqTrainer
+    from .trainer import Trainer
+    from .trainer_factory import TrainerFactory
+    from .utils import (calculate_max_steps, disable_gradient_checkpointing, dynamic_gradient_checkpointing,
+                        per_token_loss_func)
+else:
+    _import_structure = {
+        'arguments': ['TrainArgumentsMixin', 'Seq2SeqTrainingArguments', 'TrainingArguments'],
+        'embedding_trainer': ['EmbeddingTrainer'],
+        'mixin': ['DataLoaderMixin', 'SwiftMixin'],
+        'reranker_trainer': ['RerankerTrainer'],
+        'seq2seq_trainer': ['Seq2SeqTrainer'],
+        'trainer': ['Trainer'],
+        'trainer_factory': ['TrainerFactory'],
+        'utils': [
+            'disable_gradient_checkpointing', 'dynamic_gradient_checkpointing', 'per_token_loss_func',
+            'calculate_max_steps'
+        ]
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/trainers/arguments.py

@@ -0,0 +1,299 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import math
+import os
+import platform
+from dataclasses import dataclass, field
+from transformers.training_args import TrainingArguments as HfTrainingArguments
+from transformers.training_args_seq2seq import Seq2SeqTrainingArguments as HfSeq2SeqTrainingArguments
+from typing import Dict, List, Literal, Optional, Union
+
+from swift.loss import loss_map
+from swift.utils import get_dist_setting, get_logger, is_liger_available, is_mp, json_parse_to_dict
+
+logger = get_logger()
+
+
+@dataclass
+class TrainArgumentsMixin:
+    """A dataclass mixin for configuring model training parameters.
+
+    Args:
+        per_device_train_batch_size (int): The batch size per GPU/TPU core for training. Defaults to 1.
+        per_device_eval_batch_size (int): The batch size per GPU/TPU core for evaluation. Defaults to 1.
+        gradient_accumulation_steps (Optional[int]): The number of update steps to accumulate gradients for before
+            performing an optimizer step.
+        tuner_backend (Optional[str]): The backend to use for parameter-efficient fine-tuning (e.g., 'peft'). Defaults
+            to None.
+        gradient_checkpointing (bool): If True, use gradient checkpointing to save memory at the cost of a slower
+            backward pass. Defaults to True.
+        vit_gradient_checkpointing (Optional[bool]): A specific gradient checkpointing setting for the Vision
+            Transformer part of the model. Defaults to None.
+        gradient_checkpointing_kwargs (Optional[Union[dict, str]]): Keyword arguments for
+            `torch.utils.checkpoint.checkpoint`. Defaults to None.
+        logging_first_step (bool): Whether to log the first global step. Defaults to True.
+        logging_steps (int): Log every `logging_steps` global steps. Defaults to 5.
+        router_aux_loss_coef (float): The coefficient for the router auxiliary loss in Mixture-of-Experts models.
+            Defaults to 0.0.
+        enable_dft_loss (bool): Whether to enable Diversity-from-Diversity (DFD) loss.
+            See https://arxiv.org/abs/2508.05629. Defaults to False.
+        enable_channel_loss (bool): Whether to enable channel loss. Defaults to False.
+        weight_decay (float): The weight decay to apply (if not zero) to all layers except bias and LayerNorm weights.
+            Defaults to 0.1.
+        adam_beta2 (float): The beta2 hyperparameter for the AdamW optimizer. Defaults to 0.95.
+        lr_scheduler_type (str): The learning rate scheduler type to use. Defaults to 'cosine'.
+        lr_scheduler_kwargs (Optional[Union[dict, str]]): Additional keyword arguments for the learning rate scheduler,
+            passed as a JSON string or a dictionary. Defaults to None.
+        report_to (List[str]): The list of integrations to report results to (e.g., 'tensorboard', 'wandb'). Defaults
+            to ['tensorboard']. If you specify `--report_to wandb`, you can set the project name through `WANDB_PROJECT`
+            and specify the API KEY corresponding to your account through `WANDB_API_KEY`.
+        dataloader_num_workers (Optional[int]): The number of subprocesses to use for data loading. Defaults to None.
+        dataloader_persistent_workers (bool): If True, the data loader workers will not be shut down after a dataset
+            has been consumed once. Defaults to False.
+        dataloader_prefetch_factor (Optional[int]): The number of batches loaded in advance by each worker. Defaults
+            to None.
+        use_liger_kernel (bool): Whether to use the Liger kernel for optimization. Defaults to False.
+        check_model (bool): If True, checks local model files for corruption or modification and provides a warning.
+            Should be set to False in an offline environment. Defaults to True.
+        acc_strategy (Literal['token', 'seq']): The strategy for calculating accuracy during training and validation.
+            Can be 'token' for token-level accuracy or 'seq' for sequence-level accuracy. Defaults to 'token'.
+        train_dataloader_shuffle (bool): Whether to shuffle the training data. Defaults to True.
+        group_by_length (bool): Whether to group samples with approximately the same length together in the
+            training dataset (with a random factor).
+        max_epochs (Optional[int]): The total number of training epochs to perform. Overrides `num_train_epochs`.
+            Defaults to None.
+        aligner_lr (Optional[float]): A specific learning rate for the aligner part of the model. Defaults to None.
+        vit_lr (Optional[float]): A specific learning rate for the Vision Transformer part of the model. Defaults to
+            None.
+        use_logits_to_keep (Optional[bool]): If enabled, reduces VRAM usage and speeds up training by calculating and
+            storing only the necessary logits based on the labels during the forward pass. If None, the behavior is
+            automatically determined. Defaults to None.
+        ds3_gather_for_generation (bool): In DeepSpeed ZeRO-3, whether to gather model parameters for generation.
+            Defaults to True.
+        resume_only_model (bool): When resuming from a checkpoint, whether to load only the model weights and not the
+            optimizer/scheduler states. Defaults to False.
+
+        optimizer (Optional[str]):The optimizer plugin to use (takes priority over `--optim`), default is None.
+            Available optimizers can be found in `optimizers/mapping.py`
+        loss_type (Optional[str]): Custom loss_type name. Default is None, uses the model's built-in loss function.
+            Available loss options can be found in `loss/mapping.py`
+        metric (Optional[str]): Custom eval metric name. Default is None. Available eval_metric options can be found
+            in `metrics/mapping.py`.
+        callbacks (List[str]): Custom trainer callbacks, default is `[]`. Available callbacks can be found
+            in `callbacks/mapping.py`.
+        early_stop_interval (Optional[int]): The interval for early stopping. Training will be terminated if the
+            `best_metric` does not improve for `early_stop_interval` evaluation periods (based on `save_steps`). It is
+            recommended to set `eval_steps` and `save_steps` to the same value. The implementation can be found in the
+            callback plugin. For more complex requirements, you can directly override the implementation in
+            `callback.py`. Defaults to None.
+
+        eval_use_evalscope (bool): Whether to use EvalScope for evaluation during training. Must be set to `True` to
+            enable it. Refer to examples for usage details. Defaults to False.
+        eval_dataset (List[str]): A list of evaluation dataset names. Multiple datasets can be specified, separated
+            by spaces.
+        eval_dataset_args (Optional[Union[str, dict]]): Arguments for the evaluation dataset(s), provided as a JSON
+            string or a dictionary.
+        eval_limit (Optional[int]): The maximum number of samples to use from the evaluation dataset. Defaults to None.
+        eval_generation_config (Optional[Union[str, dict]]): Model inference configuration for evaluation, provided as
+            a JSON string or a dictionary, e.g., `{'max_tokens': 512}`. Defaults to None.
+        extra_eval_args (Optional[Union[str, dict]]): Extra arguments for evaluation, provided as a JSON string or a
+            dictionary.
+
+        use_galore (bool): Flag to indicate if Galore is used. Default is False.
+        galore_target_modules (Optional[List[str]]): List of target modules for Galore. Default is None.
+        galore_rank (int): Rank for Galore. Default is 128.
+        galore_update_proj_gap (int): Update projection gap for Galore. Default is 50.
+        galore_scale (float): Scaling factor for Galore. Default is 1.0.
+        galore_proj_type (str): Projection type for Galore. Default is 'std'.
+        galore_optim_per_parameter (bool): Flag to indicate if optimization is per parameter for Galore.
+            Default is False.
+        galore_with_embedding (bool): Flag to indicate if embedding is used with Galore. Default is False.
+        galore_quantization (bool): Flag to indicate if use Q-Galore. Default is False.
+        galore_proj_quant (bool): Flag to indicate if projection quantization is used for Galore. Default is False.
+        galore_proj_bits (int): Number of bits for projection quantization. Default is 4.
+        galore_proj_group_size (int): Group size for projection quantization. Default is 256.
+        galore_cos_threshold (float): Cosine threshold for projection quantization. Default is 0.4.
+        galore_gamma_proj (int): Gamma for projection quantization. Default is 2.
+        galore_queue_size (int): Queue size for projection quantization. Default is 5.
+        lisa_activated_layers (int): Number of activated layers for LISA. Default is 0.
+        lisa_step_interval (int): Step interval for LISA activation. Default is 20.
+
+        use_flash_ckpt (bool): Whether to enable DLRover Flash Checkpoint. When enabled, weights are first saved to
+            shared memory and then asynchronously persisted to disk. Currently does not support the safetensors format.
+            It is recommended to use this with `PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"` to prevent CUDA OOM
+            errors during training. Defaults to False.
+    """
+    per_device_train_batch_size: int = 1
+    per_device_eval_batch_size: int = 1
+    gradient_accumulation_steps: Optional[int] = None
+    tuner_backend: Optional[str] = None
+
+    gradient_checkpointing: bool = True
+    vit_gradient_checkpointing: Optional[bool] = None
+    gradient_checkpointing_kwargs: Optional[Union[dict, str]] = None
+    logging_first_step: bool = True
+    logging_steps: int = 5
+    router_aux_loss_coef: float = 0.
+    enable_dft_loss: bool = False  # https://arxiv.org/abs/2508.05629
+    enable_channel_loss: bool = False
+    safe_serialization: bool = True
+    max_shard_size: str = '5GB'
+
+    weight_decay: float = 0.1
+    adam_beta2: float = 0.95
+    lr_scheduler_type: str = 'cosine'
+    lr_scheduler_kwargs: Optional[Union[dict, str]] = None
+    report_to: List[str] = field(default_factory=lambda: ['tensorboard'])
+    dataloader_num_workers: Optional[int] = None
+    dataloader_persistent_workers: bool = False
+    dataloader_prefetch_factor: Optional[int] = None
+    use_liger_kernel: bool = False
+
+    # extra
+    check_model: bool = True
+    acc_strategy: Literal['token', 'seq'] = 'token'
+    train_dataloader_shuffle: bool = True
+    group_by_length: bool = False
+    max_epochs: Optional[int] = None
+    aligner_lr: Optional[float] = None
+    vit_lr: Optional[float] = None
+    use_logits_to_keep: Optional[bool] = None
+    ds3_gather_for_generation: bool = True
+    resume_only_model: bool = False
+
+    # plugins
+    optimizer: Optional[str] = None
+    loss_type: Optional[str] = field(default=None, metadata={'help': f'loss_func choices: {list(loss_map.keys())}'})
+    # embedding (Matryoshka Representation Learning)
+    # Dict[int, float], where the key is the embedding dimension and the value is the corresponding loss weight,
+    # e.g. '{"32": 1.0, "64": 1.0, "128": 1.0}'.
+    mrl_dims: Optional[Union[dict, str]] = None
+    eval_metric: Optional[str] = None
+    callbacks: List[str] = field(default_factory=list)
+    # early_step
+    early_stop_interval: Optional[int] = None
+
+    # train-eval loop args
+    eval_use_evalscope: bool = False
+    eval_dataset: List[str] = field(default_factory=list)
+    eval_dataset_args: Optional[Union[str, dict]] = None
+    eval_limit: Optional[int] = None
+    eval_generation_config: Optional[Union[str, dict]] = None
+    extra_eval_args: Optional[Union[str, dict]] = None
+
+    # Value copied from SftArguments
+    tuner_type: Optional[str] = None
+
+    # galore
+    use_galore: bool = False
+    galore_target_modules: Optional[List[str]] = None
+    galore_rank: int = 128
+    galore_update_proj_gap: int = 50
+    galore_scale: float = 1.0
+    galore_proj_type: str = 'std'
+    galore_optim_per_parameter: bool = False
+    galore_with_embedding: bool = False
+    galore_quantization: bool = False
+    galore_proj_quant: bool = False
+    galore_proj_bits: int = 4
+    galore_proj_group_size: int = 256
+    galore_cos_threshold: float = 0.4
+    galore_gamma_proj: int = 2
+    galore_queue_size: int = 5
+    # lisa
+    lisa_activated_layers: int = 0
+    lisa_step_interval: int = 20
+
+    # dlrover flash_checkpoint
+    use_flash_ckpt: bool = False
+
+    @staticmethod
+    def _patch_liger_kernel():
+        # fix logits_to_keep
+        from liger_kernel.transformers.model import loss_utils
+        origin_LigerForCausalLMLoss = loss_utils.LigerForCausalLMLoss
+
+        def LigerForCausalLMLoss(hidden_states, *args, **kwargs):
+            hidden_states = hidden_states.contiguous()
+            for key in ['cu_seq_lens_q', 'cu_seq_lens_k', 'max_length_q', 'max_length_k']:
+                kwargs.pop(key, None)
+            return origin_LigerForCausalLMLoss(hidden_states, *args, **kwargs)
+
+        loss_utils.LigerForCausalLMLoss = LigerForCausalLMLoss
+
+    def _init_liger(self):
+        if self.use_liger_kernel:
+            assert is_liger_available(), 'use_liger_kernel requires liger_kernels, try `pip install liger-kernel`'
+            try:
+                self._patch_liger_kernel()
+            except Exception:
+                logger.warning('Failed to patch liger_kernel')
+
+    def _init_callbacks(self):
+        if self.lisa_activated_layers > 0:
+            self.callbacks.append('lisa')
+        if self.tuner_type == 'adalora':
+            self.callbacks.append('adalora')
+        if self.early_stop_interval is not None and self.early_stop_interval > 0:
+            self.callbacks.append('early_stop')
+        fsdp_config = getattr(self, 'fsdp_config', {})
+        if isinstance(fsdp_config, dict) and fsdp_config.get('activation_cpu_offload', False):
+            self.callbacks.append('activation_cpu_offload')
+
+    def __post_init__(self):
+        if hasattr(self, 'output_dir'):
+            self.output_dir = os.path.abspath(os.path.expanduser(self.output_dir))
+        if is_mp() and self.use_liger_kernel:
+            raise ValueError('liger_kernel does not support device_map. '
+                             'Please use DDP/DeepSpeed for multi-GPU training.')
+
+        if self.optimizer is None and (self.vit_lr is not None or self.aligner_lr is not None):
+            self.optimizer = 'multimodal'
+        self._init_callbacks()
+        if self.gradient_accumulation_steps is None:
+            world_size = get_dist_setting()[2]
+            self.gradient_accumulation_steps = max(1, math.ceil(16 / self.per_device_train_batch_size / world_size))
+            logger.info(f'Setting args.gradient_accumulation_steps: {self.gradient_accumulation_steps}')
+        if self.lr_scheduler_kwargs:
+            self.lr_scheduler_kwargs = json_parse_to_dict(self.lr_scheduler_kwargs)
+        if 'wandb' in self.report_to:
+            os.environ.setdefault('WANDB_PROJECT', 'ms-swift')
+        if self.vit_gradient_checkpointing is None:
+            self.vit_gradient_checkpointing = self.gradient_checkpointing
+        if self.gradient_checkpointing_kwargs:
+            self.gradient_checkpointing_kwargs = json_parse_to_dict(self.gradient_checkpointing_kwargs)
+        if self.mrl_dims is not None:
+            self.mrl_dims = json_parse_to_dict(self.mrl_dims)
+            self.mrl_dims = {int(k): float(v) for k, v in self.mrl_dims.items()}
+        self._init_liger()
+        if self.dataloader_num_workers is None:
+            if platform.system() == 'Windows':
+                self.dataloader_num_workers = 0
+            else:
+                self.dataloader_num_workers = 1
+            logger.info(f'Setting args.dataloader_num_workers: {self.dataloader_num_workers}')
+        if self.dataloader_prefetch_factor is None and self.dataloader_num_workers > 0:
+            self.dataloader_prefetch_factor = 2
+        if self.eval_use_evalscope:
+            try:
+                import evalscope
+            except ImportError:
+                raise ImportError('evalscope is not installed, please install it by `pip install evalscope`')
+            self.eval_dataset_args = json_parse_to_dict(self.eval_dataset_args)
+            self.eval_generation_config = json_parse_to_dict(self.eval_generation_config)
+            self.extra_eval_args = json_parse_to_dict(self.extra_eval_args)
+
+
+@dataclass
+class TrainingArguments(TrainArgumentsMixin, HfTrainingArguments):
+
+    def __post_init__(self):
+        TrainArgumentsMixin.__post_init__(self)
+        HfTrainingArguments.__post_init__(self)
+
+
+@dataclass
+class Seq2SeqTrainingArguments(TrainArgumentsMixin, HfSeq2SeqTrainingArguments):
+
+    def __post_init__(self):
+        TrainArgumentsMixin.__post_init__(self)
+        HfSeq2SeqTrainingArguments.__post_init__(self)

swift/trainers/embedding_trainer.py

@@ -0,0 +1,40 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch.nn.functional as F
+
+from swift.utils import get_logger
+from .trainer import Trainer
+from .utils import gather_for_unpadded_tensors
+
+logger = get_logger()
+
+
+class EmbeddingTrainer(Trainer):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gather_function = gather_for_unpadded_tensors
+        mrl_dims = self.args.mrl_dims
+        if mrl_dims and self.compute_loss_func is not None:
+            origin_loss_func = self.compute_loss_func
+
+            def mrl_loss_func(outputs, labels, **kwargs):
+                # Matryoshka Representation Learning: compute loss on each truncated dimension
+                # and aggregate with the corresponding weights.
+                last_hidden_state = outputs['last_hidden_state']
+                loss = None
+                for dim, weight in mrl_dims.items():
+                    if dim > last_hidden_state.shape[-1]:
+                        logger.warning_once(f'MRL: skipping dimension {dim} because it exceeds the model hidden size '
+                                            f'({last_hidden_state.shape[-1]}).')
+                        continue
+                    sliced = F.normalize(last_hidden_state[..., :dim], p=2, dim=-1)
+                    cur_loss = weight * origin_loss_func({'last_hidden_state': sliced}, labels, **kwargs)
+                    loss = cur_loss if loss is None else loss + cur_loss
+                return loss
+
+            self.compute_loss_func = mrl_loss_func
+
+    def evaluation_loop(self, *args, **kwargs):
+        output = super().evaluation_loop(*args, **kwargs)
+        self.gather_function = gather_for_unpadded_tensors
+        return output

swift/trainers/patcher.py

@@ -0,0 +1,108 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import math
+import os
+import time
+from tqdm import tqdm
+from transformers import trainer
+from transformers.trainer_callback import (DefaultFlowCallback, PrinterCallback, ProgressCallback, TrainerControl,
+                                           TrainerState)
+from transformers.trainer_utils import IntervalStrategy, has_length
+
+from swift.utils import append_to_jsonl, format_time, get_logger, get_max_reserved_memory, is_pai_training_job
+from .arguments import TrainingArguments
+
+logger = get_logger()
+
+
+def add_train_message(logs, state, start_time, start_step) -> None:
+    logs['global_step/max_steps'] = f'{state.global_step}/{state.max_steps}'
+    elapsed = time.time() - start_time
+    logs['elapsed_time'] = format_time(elapsed)
+    n_steps = state.global_step - start_step
+    train_speed = elapsed / n_steps if n_steps > 0 else 0.0
+    logs['remaining_time'] = format_time((state.max_steps - state.global_step) * train_speed)
+    for k, v in logs.items():
+        if isinstance(v, float):
+            logs[k] = round(logs[k], 8)
+    state.max_memory = max(getattr(state, 'max_memory', 0), get_max_reserved_memory())
+    if state.max_memory:
+        logs['memory(GiB)'] = round(state.max_memory, 2)
+    logs['train_speed(s/it)'] = round(train_speed, 6)
+
+
+class ProgressCallbackNew(ProgressCallback):
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        if state.is_world_process_zero:
+            self.training_bar = tqdm(desc='Train', total=state.max_steps, dynamic_ncols=True)
+        self.start_step = state.global_step
+        self.current_step = 0
+        self.start_time = time.time()
+
+    def on_prediction_step(self, args, state: TrainerState, control, eval_dataloader=None, **kwargs):
+        if state.is_world_process_zero and has_length(eval_dataloader):
+            if self.prediction_bar is None:
+                if self.training_bar is not None:
+                    self.training_bar.fp.write('\n')
+                self.prediction_bar = tqdm(
+                    desc='Val', total=len(eval_dataloader), leave=True, dynamic_ncols=True, position=0)
+            self.prediction_bar.update()
+
+    def on_log(self, args: TrainingArguments, state: TrainerState, control, logs=None, **kwargs):
+        add_train_message(logs, state, self.start_time, self.start_step)
+        if not is_pai_training_job() and state.is_world_process_zero:
+            jsonl_path = os.path.join(args.output_dir, 'logging.jsonl')
+            append_to_jsonl(jsonl_path, logs)
+        super().on_log(args, state, control, logs, **kwargs)
+        if state.is_world_process_zero and self.training_bar is not None:
+            self.training_bar.refresh()
+
+
+class DefaultFlowCallbackNew(DefaultFlowCallback):
+
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        control = super().on_step_end(args, state, control, **kwargs)
+        # save the last ckpt
+        evaluation_strategy = args.eval_strategy if hasattr(args, 'eval_strategy') else args.evaluation_strategy
+        if state.global_step == state.max_steps:
+            if evaluation_strategy != IntervalStrategy.NO:
+                control.should_evaluate = True
+            if args.save_strategy != IntervalStrategy.NO:
+                control.should_save = True
+        return control
+
+    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        control = super().on_epoch_end(args, state, control, **kwargs)
+        evaluation_strategy = args.eval_strategy if hasattr(args, 'eval_strategy') else args.evaluation_strategy
+        if args.max_epochs is not None and args.max_epochs <= math.ceil(state.epoch):
+            logger.info('Training has reached `max_epochs`. The model will be saved and the training will be exited.')
+            if evaluation_strategy != IntervalStrategy.NO:
+                control.should_evaluate = True
+            if args.save_strategy != IntervalStrategy.NO:
+                control.should_save = True
+            control.should_training_stop = True
+        return control
+
+
+class PrinterCallbackNew(PrinterCallback):
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        self.start_time = time.time()
+        self.start_step = state.global_step
+        return super().on_train_begin(args, state, control, **kwargs)
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        add_train_message(logs, state, self.start_time, self.start_step)
+        if not is_pai_training_job() and state.is_world_process_zero:
+            jsonl_path = os.path.join(args.output_dir, 'logging.jsonl')
+            append_to_jsonl(jsonl_path, logs)
+
+        _ = logs.pop('total_flos', None)
+        if state.is_world_process_zero:
+            print(logs, flush=True)
+
+
+# monkey patching
+trainer.DEFAULT_PROGRESS_CALLBACK = ProgressCallbackNew
+trainer.DEFAULT_CALLBACKS = [DefaultFlowCallbackNew]
+trainer.PrinterCallback = PrinterCallbackNew

swift/trainers/reranker_trainer.py

@@ -0,0 +1,50 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+
+from swift.utils import get_last_valid_indices, get_logger
+from .trainer import Trainer
+from .utils import gather_for_unpadded_tensors
+
+logger = get_logger()
+
+
+class RerankerTrainer(Trainer):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gather_function = gather_for_unpadded_tensors
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # Check if we have a custom loss function
+        if self.compute_loss_func is not None:
+            # Get labels and compute outputs
+            labels = inputs.pop('labels', None)
+            outputs = model(**inputs)
+            if self.task_type == 'generative_reranker':
+                logits = outputs.logits
+                attention_mask = inputs.get('attention_mask')
+                last_valid_indices = -1 if attention_mask is None else get_last_valid_indices(attention_mask)
+                batch_indices = torch.arange(logits.shape[0], device=logits.device)
+                outputs.logits = logits[batch_indices, last_valid_indices]
+
+            if labels is not None:
+                # Call custom loss function
+                loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch)
+            else:
+                # Fallback to model's loss
+                loss = outputs.loss
+
+            if num_items_in_batch is not None and self.model_accepts_loss_kwargs:
+                loss = loss / self.args.gradient_accumulation_steps
+
+            if labels is not None:
+                self._compute_acc(outputs, labels)
+
+            return (loss, outputs) if return_outputs else loss
+        else:
+            return super().compute_loss(model, inputs, return_outputs, num_items_in_batch)
+
+    def evaluation_loop(self, *args, **kwargs):
+        output = super().evaluation_loop(*args, **kwargs)
+        self.gather_function = gather_for_unpadded_tensors
+        return output

swift/trainers/seq2seq_trainer.py

@@ -0,0 +1,233 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Part of the implementation is borrowed from huggingface/transformers.
+import inspect
+import os
+import torch
+import torch.distributed as dist
+from contextlib import contextmanager, nullcontext
+from peft import PeftModel
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+from transformers import Seq2SeqTrainer as HfSeq2SeqTrainer
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.utils import is_peft_available
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from swift.infer_engine import InferRequest, RequestConfig, TransformersEngine
+from swift.sequence_parallel import sequence_parallel
+from swift.utils import HfConfigFactory, JsonlWriter, Serializer, gc_collect, get_logger, unwrap_model_for_generation
+from .arguments import Seq2SeqTrainingArguments
+from .mixin import DataLoaderMixin, SwiftMixin
+from .utils import per_token_loss_func, per_token_loss_func_sp
+
+logger = get_logger()
+
+
+class Seq2SeqTrainer(SwiftMixin, DataLoaderMixin, HfSeq2SeqTrainer):
+    args: Seq2SeqTrainingArguments
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model_accepts_loss_kwargs = True  # fix transformers>=4.46.2
+        if self.template.model_accepts_loss_kwargs is not None:
+            self.model_accepts_loss_kwargs = self.template.model_accepts_loss_kwargs
+        if self.args.predict_with_generate:
+            self.infer_engine = TransformersEngine(
+                self.model, template=self.template, max_batch_size=self.args.per_device_eval_batch_size)
+        self.jsonl_writer = JsonlWriter(os.path.join(self.args.output_dir, 'predict.jsonl'))
+
+    @staticmethod
+    def _predict_data_collator(batch):
+        return {'_data': batch}
+
+    @contextmanager
+    def _patch_predict_with_generate(self):
+        origin_data_collator = self.data_collator
+        self.data_collator = self._predict_data_collator
+        packing = self.template.packing
+        padding_free = self.template.padding_free
+        self.template.packing = False
+        self.template.padding_free = False
+        try:
+            yield
+        finally:
+            self.template.packing = packing
+            self.template.padding_free = padding_free
+            self.data_collator = origin_data_collator
+
+    def evaluate(self, *args, **kwargs):
+        context = self._patch_predict_with_generate() if self.args.predict_with_generate else nullcontext()
+        with context:
+            res = super().evaluate(*args, **kwargs)
+            gc_collect()
+            return res
+
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+        **gen_kwargs,
+    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        if not self.args.predict_with_generate or prediction_loss_only:
+            with self.template.forward_context(self.model, inputs):
+                return super().prediction_step(
+                    model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys)
+        data_list = inputs['_data']
+        labels_list = [InferRequest.remove_response(data['messages']) for data in data_list]
+        with unwrap_model_for_generation(
+                self.model_wrapped, self.accelerator,
+                gather_deepspeed3_params=self.args.ds3_gather_for_generation), self.template.generate_context():
+            resp_list = self.infer_engine.infer(
+                data_list,
+                RequestConfig(max_tokens=self.model.generation_config.max_new_tokens),
+                use_tqdm=False,
+            )
+
+        response_list = []
+        jsonl_cache = []
+        device = self.args.device
+        for data, resp, labels in zip(data_list, resp_list, labels_list):
+            response = resp.choices[0].message.content
+            jsonl_cache.append({'response': response, 'labels': labels, **data})
+            response_list.append(Serializer.to_tensor(resp.choices[0].message.content).to(device=device))
+        self.jsonl_writer.append(jsonl_cache, gather_obj=True)
+        labels_list = [Serializer.to_tensor(labels).to(device=device) for labels in labels_list]
+        response_list = pad_sequence(response_list, batch_first=True, padding_value=0)
+        labels_list = pad_sequence(labels_list, batch_first=True, padding_value=0)
+        return None, response_list, labels_list
+
+    def _prepare_inputs(self, inputs):
+        args = self.args
+        inputs = super()._prepare_inputs(inputs)
+        if self.template.sequence_parallel_size > 1:
+            sequence_parallel.prepare_inputs(inputs)
+
+        use_logits_to_keep = self.get_use_logits_to_keep(self.template.sequence_parallel_size == 1)
+        if use_logits_to_keep:
+            self.prepare_logits_to_keep(inputs)
+            if args.tuner_backend == 'unsloth' and isinstance(inputs['logits_to_keep'], torch.Tensor):
+                inputs['logits_to_keep'] = int(inputs['logits_to_keep'].sum())
+
+        base_model = self.template.get_base_model(self.model)
+        forward_params = inspect.signature(base_model.forward).parameters
+        if self.model.model_info.is_moe_model and any(key in forward_params
+                                                      for key in ['output_router_logits', 'kwargs']):
+            HfConfigFactory.set_config_attr(base_model.config, 'router_aux_loss_coef', args.router_aux_loss_coef)
+            base_model.router_aux_loss_coef = args.router_aux_loss_coef
+            logger.info_once(f'router_aux_loss_coef: {args.router_aux_loss_coef}')
+            if args.router_aux_loss_coef > 0:
+                inputs['output_router_logits'] = True
+        inputs['compute_loss_func'] = self.compute_loss_func
+        return inputs
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        labels = None
+        compute_loss_func: Callable = inputs.pop('compute_loss_func', None)
+        loss_scale = inputs.pop('loss_scale', None)
+        text_position_ids = inputs.pop('text_position_ids', None)
+        if text_position_ids is None:
+            text_position_ids = inputs.get('position_ids')
+        channels = inputs.pop('channel', None)
+
+        if (self.label_smoother is not None or compute_loss_func is not None or loss_scale is not None
+                or self.args.enable_dft_loss or self.args.enable_channel_loss
+                or self.template.sequence_parallel_size > 1) and 'labels' in inputs:
+            if self.args.use_liger_kernel:
+                logger.warning_once('The cross_entropy loss function defined in Liger Kernel will not '
+                                    'take effect, potentially leading to increased GPU memory consumption.')
+            labels = inputs.pop('labels')
+        outputs = self.template.compute_sft_loss(model, inputs, num_items_in_batch=num_items_in_batch, trainer=self)
+        mode = 'train' if self.model.training else 'eval'
+        if getattr(outputs, 'aux_loss', None) is not None:
+            self.custom_metrics[mode]['aux_loss'].update(outputs.aux_loss)
+        # Save past state if it exists
+        # TODO: this needs to be fixed and made cleaner later.
+        if hasattr(self.args, 'past_index') and self.args.past_index >= 0:
+            self._past = outputs[self.args.past_index]
+
+        if labels is None:
+            labels = inputs['labels']
+            if isinstance(outputs, dict) and 'loss' not in outputs:
+                raise ValueError(
+                    'The model did not return a loss from the inputs, only the following keys: '
+                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}.")
+            # We don't use .loss here since the model may return tuples instead of ModelOutput.
+            loss = outputs['loss'] if isinstance(outputs, dict) else outputs[0]
+        else:
+            outputs.loss = None
+            if (self.args.enable_dft_loss or loss_scale is not None or self.args.enable_channel_loss
+                    or self.template.sequence_parallel_size > 1):
+                if self.template.sequence_parallel_size > 1:
+                    outputs.loss = per_token_loss_func_sp(outputs, labels, enable_dft_loss=self.args.enable_dft_loss)
+                else:
+                    outputs.loss = per_token_loss_func(outputs, labels, enable_dft_loss=self.args.enable_dft_loss)
+
+                if loss_scale is not None:
+                    loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1).view(-1)
+                    outputs.loss = outputs.loss * loss_scale
+
+                if self.args.enable_channel_loss:
+                    metrics = self.custom_metrics[mode]
+                    masks = torch.roll(labels, shifts=-1, dims=-1).view(-1) != -100
+                    if self.template.padding_free:
+                        cu_seqlens = self.get_cu_seqlens(text_position_ids, inputs.get('logits_to_keep'))
+                    else:
+                        cu_seqlens = torch.arange(0, labels.shape[0] + 1) * labels.shape[1]
+                    for i in range(cu_seqlens.shape[0] - 1):
+                        channel = None if channels is None else channels[i]
+                        slice_ = slice(cu_seqlens[i], cu_seqlens[i + 1])
+                        metrics[f'loss_{channel}'].update(outputs.loss[slice_][masks[slice_]])
+
+            unwrapped_model = self.accelerator.unwrap_model(model)
+            if is_peft_available() and isinstance(unwrapped_model, PeftModel):
+                model_name = unwrapped_model.model._get_name()
+            else:
+                model_name = unwrapped_model._get_name()
+            # User-defined compute_loss function
+            if compute_loss_func is not None:
+                loss = compute_loss_func(
+                    outputs, labels, num_items_in_batch=num_items_in_batch, loss_scale=loss_scale, trainer=self)
+            elif self.label_smoother is None:
+                # Handle the outputs.loss generated by loss_scale.
+                if num_items_in_batch is None:
+                    # https://github.com/huggingface/transformers/blob/9dff7ca5c9693f4c02cdd2a9c2abc4772fcea5da/src/transformers/trainer.py#L2137
+                    num_items_in_batch = (labels != -100).sum()  # compat SP
+                    if self.template.sequence_parallel_size > 1:
+                        # labels are sharded by SP; outputs.loss was gathered
+                        # to full length via GatherLoss. Reduce the denominator
+                        # across the SP group so it matches the gathered loss.
+                        dist.all_reduce(num_items_in_batch, op=dist.ReduceOp.SUM)
+                loss = outputs.loss.sum() / num_items_in_batch
+            else:
+                if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
+                    loss = self.label_smoother(outputs, labels, shift_labels=True)
+                else:
+                    loss = self.label_smoother(outputs, labels)
+
+            if self.model.model_info.is_moe_model and self.args.router_aux_loss_coef is not None:
+                aux_loss = outputs.get('aux_loss')
+                if aux_loss is not None:
+                    if num_items_in_batch is not None:
+                        aux_loss = aux_loss * ((labels[:, 1:] != -100).sum() / num_items_in_batch)
+                    loss = loss + self.args.router_aux_loss_coef * aux_loss.to(loss.device)
+
+        if getattr(self.args, 'average_tokens_across_devices',
+                   False) and self.model_accepts_loss_kwargs and num_items_in_batch is not None:
+            loss *= self.accelerator.num_processes
+            if mode == 'eval' and self.template.sequence_parallel_size > 1:
+                loss /= self.template.sequence_parallel_size
+
+        if (outputs.logits is not None and labels is not None and self.args.tuner_backend != 'unsloth'):
+            cu_seqlens = None
+            if self.template.padding_free and self.args.acc_strategy == 'seq':
+                cu_seqlens = self.get_cu_seqlens(text_position_ids, inputs.get('logits_to_keep'))
+            # Liger does not have logits
+            # Unsloth has a bug with output logits
+            self._compute_acc(outputs, labels, cu_seqlens=cu_seqlens)
+        return (loss, outputs) if return_outputs else loss
+
+    def training_step(self, model, inputs, *args, **kwargs):
+        with self.template.forward_context(self.model, inputs):
+            return super().training_step(model, inputs, *args, **kwargs)

swift/trainers/trainer.py

@@ -0,0 +1,72 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Part of the implementation is borrowed from huggingface/transformers.
+import torch
+from contextlib import contextmanager
+from functools import wraps
+from peft import PeftModel
+from transformers import Trainer as HfTrainer
+
+from swift.sequence_parallel import sequence_parallel
+from swift.utils import get_logger
+from .arguments import TrainingArguments
+from .mixin import DataLoaderMixin, SwiftMixin
+
+logger = get_logger()
+
+
+class Trainer(SwiftMixin, DataLoaderMixin, HfTrainer):
+    args: TrainingArguments
+
+    def _prepare_inputs(self, inputs):
+        inputs = super()._prepare_inputs(inputs)
+        # For tasks whose `labels` are per-sample (e.g. seq_cls/reranker/embedding), we must NOT let
+        # SP code treat them as token labels. We detect that case by `labels.dim() == 1` and temporarily
+        # remove labels during `prepare_inputs`.
+        if self.template.sequence_parallel_size > 1:
+            labels = inputs.get('labels', None)
+            pop_labels = isinstance(labels, torch.Tensor) and labels.dim() == 1
+            if pop_labels:
+                labels = inputs.pop('labels', None)
+            try:
+                sequence_parallel.prepare_inputs(inputs)
+            finally:
+                if pop_labels and labels is not None:
+                    inputs['labels'] = labels
+        return inputs
+
+    @contextmanager
+    def _patch_loss_function(self):
+        model = self.model
+        if isinstance(model, PeftModel):
+            model = model.model
+        model_cls = model.__class__
+        if not hasattr(model_cls, 'loss_function'):
+            yield
+            return
+
+        loss_function = model.loss_function
+        _old_loss_function = model_cls.loss_function
+
+        @staticmethod
+        @wraps(loss_function)
+        def new_loss_function(logits, labels, **kwargs):
+            labels = labels.to(logits.device)  # fix device_map
+            return loss_function(logits=logits, labels=labels, **kwargs)
+
+        model_cls.loss_function = new_loss_function
+        try:
+            yield
+        finally:
+            model_cls.loss_function = _old_loss_function
+
+    def train(self, *args, **kwargs):
+        with self._patch_loss_function():
+            return super().train(*args, **kwargs)
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        loss, outputs = super().compute_loss(model, inputs, return_outputs=True)
+        if inputs.get('labels') is not None:
+            self._compute_acc(outputs, inputs['labels'])
+        if num_items_in_batch is not None and self.model_accepts_loss_kwargs:
+            loss = loss / self.args.gradient_accumulation_steps
+        return (loss, outputs) if return_outputs else loss

swift/trainers/trainer_factory.py

@@ -0,0 +1,73 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib.util
+import inspect
+from dataclasses import asdict
+from typing import Dict
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+class TrainerFactory:
+    TRAINER_MAPPING = {
+        'causal_lm': 'swift.trainers.Seq2SeqTrainer',
+        'seq_cls': 'swift.trainers.Trainer',
+        'embedding': 'swift.trainers.EmbeddingTrainer',
+        'reranker': 'swift.trainers.RerankerTrainer',
+        'generative_reranker': 'swift.trainers.RerankerTrainer',
+        # rlhf
+        'dpo': 'swift.rlhf_trainers.DPOTrainer',
+        'orpo': 'swift.rlhf_trainers.ORPOTrainer',
+        'kto': 'swift.rlhf_trainers.KTOTrainer',
+        'cpo': 'swift.rlhf_trainers.CPOTrainer',
+        'rm': 'swift.rlhf_trainers.RewardTrainer',
+        'ppo': 'swift.rlhf_trainers.PPOTrainer',
+        'grpo': 'swift.rlhf_trainers.GRPOTrainer',
+        'gkd': 'swift.rlhf_trainers.GKDTrainer',
+    }
+
+    TRAINING_ARGS_MAPPING = {
+        'causal_lm': 'swift.trainers.Seq2SeqTrainingArguments',
+        'seq_cls': 'swift.trainers.TrainingArguments',
+        'embedding': 'swift.trainers.TrainingArguments',
+        'reranker': 'swift.trainers.TrainingArguments',
+        'generative_reranker': 'swift.trainers.TrainingArguments',
+        # rlhf
+        'dpo': 'swift.rlhf_trainers.DPOConfig',
+        'orpo': 'swift.rlhf_trainers.ORPOConfig',
+        'kto': 'swift.rlhf_trainers.KTOConfig',
+        'cpo': 'swift.rlhf_trainers.CPOConfig',
+        'rm': 'swift.rlhf_trainers.RewardConfig',
+        'ppo': 'swift.rlhf_trainers.PPOConfig',
+        'grpo': 'swift.rlhf_trainers.GRPOConfig',
+        'gkd': 'swift.rlhf_trainers.GKDConfig',
+    }
+
+    @staticmethod
+    def get_cls(args, mapping: Dict[str, str]):
+        if hasattr(args, 'rlhf_type'):
+            train_method = args.rlhf_type
+        else:
+            train_method = args.task_type
+        module_path, class_name = mapping[train_method].rsplit('.', 1)
+        module = importlib.import_module(module_path)
+        return getattr(module, class_name)
+
+    @classmethod
+    def get_trainer_cls(cls, args):
+        return cls.get_cls(args, cls.TRAINER_MAPPING)
+
+    @classmethod
+    def get_training_args(cls, args):
+        training_args_cls = cls.get_cls(args, cls.TRAINING_ARGS_MAPPING)
+        args_dict = asdict(args)
+        parameters = inspect.signature(training_args_cls).parameters
+
+        for k in list(args_dict.keys()):
+            if k not in parameters:
+                args_dict.pop(k)
+
+        args._prepare_training_args(args_dict)
+        training_args = training_args_cls(**args_dict)
+        return training_args

swift/trainers/utils.py

@@ -0,0 +1,410 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Part of the implementation is borrowed from huggingface/transformers.
+import inspect
+import math
+import os
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from contextlib import contextmanager
+from modelscope.hub.api import HubApi
+from peft import PeftModel
+from torch import nn
+from torch.nn import CrossEntropyLoss, Module
+from transformers import PreTrainedModel
+from types import FunctionType, MethodType
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from swift.model import ModelMeta
+from swift.sequence_parallel import ChunkedCrossEntropyLoss, GatherLoss, sequence_parallel
+from swift.utils import deep_getattr, get_dist_setting, get_logger
+
+if TYPE_CHECKING:
+    from .arguments import TrainingArguments
+
+logger = get_logger()
+
+
+def _get_deepspeed_elastic_world_size():
+    if dist.is_available() and dist.is_initialized():
+        return dist.get_world_size()
+    return get_dist_setting()[2]
+
+
+def _enable_load_universal(ds_config):
+    if isinstance(ds_config, dict):
+        checkpoint = ds_config.get('checkpoint')
+        if not isinstance(checkpoint, dict):
+            checkpoint = {}
+            ds_config['checkpoint'] = checkpoint
+        checkpoint['load_universal'] = True
+
+
+def enable_deepspeed_load_universal(args: 'TrainingArguments', trainer=None):
+    _enable_load_universal(getattr(args, 'deepspeed', None))
+
+    hf_ds_config = getattr(args, 'hf_deepspeed_config', None)
+    _enable_load_universal(getattr(hf_ds_config, 'config', None))
+
+    deepspeed_plugin = getattr(args, 'deepspeed_plugin', None)
+    if trainer is not None and deepspeed_plugin is None:
+        accelerator = getattr(trainer, 'accelerator', None)
+        state = getattr(accelerator, 'state', None)
+        deepspeed_plugin = getattr(state, 'deepspeed_plugin', None)
+    if deepspeed_plugin is not None:
+        _enable_load_universal(getattr(deepspeed_plugin, 'deepspeed_config', None))
+        plugin_hf_ds_config = getattr(deepspeed_plugin, 'hf_ds_config', None)
+        _enable_load_universal(getattr(plugin_hf_ds_config, 'config', None))
+
+
+def prepare_deepspeed_elastic_config(args: 'TrainingArguments', state=None):
+    ds_config = args.deepspeed
+    if not ds_config:
+        return
+    if not isinstance(ds_config, dict):
+        logger.warning('DeepSpeed elastic expects args.deepspeed to be a dict, but got '
+                       f'{type(ds_config).__name__}. Skip elastic config.')
+        return
+
+    from deepspeed.elasticity import compute_elastic_config
+    from deepspeed.git_version_info import version as __version__
+
+    enable_deepspeed_load_universal(args)
+    elasticity = ds_config.get('elasticity') or {}
+    if not elasticity:
+        logger.warning_once('DeepSpeed elastic callback is enabled, but no `elasticity` section is found in '
+                            'the DeepSpeed config. Only `checkpoint.load_universal` is enabled.')
+        return
+    if elasticity.get('enabled') is False:
+        return
+
+    world_size = _get_deepspeed_elastic_world_size()
+    final_batch_size, _, micro_batch_size = compute_elastic_config(
+        ds_config=ds_config,
+        target_deepspeed_version=__version__,
+        world_size=world_size,
+    )
+    if world_size <= 0 or micro_batch_size <= 0:
+        raise ValueError('DeepSpeed elastic config produced invalid batch settings: '
+                         f'world_size={world_size}, micro_batch_size={micro_batch_size}.')
+    gradient_accu_steps = max(1, final_batch_size // (micro_batch_size * world_size))
+    args.per_device_train_batch_size = micro_batch_size
+    args.gradient_accumulation_steps = gradient_accu_steps
+    if state is not None:
+        state.train_batch_size = args.per_device_train_batch_size * max(1, args.n_gpu)
+    logger.info_once('DeepSpeed elastic config is enabled. '
+                     f'world_size: {world_size}, '
+                     f'per_device_train_batch_size: {args.per_device_train_batch_size}, '
+                     f'gradient_accumulation_steps: {args.gradient_accumulation_steps}')
+
+
+def can_return_loss(model: Module) -> bool:
+    """Check if a given model can return loss."""
+    if isinstance(model, PeftModel):
+        signature = inspect.signature(model.model.forward)
+    else:
+        signature = inspect.signature(model.forward)
+    for p in signature.parameters:
+        if p == 'return_loss' and signature.parameters[p].default is True:
+            return True
+    return False
+
+
+def find_labels(model: Module) -> List[str]:
+    """Find the labels used by a given model."""
+    model_name = model.__class__.__name__
+    if isinstance(model, PeftModel):
+        signature = inspect.signature(model.model.forward)
+    else:
+        signature = inspect.signature(model.forward)
+    if 'QuestionAnswering' in model_name:
+        return [p for p in signature.parameters if 'label' in p or p in ('start_positions', 'end_positions')]
+    else:
+        return [p for p in signature.parameters if 'label' in p]
+
+
+def get_function(method_or_function: Union[MethodType, FunctionType]) -> FunctionType:
+    if isinstance(method_or_function, MethodType):
+        method_or_function = method_or_function.__func__
+    return method_or_function
+
+
+def is_instance_of_ms_model(model: Module) -> bool:
+    """avoid import modelscope: circular dependency problem"""
+    for m_cls in model.__class__.__mro__:
+        cls_name = m_cls.__name__
+        cls_module = m_cls.__module__
+        if cls_name == 'Model' and cls_module.startswith('modelscope'):
+            return True
+    return False
+
+
+def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, **kwargs) -> torch.Tensor:
+    """Common loss function for sequence parallel training"""
+    if hasattr(outputs, 'logits'):
+        logits = outputs.logits
+    else:
+        logits = outputs
+    device = logits.device
+
+    batch_size = logits.shape[0]
+    logits = logits.view(-1, logits.shape[-1])
+    labels = labels.flatten().to(device)
+    sploss_parallel_size = int(os.environ.get('CELOSS_PARALLEL_SIZE', '0'))
+    if sploss_parallel_size > 0:
+        loss = ChunkedCrossEntropyLoss.apply(logits, labels, sploss_parallel_size)
+    else:
+        loss_fct = CrossEntropyLoss(reduction='none')
+        loss = loss_fct(logits, labels)
+    if enable_dft_loss:
+        with torch.no_grad():
+            target_probs = torch.exp(-loss)
+        loss *= target_probs
+    position_ids = sequence_parallel.real_position_ids
+    if position_ids is not None:
+        position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
+    loss, labels = GatherLoss.apply(loss.reshape(batch_size, -1), labels.reshape(batch_size, -1), 1, position_ids)
+    if position_ids is not None and position_ids.min() == -1:
+        _pos_mask = position_ids >= 0
+        loss = loss[_pos_mask].contiguous()
+
+    return loss
+
+
+def per_token_loss_func(outputs, labels, enable_dft_loss: bool = False, **kwargs):
+    logits = outputs.logits
+    # Upcast to float if we need to compute the loss to avoid potential precision issues
+    logits = logits.float()
+    labels = torch.roll(labels, shifts=-1, dims=-1).view(-1)
+
+    # Flatten the tokens
+    logits = logits.view(-1, logits.shape[-1])
+    # Enable model parallelism
+    labels = labels.to(logits.device)
+    loss = F.cross_entropy(logits, labels, ignore_index=-100, reduction='none')
+    if enable_dft_loss:
+        with torch.no_grad():
+            target_probs = torch.exp(-loss)
+        loss *= target_probs
+    return loss
+
+
+def _kwargs_to_args(func, args, kwargs) -> Optional[List[Any]]:
+    parameters = inspect.signature(func).parameters
+    args = list(args)
+    parameters = list(parameters.items())[len(args):]
+    for key, param in parameters:
+        if key in kwargs:
+            args.append(kwargs[key])
+        elif param.default != param.empty:
+            args.append(param.default)
+        else:
+            return
+    return args
+
+
+def _add_gradient_checkpointing(module_list):
+
+    requires_grad = None
+
+    def _new_forward(self, *args, **kwargs):
+        nonlocal requires_grad
+        if requires_grad is None:
+            requires_grad = any(p.requires_grad for p in self.parameters())
+
+        new_args = _kwargs_to_args(self.__old_forward, args, kwargs)
+        if new_args is not None and self.gradient_checkpointing and self.training:
+            if new_args and isinstance(new_args[0], torch.Tensor) and requires_grad and not new_args[0].requires_grad:
+                new_args[0].requires_grad_(True)
+            layer_ret = self._gradient_checkpointing_func(self.__old_forward, *new_args)
+            logger.info_once('Successfully using dynamic gradient checkpointing.')
+        else:
+            layer_ret = self.__old_forward(*args, **kwargs)
+        return layer_ret
+
+    for module in module_list:
+        module.gradient_checkpointing = False
+        if hasattr(module, '_old_forward'):  # device_map
+            __old_forward = module._old_forward
+            module._old_forward = MethodType(_new_forward, module)
+        else:
+            __old_forward = module.forward
+            module.forward = MethodType(_new_forward, module)
+        module.__old_forward = __old_forward
+
+
+def find_module_list(model) -> Optional[nn.ModuleList]:
+    module_lists = []
+    for m in model.modules():
+        if hasattr(m, 'gradient_checkpointing') or m.__class__.__name__ == 'CheckpointWrapper':
+            return
+        if (isinstance(m, (nn.ModuleList, nn.Sequential)) and len(m) >= 10
+                and 'mlp' not in m[0].__class__.__name__.lower()):  # fix moe
+            module_lists.append(m)
+    if module_lists:
+        return max(module_lists, key=lambda x: len(x))
+
+
+def dynamic_gradient_checkpointing(model, including_vit: bool = False) -> None:
+    if isinstance(model, PeftModel):
+        model = model.model
+    model_meta: ModelMeta = getattr(model, 'model_meta', None)
+    if model_meta is not None and model_meta.is_multimodal and model_meta.model_arch:
+        tower_names = model_meta.model_arch.language_model.copy()
+        if including_vit:
+            tower_names += model_meta.model_arch.vision_tower
+    else:
+        tower_names = [None]
+
+    model.supports_gradient_checkpointing = True
+    for tower_name in tower_names:
+        if tower_name is None:
+            model_tower = model
+        else:
+            model_tower = deep_getattr(model, tower_name)
+        if model_tower is None:
+            continue
+        model_tower.supports_gradient_checkpointing = True
+        module_list = find_module_list(model_tower)
+        if module_list is None:
+            continue
+        _add_gradient_checkpointing(module_list)
+        logger.info(f'Automatically add gradient_checkpointing to {model_tower.__class__}.')
+
+
+@contextmanager
+def disable_gradient_checkpointing(model: PreTrainedModel, gradient_checkpointing_kwargs: Optional[Dict] = None):
+    """
+    Temporarily disable gradient checkpointing, restoring the previous state afterward.
+
+    When gradient checkpointing is enabled with use_reentrant=True (default), calling the model inside a
+    torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True").
+    Temporarily disable checkpointing to avoid this warning during inference.
+
+    Args:
+        model (`PreTrainedModel`):
+            Model for which to temporarily disable gradient checkpointing.
+        gradient_checkpointing_kwargs (`dict` or `None`, *optional*):
+            Additional kwargs for gradient checkpointing enabling.
+    """
+    was_enabled = getattr(model, 'is_gradient_checkpointing', False)
+    if was_enabled:
+        model.gradient_checkpointing_disable()
+    try:
+        yield
+    finally:
+        if was_enabled:
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs)
+
+
+def gather_for_unpadded_tensors(input_data, use_gather_object=False):
+    from accelerate.utils import gather_object
+    if getattr(sequence_parallel, 'dp_group', None) is not None:
+        input_data = sequence_parallel._gather_object_dp(input_data)
+    else:
+        input_data = gather_object(input_data)
+    output = []
+    for _data in input_data:
+        if len(_data.shape) == 0:
+            _data = _data.unsqueeze(0)
+        _data = _data.cpu()
+        output.append(_data)
+    if len(output[0].shape) == 1 and output[0].shape[0] > 1:
+        data = torch.stack(output, dim=0)
+    else:
+        data = torch.concat(output, dim=0)
+    return data
+
+
+def calculate_max_steps(args: 'TrainingArguments', dataset) -> int:
+    if args.max_steps and args.max_steps > 0:
+        max_steps = args.max_steps
+    else:
+        len_dataset = len(dataset)
+        _, _, world_size, _ = get_dist_setting()
+        total_train_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps * world_size
+        num_update_steps_per_epoch = len_dataset // total_train_batch_size
+        num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
+        max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
+    return max_steps
+
+
+def extract_version(name: str) -> Optional[int]:
+    if not name.startswith('v'):
+        return None
+    try:
+        num = name[1:].split('-', 1)[0]
+        return int(num)
+    except ValueError:
+        return None
+
+
+def get_previous_version_from_path(current_path: str) -> Optional[str]:
+    from pathlib import Path
+    current = Path(current_path)
+    parent = current.parent
+    current_name = current.name
+
+    candidates = [d for d in parent.iterdir() if d.is_dir()]
+
+    valid = [(d.name, extract_version(d.name)) for d in candidates]
+    valid = [(name, ver) for name, ver in valid if ver is not None]
+
+    valid.sort(key=lambda x: x[1])
+    names = [name for name, _ in valid]
+
+    if current_name not in names:
+        return None
+
+    idx = names.index(current_name)
+    if idx == 0:
+        return None
+
+    prev_name = names[idx - 1]
+    return str(parent / prev_name)
+
+
+def get_resume_dir(output_dir):
+    return get_previous_version_from_path(output_dir)
+
+
+def replace_index_file(output_dir: str):
+    import json
+    import os
+    from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, WEIGHTS_INDEX_NAME
+    index_file = os.path.join(output_dir, WEIGHTS_INDEX_NAME)
+
+    if not os.path.exists(index_file):
+        return
+    with open(index_file, 'r', encoding='utf-8') as f:
+        bin_data = json.load(f)
+    if 'weight_map' not in bin_data:
+        return
+    bin_data['weight_map'] = {
+        k: v.replace('pytorch_model', 'model').replace('.bin', '.safetensors')
+        for k, v in bin_data['weight_map'].items()
+    }
+    safe_path = os.path.join(output_dir, SAFE_WEIGHTS_INDEX_NAME)
+    with open(safe_path, 'w', encoding='utf-8') as f:
+        json.dump(bin_data, f, indent=2)
+    from contextlib import suppress
+    with suppress(FileNotFoundError):
+        os.remove(os.path.join(output_dir, WEIGHTS_INDEX_NAME))
+
+
+@contextmanager
+def patch_modelscope_hub_timeout():
+    __init__ = HubApi.__init__
+
+    def __new_init__(self, *args, **kwargs):
+        timeout = kwargs.get('timeout')
+        if timeout is not None and timeout > 5:
+            kwargs['timeout'] = 5
+        __init__(self, *args, **kwargs)
+
+    HubApi.__init__ = __new_init__
+
+    try:
+        yield
+    finally:
+        HubApi.__init__ = __init__