Add files using upload-large-folder tool

ms-swift/.ipynb_checkpoints/clean_transcripts-checkpoint.py

@@ -0,0 +1,95 @@
+import json
+import re
+from typing import List, Dict, Tuple
+
+def parse_timestamp(timestamp: str) -> Tuple[int, int]:
+    """Convert timestamp string like '00:15' to seconds."""
+    minutes, seconds = map(int, timestamp.split(':'))
+    return minutes * 60 + seconds
+
+def extract_time_and_speaker(line: str) -> Tuple[Tuple[int, int], str]:
+    """Extract time range and speaker from a line."""
+    # Extract time range
+    time_match = re.match(r'\[(\d{2}:\d{2}) - (\d{2}:\d{2})\] (Speaker [A-Z]):', line)
+    if not time_match:
+        return None, None
+    
+    start_time = parse_timestamp(time_match.group(1))
+    end_time = parse_timestamp(time_match.group(2))
+    speaker = time_match.group(3)
+    
+    return (start_time, end_time), speaker
+
+def has_overlap(range1: Tuple[int, int], range2: Tuple[int, int]) -> bool:
+    """Check if two time ranges overlap."""
+    start1, end1 = range1
+    start2, end2 = range2
+    return not (end1 <= start2 or end2 <= start1)
+
+def has_same_speaker_overlap(transcript: str) -> bool:
+    """Check if a transcript contains overlapping timestamps for the same speaker."""
+    lines = transcript.split('\n')
+    # Dictionary to store time ranges for each speaker
+    speaker_ranges = {}
+    
+    for line in lines:
+        if not line.strip():
+            continue
+            
+        time_range, speaker = extract_time_and_speaker(line)
+        if time_range is None or speaker is None:
+            continue
+            
+        # Check for overlaps with existing ranges of the same speaker
+        if speaker in speaker_ranges:
+            for existing_range in speaker_ranges[speaker]:
+                if has_overlap(time_range, existing_range):
+                    return True
+            
+            speaker_ranges[speaker].append(time_range)
+        else:
+            speaker_ranges[speaker] = [time_range]
+    
+    return False
+
+def process_file(input_file: str, output_file: str, delete_file: str):
+    """Process the JSON file and separate entries with same-speaker overlapping timestamps."""
+    with open(input_file, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+        
+    if isinstance(data, dict):
+        data = [data]
+        
+    cleaned_data = []
+    deleted_data = []
+    removed_count = 0
+    
+    for entry in data:
+        if 'model_output' in entry:
+            if not has_same_speaker_overlap(entry['model_output']):
+                cleaned_data.append(entry)
+            else:
+                deleted_data.append(entry)
+                removed_count += 1
+                print(f"Removing entry with key: {entry.get('key', 'unknown')}")
+    
+    # Save cleaned data
+    with open(output_file, 'w', encoding='utf-8') as f:
+        json.dump(cleaned_data, f, ensure_ascii=False, indent=2)
+    
+    # Save deleted data
+    with open(delete_file, 'w', encoding='utf-8') as f:
+        json.dump(deleted_data, f, ensure_ascii=False, indent=2)
+    
+    print(f"\nProcessing Summary:")
+    print(f"Processed {len(data)} entries")
+    print(f"Removed {removed_count} entries with same-speaker overlapping timestamps")
+    print(f"Remaining entries: {len(cleaned_data)}")
+
+if __name__ == '__main__':
+    input_file = 'silence_overlaps/transcriptions.json'
+    output_file = 'silence_overlaps/cleaned_transcriptions2.json'
+    delete_file = 'silence_overlaps/delete_transcript2.json'
+    process_file(input_file, output_file, delete_file)
+    print(f"\nCleaned transcriptions have been saved to {output_file}")
+    print(f"Deleted entries have been saved to {delete_file}") 

ms-swift/swift/megatron/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+try:
+    from .init import init_megatron_env
+    init_megatron_env()
+except Exception:
+    # allows lint pass.
+    raise
+
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .train import megatron_sft_main, megatron_pt_main
+    from .utils import convert_hf2mcore, convert_mcore2hf
+    from .argument import MegatronTrainArguments
+    from .model import MegatronModelType, MegatronModelMeta, get_megatron_model_meta, register_megatron_model
+else:
+    _import_structure = {
+        'train': ['megatron_sft_main', 'megatron_pt_main'],
+        'utils': ['convert_hf2mcore', 'convert_mcore2hf'],
+        'argument': ['MegatronTrainArguments'],
+        'model': ['MegatronModelType', 'MegatronModelMeta', 'get_megatron_model_meta', 'register_megatron_model']
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

ms-swift/swift/megatron/argument/megatron_args.py

@@ -0,0 +1,253 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import sys
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+import torch
+from transformers.utils.versions import require_version
+
+from swift.llm.argument.base_args import to_abspath
+
+
+@dataclass
+class ExtraMegatronArguments:
+    padded_vocab_size: Optional[int] = None
+    rope_scaling: Optional[Union[dict, str]] = None
+    torch_dtype: Optional[torch.dtype] = None
+
+    dataloader_persistent_workers: bool = True
+    dataloader_prefetch_factor: int = 10
+
+    model_type: Optional[str] = None
+    max_epochs: Optional[int] = None
+
+
+@dataclass
+class MegatronArguments(ExtraMegatronArguments):
+    # training
+    micro_batch_size: int = 1
+    global_batch_size: int = 16
+    recompute_granularity: Literal['selective', 'full'] = 'selective'
+    recompute_method: Literal['uniform', 'block'] = None
+    recompute_num_layers: Optional[int] = None
+    recompute_modules: List[str] = field(default_factory=lambda: ['core_attn'])
+    use_cpu_initialization: bool = False
+    deterministic_mode: bool = False
+    train_iters: Optional[int] = None
+    log_interval: int = 5
+    tensorboard_dir: Optional[str] = None
+    no_masked_softmax_fusion: bool = False
+    no_bias_dropout_fusion: bool = False
+    no_bias_swiglu_fusion: bool = False
+    no_rope_fusion: bool = False
+    no_gradient_accumulation_fusion: bool = False
+    cross_entropy_loss_fusion: bool = False
+    calculate_per_token_loss: bool = True
+    use_flash_attn: bool = False
+    attention_backend: str = 'auto'  # flash, fused, unfused, local, auto
+    optimizer: Literal['adam', 'sgd'] = 'adam'
+    dataloader_type: Literal['single', 'cyclic', 'external'] = 'cyclic'
+    manual_gc: bool = False
+    manual_gc_interval: int = 0
+
+    # learning rate
+    lr: float = 1e-5
+    lr_decay_style: Literal['cosine', 'linear', 'constant'] = 'cosine'
+    # The default is None, which will be set to `train_iters`.
+    lr_decay_iters: Optional[int] = None
+    lr_warmup_iters: int = 0
+    min_lr: float = 0
+
+    # regularization
+    weight_decay: float = 0.1
+    clip_grad: float = 1.
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.95
+    adam_eps: float = 1e-8
+    sgd_momentum: float = 0.9
+
+    # checkpoint
+    save: Optional[str] = None
+    save_interval: int = 500
+    no_save_optim: bool = False
+    no_save_rng: bool = False
+    load: Optional[str] = None
+    no_load_optim: bool = False
+    no_load_rng: bool = False
+    finetune: bool = False
+    ckpt_format: Literal['torch', 'torch_dist', 'zarr'] = 'torch_dist'
+    no_initialization: bool = True
+    auto_detect_ckpt_format: bool = True
+    exit_on_missing_checkpoint: bool = True
+
+    # dist
+    distributed_backend: Literal['nccl', 'gloo'] = 'nccl'
+    use_distributed_optimizer: bool = True
+    tensor_model_parallel_size: int = 1
+    pipeline_model_parallel_size: int = 1
+    decoder_first_pipeline_num_layers: Optional[int] = None
+    decoder_last_pipeline_num_layers: Optional[int] = None
+    sequence_parallel: bool = False
+    context_parallel_size: int = 1
+    tp_comm_overlap: bool = False
+    overlap_grad_reduce: bool = False
+    overlap_param_gather: bool = False
+    distributed_timeout_minutes: int = 60
+
+    # model
+    num_layers: Optional[int] = None
+    hidden_size: Optional[int] = None
+    ffn_hidden_size: Optional[int] = None
+    num_attention_heads: Optional[int] = None
+    group_query_attention: Optional[bool] = None
+    num_query_groups: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    position_embedding_type: Literal['learned_absolute', 'rope', 'relative', 'none'] = 'rope'
+    rotary_base: Optional[int] = None
+    rotary_percent: float = 1.
+    normalization: Literal['LayerNorm', 'RMSNorm'] = 'RMSNorm'
+    norm_epsilon: Optional[float] = None
+    swiglu: Optional[bool] = None
+    untie_embeddings_and_output_weights: Optional[bool] = None
+    disable_bias_linear: Optional[bool] = None
+    add_qkv_bias: Optional[bool] = None
+    attention_dropout: Optional[float] = None
+    hidden_dropout: float = 0.
+    kv_channels: Optional[int] = None
+    qk_layernorm: Optional[bool] = None
+    transformer_impl: Literal['local', 'transformer_engine'] = 'transformer_engine'
+
+    # moe
+    num_experts: Optional[int] = None
+    moe_ffn_hidden_size: Optional[int] = None
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_router_topk: Optional[int] = None
+    moe_router_pre_softmax: Optional[bool] = None
+    moe_aux_loss_coeff: Optional[float] = None
+
+    expert_model_parallel_size: int = 1
+    moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'alltoall_seq'] = 'alltoall'
+    moe_grouped_gemm: bool = False
+    moe_router_load_balancing_type: Literal['aux_loss', 'seq_aux_loss', 'sinkhorn', 'none'] = 'aux_loss'
+    moe_z_loss_coeff: Optional[float] = None
+    moe_expert_capacity_factor: Optional[float] = None
+    moe_shared_expert_overlap: bool = False
+
+    # mixed precision
+    fp16: Optional[bool] = None
+    bf16: Optional[bool] = None
+    apply_query_key_layer_scaling: Optional[bool] = None
+    attention_softmax_in_fp32: bool = True
+
+    # logging
+    log_params_norm: bool = False
+    log_throughput: bool = True
+    tensorboard_log_interval: int = 1
+    tensorboard_queue_size: int = 50
+    log_timers_to_tensorboard: bool = True
+    no_log_learning_rate_to_tensorboard: bool = False
+    log_validation_ppl_to_tensorboard: bool = True
+    log_memory_to_tensorboard: bool = True
+    logging_level: Optional[str] = None
+    wandb_project: Optional[str] = None
+    wandb_exp_name: Optional[str] = None
+    wandb_save_dir: Optional[str] = None
+
+    # evaluate
+    eval_iters: int = 100
+    eval_interval: Optional[int] = None
+
+    # other
+    seed: int = 42
+    seq_length: Optional[int] = None
+    num_workers: int = 4
+    no_create_attention_mask_in_dataloader: bool = True
+
+    def _set_default(self):
+        if self.num_query_groups is None:
+            self.num_query_groups = 1
+        if self.norm_epsilon is None:
+            self.norm_epsilon = 1e-5
+        if self.rotary_base is None:
+            self.rotary_base = 10000
+        if self.attention_dropout is None:
+            self.attention_dropout = 0.
+        if self.untie_embeddings_and_output_weights is None:
+            self.untie_embeddings_and_output_weights = True
+        if self.swiglu is None:
+            self.swiglu = True
+        if self.add_qkv_bias is None:
+            self.add_qkv_bias = True
+        if self.disable_bias_linear is None:
+            self.disable_bias_linear = True
+        if self.moe_router_topk is None:
+            self.moe_router_topk = 2
+        if self.moe_router_pre_softmax is None:
+            self.moe_router_pre_softmax = False
+        if self.moe_aux_loss_coeff is None:
+            self.moe_aux_loss_coeff = 0.
+        if self.qk_layernorm is None:
+            self.qk_layernorm = False
+
+    def _init_mixed_precision(self):
+        from swift.llm.argument.base_args.model_args import ModelArguments
+        ModelArguments._init_mixed_precision(self)
+        if self.apply_query_key_layer_scaling is None:
+            self.apply_query_key_layer_scaling = self.fp16
+        if self.apply_query_key_layer_scaling:
+            os.environ['NVTE_APPLY_QK_LAYER_SCALING'] = '1'
+
+    def _init_moe(self):
+        if self.moe_shared_expert_intermediate_size == 0:
+            self.moe_shared_expert_intermediate_size = None
+        if self.moe_ffn_hidden_size is None:
+            self.moe_ffn_hidden_size = self.ffn_hidden_size
+        else:
+            self.ffn_hidden_size = self.moe_ffn_hidden_size
+
+    def __post_init__(self):
+        from swift.llm.argument.base_args.model_args import ModelArguments
+        if self.use_flash_attn or self.attention_backend == 'flash':
+            require_version('flash-attn')
+        os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1'
+        self._set_default()
+        self.group_query_attention = self.num_query_groups > 1
+        if self.rope_scaling is not None:
+            self.rope_scaling = ModelArguments.parse_to_dict(self.rope_scaling)
+        if self.eval_interval is None:
+            self.eval_interval = self.save_interval
+        if self.seq_length is None:
+            self.seq_length = self.max_position_embeddings
+        if self.tensorboard_dir is None and self.save is not None:
+            self.tensorboard_dir = f'{self.save}/runs'
+        self._init_moe()
+        self._init_mixed_precision()
+
+        self.tensorboard_dir = to_abspath(self.tensorboard_dir)
+
+    def _args_to_argv(self) -> Tuple[List[Any], Dict[str, Any]]:
+        new_args = []
+        args_dict = asdict(self)
+        extra_args = {}
+        for k, value in args_dict.items():
+            if k not in MegatronArguments.__annotations__:
+                extra_args[k] = value
+                continue
+            if value is None or value is False:
+                continue
+            new_args.append(f"--{k.replace('_', '-')}")
+            if isinstance(value, list):
+                new_args += [str(v) for v in value]
+            elif value is not True:
+                new_args.append(str(value))
+
+        return new_args, extra_args
+
+    def parse_to_megatron(self):
+        new_args, extra_args = self._args_to_argv()
+        sys._old_argv = sys.argv
+        sys.argv = sys.argv[:1] + new_args
+        # parameter conflict
+        extra_args.pop('loss_scale', None)
+        return extra_args

ms-swift/swift/megatron/model/gpt/mcore2hf.py

@@ -0,0 +1,70 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from megatron.training import get_args
+
+
+def set_attn_state(args, mg_attn, hf_attn):
+    num_query_groups = (args.num_query_groups if args.group_query_attention else args.num_attention_heads)
+    # Copy weights
+    mg_attn_weight = mg_attn.linear_qkv.weight.reshape((num_query_groups, -1, args.hidden_size))
+    q_dim, kv_dim = hf_attn.q_proj.weight.shape[0] // num_query_groups, hf_attn.k_proj.weight.shape[
+        0] // num_query_groups
+    hf_attn.q_proj.weight.data.copy_(mg_attn_weight[:, :q_dim, :].reshape(-1, args.hidden_size))
+    hf_attn.k_proj.weight.data.copy_(mg_attn_weight[:, q_dim:-kv_dim, :].reshape(-1, args.hidden_size))
+    hf_attn.v_proj.weight.data.copy_(mg_attn_weight[:, -kv_dim:, :].reshape(-1, args.hidden_size))
+    hf_attn.o_proj.weight.data.copy_(mg_attn.linear_proj.weight)
+
+    # Copy bias
+    if args.add_qkv_bias:
+        mg_attn_bias = mg_attn.linear_qkv.bias.reshape((num_query_groups, -1))
+        hf_attn.q_proj.bias.data.copy_(mg_attn_bias[:, :q_dim].reshape(-1))
+        hf_attn.k_proj.bias.data.copy_(mg_attn_bias[:, q_dim:-kv_dim].reshape(-1))
+        hf_attn.v_proj.bias.data.copy_(mg_attn_bias[:, -kv_dim:].reshape(-1))
+
+    if args.qk_layernorm:
+        hf_attn.q_norm.weight.data.copy_(mg_attn.q_layernorm.weight)
+        hf_attn.k_norm.weight.data.copy_(mg_attn.k_layernorm.weight)
+
+
+def _set_mlp_state(mg_mlp, hf_mlp):
+    ffn_hidden_size = hf_mlp.gate_proj.weight.shape[0]
+    hf_mlp.gate_proj.weight.data.copy_(mg_mlp.linear_fc1.weight[:ffn_hidden_size])
+    hf_mlp.up_proj.weight.data.copy_(mg_mlp.linear_fc1.weight[ffn_hidden_size:])
+    hf_mlp.down_proj.weight.data.copy_(mg_mlp.linear_fc2.weight)
+
+
+def set_mlp_state(args, mg_mlp, hf_mlp):
+    if args.num_experts:
+        hf_mlp.gate.weight.data.copy_(mg_mlp.router.weight)
+        if mg_mlp.shared_experts is not None:
+            hf_mlp.shared_expert_gate.weight.data.copy_(mg_mlp.shared_experts.gate_weight)
+        for expert_idx in range(args.num_experts):
+            _set_mlp_state(mg_mlp.experts.local_experts[expert_idx], hf_mlp.experts[expert_idx])
+
+        if mg_mlp.shared_experts is not None:
+            _set_mlp_state(mg_mlp.shared_experts, hf_mlp.shared_expert)
+    else:
+        _set_mlp_state(mg_mlp, hf_mlp)
+
+
+def set_layer_state(args, mg_model, hf_model, layer_idx):
+    mg_layer = mg_model.decoder.layers[layer_idx]
+    hf_layer = hf_model.model.layers[layer_idx]
+    set_attn_state(args, mg_layer.self_attention, hf_layer.self_attn)
+    set_mlp_state(args, mg_layer.mlp, hf_layer.mlp)
+
+    post_attention_layernorm_weight = hf_layer.post_attention_layernorm.weight
+    if args.num_experts:
+        post_attention_layernorm_weight.data.copy_(mg_layer.pre_mlp_layernorm.weight)
+    else:
+        post_attention_layernorm_weight.data.copy_(mg_layer.mlp.linear_fc1.layer_norm_weight)
+    hf_layer.input_layernorm.weight.data.copy_(mg_layer.self_attention.linear_qkv.layer_norm_weight)
+
+
+def convert_mcore2hf(hf_model, mg_model):
+    args = get_args()
+    hf_model.model.embed_tokens.weight.data.copy_(mg_model.embedding.word_embeddings.weight)
+    if args.untie_embeddings_and_output_weights:
+        hf_model.lm_head.weight.data.copy_(mg_model.output_layer.weight)
+    hf_model.model.norm.weight.data.copy_(mg_model.decoder.final_layernorm.weight)
+    for layer_idx in range(args.num_layers):
+        set_layer_state(args, mg_model, hf_model, layer_idx)

ms-swift/swift/megatron/train/__init__.py

@@ -0,0 +1,2 @@
+from .pt import megatron_pt_main
+from .sft import megatron_sft_main

ms-swift/swift/megatron/train/pt.py

@@ -0,0 +1,19 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import List, Union
+
+from ..argument import MegatronTrainArguments
+from .sft import MegatronSft
+
+
+class MegatronPt(MegatronSft):
+    args_class = MegatronTrainArguments
+    args: args_class
+
+    def _prepare_template(self) -> None:
+        self.args.use_chat_template = False
+        super()._prepare_template()
+        self.template.loss_scale = 'all'
+
+
+def megatron_pt_main(args: Union[List[str], MegatronTrainArguments, None] = None):
+    return MegatronPt(args).main()

ms-swift/swift/megatron/train/sft.py

@@ -0,0 +1,65 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+from typing import List, Union
+
+from megatron.core.enums import ModelType
+from megatron.training import pretrain
+
+from swift.llm.train import SwiftSft
+from swift.utils import get_logger, is_master, plot_images
+from ..argument import MegatronTrainArguments
+from ..utils import patch_megatron_tokenizer
+from .patcher import patch_megatron_data_collator, patch_training_log
+from .utils import build_streaming_dataloader, forward_step, get_swift_datasets_provider
+
+logger = get_logger()
+
+
+class MegatronSft(SwiftSft):
+    args_class = MegatronTrainArguments
+    args: args_class
+
+    def __init__(self, args: Union[List[str], MegatronTrainArguments, None] = None) -> None:
+        self.train_msg = {}
+        super(SwiftSft, self).__init__(args)
+        args = self.args
+        _, self.processor = args.get_model_processor(load_model=False)
+        patch_megatron_tokenizer(self.processor)
+        args.init_model_args(self.processor.model_info.config)
+        self._prepare_template()
+        self.template.use_megatron = True
+        args.save_args(args.save)
+
+    def run(self):
+        args = self.args
+
+        train_dataset, val_dataset = self._get_dataset()
+        train_dataset, val_dataset = self._encode_dataset(train_dataset, val_dataset)
+        data_collator = self.template.data_collator
+        if args.streaming:
+            train_dataset = build_streaming_dataloader(args, train_dataset, data_collator)
+            if val_dataset is not None:
+                val_dataset = build_streaming_dataloader(args, val_dataset, data_collator)
+        datasets_provider = get_swift_datasets_provider(train_dataset, val_dataset)
+        datasets_provider.is_distributed = True
+
+        logging_path = os.path.join(args.save, 'logging.jsonl')
+        logger.info(f'The logging file will be saved in: {logging_path}')
+        try:
+            with patch_training_log(), patch_megatron_data_collator(data_collator):
+                pretrain(
+                    datasets_provider,
+                    args.megatron_model_meta.model_provider,
+                    ModelType.encoder_or_decoder,
+                    forward_step,
+                    args_defaults=args.extra_args)
+        finally:
+            # Visualization
+            if is_master():
+                images_dir = os.path.join(args.save, 'images')
+                logger.info(f'images_dir: {images_dir}')
+                plot_images(images_dir, args.tensorboard_dir)
+
+
+def megatron_sft_main(args: Union[List[str], MegatronTrainArguments, None] = None):
+    return MegatronSft(args).main()

ms-swift/swift/plugin/agent_template/hermes.py

@@ -0,0 +1,78 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import re
+from typing import TYPE_CHECKING, List, Tuple, Union
+
+import json
+
+from .base import BaseAgentTemplate
+
+if TYPE_CHECKING:
+    from swift.llm.infer import Function
+    from swift.llm.template import Prompt
+
+
+class HermesAgentTemplate(BaseAgentTemplate):
+
+    def get_toolcall(self, response: str) -> List['Function']:
+        from swift.llm.infer import Function
+        res_list = re.findall(r'<tool_call>(.+?)</tool_call>', response, re.DOTALL)
+        functions = []
+        for res in res_list:
+            res = self._parse_json(res)
+            if isinstance(res, dict) and 'name' in res and 'arguments' in res:
+                functions.append(Function(name=res['name'], arguments=res['arguments']))
+        if len(functions) == 0:
+            # compat react_en
+            return super().get_toolcall(response)
+        return functions
+
+    def _format_tool_responses(
+        self,
+        assistant_content: str,
+        tool_messages,
+    ) -> Tuple[str, 'Prompt']:
+        with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
+        if with_action:
+            return super()._format_tool_responses(assistant_content, tool_messages)
+        if hasattr(self, 'template_meta'):
+            prompt = self.template_meta.prompt
+            chat_sep = self.template_meta.chat_sep
+        else:
+            prompt = ['<|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n']
+            chat_sep = ['<|im_end|>\n']
+        res = chat_sep.copy()
+        res_tool = []
+        for tool_message in tool_messages:
+            tool_content = tool_message['content']
+            res_tool.append(f'<tool_response>\n{tool_content}\n</tool_response>')
+        total_tool = '\n'.join(res_tool)
+        for context in prompt:
+            if isinstance(context, str):
+                context = context.replace('{{QUERY}}', total_tool)
+            res.append(context)
+        return assistant_content, res
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: str, user_message=None) -> str:
+        tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
+        return f"""{system}
+
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+""" + '\n'.join(tool_descs) + """
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{"name": <function-name>, "arguments": <args-json-object>}
+</tool_call>"""
+
+    def _format_tool_calls(self, tool_call_messages):
+        tool_calls = []
+        for message in tool_call_messages:
+            tool_call = self._parse_tool_call(message['content'])
+            tool_calls.append(f'<tool_call>\n{json.dumps(tool_call, ensure_ascii=False)}\n</tool_call>')
+        return '\n'.join(tool_calls)

ms-swift/swift/plugin/agent_template/react.py

@@ -0,0 +1,66 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import List, Union
+
+from .base import BaseAgentTemplate
+
+
+class ReactEnAgentTemplate(BaseAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: str, user_message=None) -> str:
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'en')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(
+                f'{tool_desc.name_for_model}: Call this tool to interact with the {tool_desc.name_for_human} API. '
+                f'What is the {tool_desc.name_for_human} API useful for? {tool_desc.description_for_model} '
+                f'Parameters: {tool_desc.parameters} {tool_desc.args_format}')
+
+        return """Answer the following questions as best you can. You have access to the following tools:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+Use the following format:
+
+Question: the input question you must answer
+Thought: you should always think about what to do
+Action: the action to take, should be one of [{','.join(tool_names)}]
+Action Input: the input to the action
+Observation: the result of the action
+... (this Thought/Action/Action Input/Observation can be repeated zero or more times)
+Thought: I now know the final answer
+Final Answer: the final answer to the original input question
+
+Begin!
+"""
+
+
+class ReactZnAgentTemplate(BaseAgentTemplate):
+
+    def _format_tools(self, tools: List[Union[str, dict]], system: str, user_message=None) -> str:
+        tool_names = []
+        tool_descs = []
+        for tool in tools:
+            tool_desc = self._parse_tool(tool, 'zh')
+            tool_names.append(tool_desc.name_for_model)
+            tool_descs.append(f'{tool_desc.name_for_model}: 调用此工具与 {tool_desc.name_for_human} API 进行交互。'
+                              f'{tool_desc.name_for_human} 有什么用？{tool_desc.description_for_model} '
+                              f'输入参数：{tool_desc.parameters} {tool_desc.args_format}')
+        return """尽可能地回答以下问题。你可以使用以下工具:
+
+""" + '\n\n'.join(tool_descs) + f"""
+
+请按照以下格式进行:
+
+Question: 需要你回答的输入问题
+Thought: 你应该总是思考该做什么
+Action: 需要使用的工具，应该是[{','.join(tool_names)}]中的一个
+Action Input: 传入工具的内容
+Observation: 行动的结果
+... (这个Thought/Action/Action Input/Observation可以重复N次)
+Thought: 我现在知道最后的答案
+Final Answer: 对原始输入问题的最终答案
+
+现在开始！
+"""

ms-swift/swift/plugin/loss_scale/__init__.py

@@ -0,0 +1 @@
+from .loss_scale import loss_scale_map

ms-swift/swift/plugin/loss_scale/config/agentflan.json

@@ -0,0 +1,22 @@
+{
+    "response":{
+        "Name:": [1.0, 3.0],
+        "Action:": [1.0, 3.0],
+        "ACTION:": [1.0,3.0],
+        "Tool:": [1.0, 3.0],
+        "Command": [1.0, 3.0],
+        "Arguments:": [1.0, 3.0],
+        "action input": [1.0, 3.0],
+        "ACTION_INPUT:":[1.0, 3.0],
+        "Action Input:": [1.0, 3.0],
+        "Thought:": [1.0, 1.0],
+        "Final Answer:": [1.0, 1.0],
+        "Observation:": [2.0, 0.0]
+    },
+    "query":{
+        "What is the tool you want to use": [3.0],
+        "What are the required parameter names": [3.0],
+        "What is the value of": [3.0],
+        "What are the required parameter names for this tool": [3.0]
+    }
+}

ms-swift/swift/plugin/loss_scale/config/hermes.json

@@ -0,0 +1,3 @@
+{
+    "<tool_call>.+?</tool_call>": [2.0]
+}

ms-swift/swift/plugin/prm.py

@@ -0,0 +1,154 @@
+import os
+from typing import Any, Dict, List, Union
+
+import json
+
+from swift.llm import InferRequest
+
+
+class PRM:
+
+    def __call__(self, **kwargs) -> List[Any]:
+        raise NotImplementedError
+
+
+SYSTEM = """
+You are a process reward model, give the reward value of the answer, you must follow the instructions below:
+
+1. Output a float reward value between -1.0 and 1.0, -1.0 means the worst answer, 1.0 means the best answer, please think step by step to give your reasons and thoughts, but the reward must appare at the end with this format: **Reward: your-reward-value**.
+
+2. The answer may be incomplete, you must give the reward by the existing part of the answer, taking into account semantic coherence, logical correctness, and clarity.
+
+3. A ground truth answer will be given to you, it may be not the best one, consider it as a reference example.
+
+Begin!
+""" # noqa
+
+QUERY = """
+The original question or the previous conversation:
+
+#query#
+
+Here is the ground truth as the reference:
+
+#ground_truth#
+
+Given the upper information, give your reward(-1.0~1.0) of the following answer:
+
+#response#
+"""
+
+
+class QwenMaxPRM(PRM):
+
+    def __call__(self, infer_requests: List[Union[InferRequest, Dict]], ground_truths: List[str],
+                 **kwargs) -> List[float]:
+        # TODO: check request_config
+        rewards = []
+
+        from openai import OpenAI
+
+        client = OpenAI(
+            api_key=os.getenv('DASHSCOPE_API_KEY'),
+            base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
+        )
+
+        for request, ground_truth in zip(infer_requests, ground_truths):
+            previous = request['messages'][:-1]
+            if previous[0]['role'] == 'system':
+                previous = previous[1:]
+
+            assert request['messages'][-1]['role'] == 'assistant'
+            query = QUERY.replace('#query#', json.dumps(previous))
+            query = query.replace('#ground_truth#', ground_truth)
+            query = query.replace('#response#', request['messages'][-1]['content'])
+            messages = [
+                {
+                    'role': 'system',
+                    'content': SYSTEM
+                },
+                {
+                    'role': 'user',
+                    'content': query
+                },
+            ]
+            completion = client.chat.completions.create(
+                model='qwen-max',
+                messages=messages,
+            )
+
+            content = completion.choices[0].message.content
+            if 'Reward:' not in content:
+                rewards.append(0.)
+            else:
+                try:
+                    reward = float(content.split('Reward:')[1].strip().replace('*', ''))
+                    rewards.append(reward)
+                except Exception:
+                    rewards.append(0.)
+
+        return rewards
+
+
+class ClientPRM(PRM):
+
+    def __init__(self, api_key=None, base_url=None, model=None):
+        from swift.llm import InferClient
+        import os
+        if api_key is None:
+            api_key = os.getenv('DASHSCOPE_API_KEY')
+        if base_url is None:
+            base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
+        if model is None:
+            model = 'qwen-plus'
+        self.infer_engine = InferClient(base_url=base_url, api_key=api_key)
+        self.infer_engine.strict = False
+        self.infer_kwargs = {
+            'model': model,
+        }
+
+    def __call__(self, infer_requests: List[Union[InferRequest, Dict]], ground_truths: List[str],
+                 **kwargs) -> List[float]:
+        prm_infer_requests = []
+        request_config = kwargs.get('request_config')
+        for request, ground_truth in zip(infer_requests, ground_truths):
+            previous = request['messages'][:-1]
+            if previous[0]['role'] == 'system':
+                previous = previous[1:]
+
+            assert request['messages'][-1]['role'] == 'assistant'
+            query = QUERY.replace('#query#', json.dumps(previous))
+            query = query.replace('#ground_truth#', ground_truth)
+            query = query.replace('#response#', request['messages'][-1]['content'])
+            messages = [
+                {
+                    'role': 'system',
+                    'content': SYSTEM
+                },
+                {
+                    'role': 'user',
+                    'content': query
+                },
+            ]
+
+            prm_infer_requests.append(InferRequest(messages=messages))
+
+        responses = self.infer_engine.infer(prm_infer_requests, request_config=request_config, **self.infer_kwargs)
+        rewards = []
+        for response in responses:
+            content = response.choices[0].message.content
+            if 'Reward:' not in content:
+                rewards.append(0.)
+            else:
+                try:
+                    reward = float(content.split('Reward:')[1].strip().replace('*', ''))
+                    rewards.append(reward)
+                except Exception:
+                    rewards.append(0.)
+        return rewards
+
+
+prms = {
+    'qwen_max': QwenMaxPRM,
+    'client': ClientPRM,
+}

ms-swift/swift/plugin/rm_plugin.py

@@ -0,0 +1,229 @@
+import re
+import textwrap
+from copy import deepcopy
+from typing import Dict, List
+
+import torch
+
+from swift.llm import PtEngine, RequestConfig, Template, to_device
+from swift.llm.infer.protocol import ChatCompletionResponse
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+class DefaultRMPlugin:
+    """
+    Default Reward Model Plugin
+
+    This class implements the default processing logic for reward models.
+    It assumes that `self.model` is a classification model with a value head(output dimmension 1).
+    The first logits value from the model's output is used as the reward score.
+    """
+
+    def __init__(self, model, template):
+        self.model = model
+        self.template: Template = template
+
+    def __call__(self, inputs):
+        batched_inputs = [self.template.encode(deepcopy(infer_request)) for infer_request in inputs]
+        reward_inputs = to_device(self.template.data_collator(batched_inputs), self.model.device)
+        reward_inputs.pop('labels')
+
+        with torch.inference_mode():
+            return self.model(**reward_inputs).logits[:, 0]
+
+
+class GenRMPlugin(DefaultRMPlugin):
+
+    def __init__(self, model, template):
+        """
+        Generative Reward Model Plugin Example.
+
+        This method sets up the reward model plugin by initializing the PtEngine for efficient inference,
+        configuring the request parameters, and defining the system prompt that guides the reward model in
+        evaluating responses.
+
+        Args:
+            model (torch.nn.Module): The generative reward model.
+            template (Template): The template used for encoding input data.
+    """
+
+        super().__init__(model, template)
+        # initilize PTEngine to infer
+        self.engine = PtEngine.from_model_template(self.model, self.template, max_batch_size=0)  # 0: no limit
+        self.request_config = RequestConfig()  # customise your request config here
+        self.system = textwrap.dedent("""
+            Based on the dialogue history, analyze in detail whether the model's response is accurate, complete, and relevant.
+            Assign a reward score between 0 and 1, where 0 indicates completely incorrect and 1 indicates fully correct.
+            Before finishing your response, please assign a reward using the following format:
+
+            Reward: {reward}
+
+            For example:
+            Reward: 0.85
+        """)  # noqa
+
+    def __call__(self, inputs):
+        """
+        Compute reward scores for the provided inputs.
+
+        This method processes each input by converting dialogue messages into a query, sending the query to the
+        reward model for inference, and extracting the reward scores from the model's responses. The final reward
+        for each input is the average of all extracted scores.
+        Args:
+            inputs (List[Dict]): A list of input requests. Each input request is a dictionary containing:
+                - 'messages' (List[Dict]): messages from the training model. Each message dictionary includes:
+                    - 'role' (str): The role of the speaker (e.g., 'user', 'assistant').
+                    - 'content' (str): The content of the message.
+                - Additional dataset columns as key-value pairs (e.g., 'solutions', 'images').
+        Returns:
+            torch.Tensor: A tensor containing the average reward scores for each input. The tensor has a shape of (N,),
+            where N is the number of input requests.
+        """
+
+        rm_inputs = self.prepare_rm_inputs(inputs)
+        results = self.engine.infer(rm_inputs, self.request_config, use_tqdm=False)
+        rewards = self.compute_rewards(results)
+        return torch.tensor(rewards, dtype=torch.float32)
+
+    def prepare_rm_inputs(self, inputs: List[Dict]) -> List[Dict]:
+        """
+        Prepare inputs for the reward model by converting messages into queries.
+
+        Args:
+            inputs (List[Dict]): A list of input requests.
+
+        Returns:
+            List[Dict]: Processed inputs for the reward model.
+        """
+        rm_inputs = []
+        for idx, infer_request in enumerate(inputs):
+            # Deep copy to prevent modification of original input
+            rm_infer_request = deepcopy(infer_request)
+
+            # Extract and convert messages to a single query string
+            messages = rm_infer_request.get('messages')
+            query = self.messages_to_query(messages)
+
+            # Construct new messages tailored for the reward model
+            rm_messages = [{'role': 'system', 'content': self.system}, {'role': 'user', 'content': query}]
+
+            # Update the messages in the reward infer request
+            rm_infer_request['messages'] = rm_messages
+            rm_inputs.append(rm_infer_request)
+        return rm_inputs
+
+    @staticmethod
+    def extract_reward(model_output: str) -> float:
+        """
+        Extract the reward score from the model's output.
+
+        Args:
+            model_output (str): The model's output string, expected to follow the format "Reward: {reward}".
+
+        Returns:
+            float: The extracted reward score.
+
+        Raises:
+            ValueError: If the reward score cannot be extracted or the format is incorrect.
+        """
+        match = re.search(r'Reward:\s*([0-1](?:\.\d+)?)', model_output)
+        if match:
+            return float(match.group(1))
+        else:
+            logger.warning("Unable to extract reward score from the model's output, set reward to 0")
+            return None
+
+    @staticmethod
+    def messages_to_query(messages):
+        """
+        Compress a list of message dictionaries into a single query string.
+
+        Args:
+            messages (list[dict]): A list of message dictionaries, each containing:
+                - 'role' (str): The role of the speaker (e.g., 'user', 'assistant').
+                - 'content' (str): The content of the message.
+
+        Returns:
+            str: A single string that concatenates all messages in a formatted manner.
+
+        Example:
+            >>> messages = [
+            ...     {'role': 'user', 'content': 'Hello, how are you?'},
+            ...     {'role': 'assistant', 'content': 'I am fine, thank you! How can I assist you today?'},
+            ...     {'role': 'user', 'content': 'Can you help me with my homework?'}
+            ... ]
+            >>> print(messages_to_query(messages))
+            User: Hello, how are you?
+            Assistant: I am fine, thank you! How can I assist you today?
+            User: Can you help me with my homework?
+        """
+        # Initialize an empty list to hold formatted messages
+        formatted_messages = []
+
+        # Define a mapping for role capitalization if needed
+        role_mapping = {
+            'user': 'User',
+            'assistant': 'Assistant',
+            'system': 'System'
+            # Add more roles here as needed
+        }
+
+        for idx, message in enumerate(messages):
+            if not isinstance(message, dict):
+                raise TypeError(f'Each message must be a dictionary. Found {type(message)} at index {idx}.')
+
+            # Extract 'role' and 'content' from each message
+            role = message.get('role')
+            content = message.get('content')
+            if not content:
+                continue
+
+            # Capitalize the role using the mapping, default to capitalized original role
+            role_formatted = role_mapping.get(role.lower(), role.capitalize())
+
+            # Append the formatted message to the list
+            formatted_messages.append(f'{role_formatted}: {content}')
+
+        # Join all formatted messages with newline characters
+        query = '\n'.join(formatted_messages)
+
+        return query
+
+    def compute_rewards(self, results: List[ChatCompletionResponse]) -> List[float]:
+        """
+        Compute average reward scores from the reward model's outputs.
+
+        Args:
+            results (List[ChatCompletionResponse]): A list of results from the reward model.
+
+        Returns:
+            List[float]: A list of average reward scores.
+        """
+        rewards = []
+        for idx, output in enumerate(results):
+            try:
+                cur_rewards = []
+                for choice in output.choices:
+                    response = choice.message.content
+                    reward = self.extract_reward(response)
+                    cur_rewards.append(reward)
+                cur_rewards = [r for r in cur_rewards if r is not None]
+                if cur_rewards:
+                    average_reward = sum(cur_rewards) / len(cur_rewards)
+                else:
+                    average_reward = 0.0
+                    logger.warning('No valid rewards extracted. Assigning reward score of 0.0.')
+
+                rewards.append(average_reward)
+            except Exception as e:
+                logger.error(f'Error computing reward: {e}')
+                rewards.append(0.0)  # Assign default reward score on failure
+        return rewards
+
+
+rm_plugins = {
+    'default': DefaultRMPlugin,
+    'genrm': GenRMPlugin,
+}

ms-swift/swift/trainers/__init__.py

@@ -0,0 +1,49 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_utils import (EvaluationStrategy, FSDPOption, HPSearchBackend, HubStrategy, IntervalStrategy,
+                                        SchedulerType)
+
+from swift.utils.import_utils import _LazyModule
+from . import callback
+
+try:
+    # https://github.com/huggingface/transformers/pull/25702
+    from transformers.trainer_utils import ShardedDDPOption
+except ImportError:
+    ShardedDDPOption = None
+
+if TYPE_CHECKING:
+    from .arguments import Seq2SeqTrainingArguments, TrainingArguments
+    from .rlhf_trainer import (CPOTrainer, DPOTrainer, KTOTrainer, ORPOTrainer, RLHFTrainerMixin, PPOTrainer,
+                               RewardTrainer, GRPOTrainer)
+    from .rlhf_arguments import DPOConfig, CPOConfig, KTOConfig, ORPOConfig, PPOConfig, RewardConfig
+    from .trainer_factory import TrainerFactory
+    from .trainers import Seq2SeqTrainer, Trainer, EmbeddingTrainer
+    from .mixin import SwiftMixin
+
+else:
+    _extra_objects = {k: v for k, v in globals().items() if not k.startswith('_')}
+    _import_structure = {
+        'arguments': ['Seq2SeqTrainingArguments', 'TrainingArguments'],
+        'rlhf_arguments':
+        ['DPOConfig', 'CPOConfig', 'KTOConfig', 'ORPOConfig', 'PPOConfig', 'RewardConfig', 'GRPOConfig'],
+        'rlhf_trainer': [
+            'CPOTrainer', 'DPOTrainer', 'KTOTrainer', 'ORPOTrainer', 'RLHFTrainerMixin', 'PPOTrainer', 'RewardTrainer',
+            'GRPOTrainer'
+        ],
+        'trainer_factory': ['TrainerFactory'],
+        'trainers': ['Seq2SeqTrainer', 'Trainer', 'EmbeddingTrainer'],
+        'mixin': ['SwiftMixin'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects=_extra_objects,
+    )

ms-swift/swift/trainers/callback.py

@@ -0,0 +1,124 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+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, speed_metrics
+
+from swift.utils import append_to_jsonl, is_pai_training_job, use_torchacc
+from ..utils.utils import format_time
+from .arguments import TrainingArguments
+
+
+def add_train_message(logs, state, start_time) -> None:
+    logs['global_step/max_steps'] = f'{state.global_step}/{state.max_steps}'
+    train_percentage = state.global_step / state.max_steps if state.max_steps else 0.
+    logs['percentage'] = f'{train_percentage * 100:.2f}%'
+    elapsed = time.time() - start_time
+    logs['elapsed_time'] = format_time(elapsed)
+    if train_percentage != 0:
+        logs['remaining_time'] = format_time(elapsed / train_percentage - elapsed)
+    for k, v in logs.items():
+        if isinstance(v, float):
+            logs[k] = round(logs[k], 8)
+
+
+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.current_step = 0
+        self.start_time = time.time()
+        if use_torchacc():
+            self.warmup_start_time = 0
+            self.warmup_metric = None
+            self.metric_warmup_step = int(args.metric_warmup_step
+                                          * state.max_steps) if args.metric_warmup_step < 1 else args.metric_warmup_step
+
+    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):
+
+        if use_torchacc():
+            if state.global_step >= self.metric_warmup_step and self.warmup_start_time == 0:
+                self.warmup_start_time = time.time()
+                self.metric_warmup_step = state.global_step
+            if state.max_steps == state.global_step and self.warmup_metric is None:
+                num_steps = state.max_steps - self.metric_warmup_step
+                num_total_samples = args.train_dataset_sample
+                num_after_warmup_samples = int(num_total_samples / state.max_steps * num_steps)
+                self.warmup_metric = speed_metrics('warmup_train', self.warmup_start_time, num_after_warmup_samples,
+                                                   num_steps)
+                self.warmup_metric['num_total_samples'] = num_total_samples
+                self.warmup_metric['num_after_warmup_samples'] = num_after_warmup_samples
+            if 'train_samples_per_second' in logs:
+                logs.update(self.warmup_metric)
+                state.log_history[-1] = logs
+
+        add_train_message(logs, state, self.start_time)
+        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):
+            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()
+        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)
+        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

ms-swift/swift/trainers/mixin.py

@@ -0,0 +1,516 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Part of the implementation is borrowed from huggingface/transformers.
+import inspect
+import os
+import shutil
+import time
+from contextlib import contextmanager
+from copy import copy
+from functools import partial
+from types import MethodType
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import safetensors
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import transformers
+from datasets import Dataset as HfDataset
+from modelscope import check_local_model_is_latest
+from packaging import version
+from peft import PeftModel
+from torch.nn import Module
+from torch.utils.data import DataLoader
+from transformers import PreTrainedModel
+from transformers.data.data_collator import DataCollator
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.modeling_utils import unwrap_model
+from transformers.trainer import TrainerCallback
+from transformers.trainer_utils import EvalPrediction, IntervalStrategy
+from transformers.utils import is_torch_npu_available
+
+from swift.hub import get_hub
+from swift.llm import BatchSamplerShard, DataLoaderDispatcher, DataLoaderShard, Template
+from swift.plugin import MeanMetric, compute_acc, extra_tuners
+from swift.tuners import SwiftModel
+from swift.utils import get_logger, is_mp_ddp, use_torchacc
+from swift.utils.torchacc_utils import ta_trim_graph
+from ..utils.torch_utils import get_device_count
+from .arguments import TrainingArguments
+from .utils import can_return_loss, find_labels, get_function, is_instance_of_ms_model
+
+try:
+    from trl import AutoModelForCausalLMWithValueHead
+except (ImportError, RuntimeError):
+    AutoModelForCausalLMWithValueHead = None
+
+logger = get_logger()
+
+
+class SwiftMixin:
+
+    def __init__(self,
+                 model: Union[PreTrainedModel, Module] = None,
+                 args: TrainingArguments = None,
+                 data_collator: Optional[DataCollator] = None,
+                 train_dataset: Optional[HfDataset] = None,
+                 eval_dataset: Optional[Union[HfDataset, Dict[str, HfDataset]]] = None,
+                 template: Optional[Template] = None,
+                 model_init: Optional[Callable[[], PreTrainedModel]] = None,
+                 compute_loss_func: Optional[Callable] = None,
+                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
+                 callbacks: Optional[List[TrainerCallback]] = None,
+                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+                 preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
+                 **kwargs) -> None:
+        if not hasattr(train_dataset, '__len__') and args.dataloader_num_workers > 1:
+            args.dataloader_num_workers = 1
+            logger.warning('Using IterableDataset, setting args.dataloader_num_workers to 1.')
+
+        if args.check_model and hasattr(model, 'model_dir'):
+            from swift.utils.logger import ms_logger_ignore_error
+            with ms_logger_ignore_error():
+                check_local_model_is_latest(
+                    model.model_dir, user_agent={
+                        'invoked_by': 'local_trainer',
+                        'third_party': 'swift',
+                    })
+        if eval_dataset is None and args:
+            args.evaluation_strategy = IntervalStrategy.NO
+            args.eval_strategy = IntervalStrategy.NO
+
+        self._custom_metrics = {}
+        self.template = template
+        self.max_memory = 0
+        self.hub = get_hub()
+
+        self.model_meta = model.model_meta
+        with self.hub.patch_hub():
+            super().__init__(
+                model=model,
+                args=args,
+                data_collator=data_collator,
+                train_dataset=train_dataset,
+                eval_dataset=eval_dataset,
+                tokenizer=template.tokenizer,
+                model_init=model_init,
+                compute_metrics=compute_metrics,
+                callbacks=callbacks,
+                optimizers=optimizers,
+                preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+                **kwargs)
+
+        self.compute_loss_func = compute_loss_func
+        if get_function(model.__class__.forward) is not get_function(model.forward):
+            self.label_names = find_labels(model)
+            self.can_return_loss = can_return_loss(model)
+        self.label_names = self.label_names or ['labels']
+        self.start_time = time.time()
+        if self.template.sequence_parallel_size > 1:
+            from swift.trainers.sequence_parallel import sequence_parallel
+            sequence_parallel.prepare_trainer(self)
+
+    def _save_initial_model(self, output_dir):
+        # pissa/olora/lora-ga
+        model = unwrap_model(self.model)
+        if isinstance(model, PeftModel):
+            config = model.peft_config.get('default')
+            init_lora_weights = getattr(config, 'init_lora_weights', None)
+            if (isinstance(init_lora_weights, str)
+                    and any(s in init_lora_weights for s in ('pissa', 'olora', 'lora-ga'))):
+                config.init_lora_weights = True
+                model.save_pretrained(os.path.join(output_dir, 'initial_model'))
+                config.init_lora_weights = init_lora_weights
+
+    def _save_converted_model(self, output_dir):
+        # pissa/olora/lora-ga
+        model = unwrap_model(self.model)
+        if isinstance(model, PeftModel):
+            config = model.peft_config.get('default')
+            init_lora_weights = getattr(config, 'init_lora_weights', None)
+            if isinstance(init_lora_weights, str):
+                config = copy(config)
+                os.makedirs(os.path.join(output_dir, 'converted'), exist_ok=True)
+                if 'lora-ga' in init_lora_weights:
+                    try:
+                        from lora_ga.entrypoint import LoraGAContext
+                        with LoraGAContext(model):
+                            model.save_pretrained(
+                                os.path.join(output_dir, 'converted', 'default'),
+                                path_initial_model_for_weight_conversion=os.path.join(
+                                    os.path.dirname(output_dir), 'initial_model'),
+                            )
+                            model.peft_config['default'] = config
+                    except ImportError as e:
+                        error_message = """
+                        Since 'LoRA-GA' is not implemented by PEFT, you will need to install it directly from GitHub.
+                        Command: 'pip install git+https://github.com/lxline/LoRA-GA.git'.
+                        """
+                        logger.info(error_message)
+                        raise RuntimeError(error_message) from e
+                elif 'pissa' in init_lora_weights or 'olora' in init_lora_weights:
+                    model.save_pretrained(
+                        os.path.join(output_dir, 'converted', 'default'),
+                        path_initial_model_for_weight_conversion=os.path.join(
+                            os.path.dirname(output_dir), 'initial_model'),
+                    )
+                    model.peft_config['default'] = config
+
+    def _load_optimizer_and_scheduler(self, *args, **kwargs):
+        super()._load_optimizer_and_scheduler(*args, **kwargs)
+        if is_mp_ddp():
+            # fix mp+ddp adamw
+            for v in self.optimizer.state.values():
+                if 'step' in v:
+                    # not on the same device
+                    device_set = set([t.device for t in v.values()]) - {v['step'].device, torch.device('cpu')}
+                    if len(device_set) >= 1:
+                        v['step'] = v['step'].to('cpu')
+
+    def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
+        # model
+        supported_classes = (SwiftModel, PreTrainedModel, PeftModel)
+        supported_names = ('SentenceTransformer')
+        if AutoModelForCausalLMWithValueHead is not None:
+            supported_classes = supported_classes + (AutoModelForCausalLMWithValueHead, )
+        save_safetensors = self.args.save_safetensors
+        if not isinstance(self.model, supported_classes) and self.model.__class__.__name__ not in supported_names:
+            if state_dict is None:
+                state_dict = self.model.state_dict()
+
+            _unwrap_model = unwrap_model(self.model)
+            if isinstance(_unwrap_model, supported_classes):
+                _unwrap_model.save_pretrained(output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
+            else:
+                logger.info('Trainer.model is not a `PreTrainedModel`, only saving its state dict.')
+                if save_safetensors:
+                    safetensors.torch.save_file(state_dict, os.path.join(output_dir, 'model.safetensors'))
+                else:
+                    torch.save(state_dict, os.path.join(output_dir, 'pytorch_model.bin'))
+        elif AutoModelForCausalLMWithValueHead and isinstance(self.model, AutoModelForCausalLMWithValueHead):
+            # save reward model
+            state_dict = self.model.state_dict()
+            decoder_state_dict, v_head_state_dict = {}, {}
+            for name, param in state_dict.items():
+                if name.startswith('v_head.'):
+                    v_head_state_dict[name] = param
+                else:
+                    decoder_state_dict[name.replace('pretrained_model.', '', 1)] = param
+            self.model.pretrained_model.save_pretrained(
+                output_dir, state_dict=decoder_state_dict or None, safe_serialization=save_safetensors)
+            if save_safetensors:
+                from safetensors.torch import save_file
+                save_file(
+                    v_head_state_dict, os.path.join(output_dir, 'value_head.safetensors'), metadata={'format': 'pt'})
+            else:
+                torch.save(v_head_state_dict, os.path.join(output_dir, 'value_head.bin'))
+        elif is_instance_of_ms_model(self.model):
+            PreTrainedModel.save_pretrained(
+                self.model, output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
+        elif self.args.train_type in extra_tuners:
+            extra_tuners[self.args.train_type].save_pretrained(
+                self.model, output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
+        else:
+            if self.model.__class__.__name__ != 'SentenceTransformer':
+                self.model.save_pretrained(output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
+            else:
+
+                @contextmanager
+                def save_context():
+                    save_pretrained = self.model[0].auto_model.save_pretrained
+                    _state_dict = {
+                        key[len('0.auto_model.'):] if 'auto_model' in key else key: value
+                        for key, value in state_dict.items()
+                    }
+                    self.model[0].auto_model.save_pretrained = partial(
+                        self.model[0].auto_model.save_pretrained, state_dict=_state_dict)
+                    yield
+                    self.model[0].auto_model.save_pretrained = save_pretrained
+
+                with save_context():
+                    self.model.save_pretrained(output_dir, safe_serialization=save_safetensors)
+                    # copy sentencetransformers files
+                    from swift.utils import copy_files_by_pattern
+                    copy_files_by_pattern(self.model.model_dir, output_dir, '*.py')
+                    copy_files_by_pattern(self.model.model_dir, output_dir, '*.json')
+
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        """Compatible with swift and peft"""
+        # If we are executing this function, we are the process zero, so we don't check for that.
+        output_dir = output_dir if output_dir is not None else self.args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        self._save_model(output_dir, state_dict)
+        # training_args.bin
+        torch.save(self.args, os.path.join(output_dir, 'training_args.bin'))
+        self._save_converted_model(output_dir)
+        # args.json
+        args_path = os.path.join(os.path.dirname(output_dir), 'args.json')
+        if os.path.exists(args_path):
+            shutil.copy(args_path, os.path.join(output_dir, 'args.json'))
+        # predict.jsonl
+        predict_jsonl = os.path.join(os.path.dirname(output_dir), 'predict.jsonl')
+        if os.path.exists(predict_jsonl):
+            shutil.move(predict_jsonl, os.path.join(output_dir, 'predict.jsonl'))
+
+        is_adapter = isinstance(self.model, (SwiftModel, PeftModel))
+        # tokenizer
+        if not is_adapter:
+            from swift.llm import save_checkpoint
+            additional_saved_files = self.model_meta.additional_saved_files
+            save_checkpoint(
+                None,
+                self.template.processor,
+                output_dir,
+                model_dirs=[self.model.model_dir],
+                additional_saved_files=additional_saved_files)
+            if getattr(self.model, 'origin_generation_config', None):
+                self.model.origin_generation_config.save_pretrained(output_dir)
+
+    def _fix_zero3_gather_all_parameters(self) -> None:
+        if is_deepspeed_zero3_enabled() and not hasattr(self.deepspeed, '_zero3_consolidated_16bit_state_dict_origin'):
+            parameters = inspect.signature(self.deepspeed._zero3_consolidated_16bit_state_dict).parameters
+            if 'exclude_frozen_parameters' in parameters:
+
+                def _zero3_consolidated_16bit_state_dict(model, exclude_frozen_parameters=False):
+                    unwrapped = unwrap_model(model)
+                    exclude_frozen_parameters = False
+                    if isinstance(unwrapped, SwiftModel) and unwrapped.has_additional_modules:
+                        exclude_frozen_parameters = True
+                    if isinstance(unwrapped, PeftModel):
+                        exclude_frozen_parameters = True
+                    return model._zero3_consolidated_16bit_state_dict_origin(exclude_frozen_parameters)
+
+                self.deepspeed._zero3_consolidated_16bit_state_dict_origin = (
+                    self.deepspeed._zero3_consolidated_16bit_state_dict)
+                self.deepspeed._zero3_consolidated_16bit_state_dict = MethodType(_zero3_consolidated_16bit_state_dict,
+                                                                                 self.deepspeed)
+
+    def _save_checkpoint(self, *args, **kwargs):
+        self.state.last_model_checkpoint = os.path.join(self.args.output_dir, f'checkpoint-{self.state.global_step}')
+        self._fix_zero3_gather_all_parameters()
+        result = super()._save_checkpoint(*args, **kwargs)
+        logger.info(f'Saving model checkpoint to {self.state.last_model_checkpoint}')
+        return result
+
+    @staticmethod
+    @contextmanager
+    def _fix_grad_norm_nan():
+        from accelerate import Accelerator
+        origin_clip_grad_norm_ = Accelerator.clip_grad_norm_
+
+        def clip_grad_norm_(self, parameters, *args, **kwargs):
+            # If NaN occurs, ignore weight updates.
+            parameters = list(parameters)
+            grad_norm = origin_clip_grad_norm_(self, parameters, *args, **kwargs)
+            if isinstance(grad_norm, torch.Tensor) and grad_norm.isnan().item():
+                for p in parameters:
+                    p.grad = None
+            return grad_norm
+
+        Accelerator.clip_grad_norm_ = clip_grad_norm_
+        try:
+            yield
+        finally:
+            Accelerator.clip_grad_norm_ = origin_clip_grad_norm_
+
+    def train(self, *args, **kwargs):
+        if self.model_meta.is_multimodal:
+            models = []
+            for model_name in ['model', 'ref_model', 'value_model']:
+                model = getattr(self, model_name, None)
+                if isinstance(model, nn.Module):
+                    models.append(model)
+
+            reward_model = getattr(self, 'reward_model', None)
+            if reward_model is not None:
+                if isinstance(reward_model, list):
+                    models.extend([m for m in reward_model if isinstance(m, nn.Module)])
+                elif isinstance(reward_model, nn.Module):
+                    models.append(reward_model)
+
+            models = list(set(models))  # Deduplicate
+            self.template.register_post_encode_hook(models)
+            logger.info(f'Successfully registered post_encode hook: {[model.__class__.__name__ for model in models]}.')
+        self._save_initial_model(self.args.output_dir)
+        with self.hub.patch_hub(), self._fix_grad_norm_nan():
+            res = super().train(*args, **kwargs)
+        self.template.remove_post_encode_hook()
+        return res
+
+    def push_to_hub(self, *args, **kwargs):
+        with self.hub.patch_hub():
+            return super().push_to_hub(*args, **kwargs)
+
+    def get_max_cuda_memory(self, device: Optional[Union[torch.device, int]] = None) -> float:
+        if device is None:
+            mems = [torch.cuda.max_memory_reserved(device=device) for device in range(get_device_count())]
+        else:
+            mems = [torch.cuda.max_memory_reserved(device=device)]
+        mem = sum(mems) / 1024**3
+        self.max_memory = max(self.max_memory, mem)
+        return mem
+
+    def _maybe_log_save_evaluate(self, tr_loss, *args, **kwargs):
+        if self.control.should_log and self.state.global_step > self._globalstep_last_logged:
+            self.control.should_log = False
+
+            # all_gather + mean() to get average loss over all processes
+            tr_loss_scalar = self._nested_gather(tr_loss).mean().item()
+            loss = tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged)
+            logs: Dict[str, float] = {'loss': loss}  # loss first
+
+            for k, metric in self._custom_metrics.items():
+                value = metric.compute()
+                if len(value) == 1:
+                    val = list(value.values())[0]
+                    logs[k] = val
+                else:
+                    for k_suffix, val in value.items():
+                        new_k = f'{k}_{k_suffix}'
+                        logs[new_k] = val
+                metric.reset()
+
+            if version.parse(transformers.__version__) >= version.parse('4.38'):
+                grad_norm = args[0]
+                if grad_norm is not None:
+                    logs['grad_norm'] = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
+            logs['learning_rate'] = self._get_learning_rate()
+            if not is_torch_npu_available():
+                logs['memory(GiB)'] = round(self.get_max_cuda_memory(), 2)
+
+            elapse_time = time.time() - self.start_time
+            logs['train_speed(iter/s)'] = round(self.state.global_step / elapse_time, 6)
+            for k in list(logs.keys()):
+                if logs[k] is None:
+                    logs.pop(k)
+            tr_loss -= tr_loss
+            self._total_loss_scalar += tr_loss_scalar
+            self._globalstep_last_logged = self.state.global_step
+            self.store_flos()
+            self.log(logs)
+
+        if self.args.eval_use_evalscope and self.control.should_evaluate:
+            self._evalscope_eval()
+        super()._maybe_log_save_evaluate(tr_loss, *args, **kwargs)
+
+    def create_optimizer_and_scheduler(self, num_training_steps: int):
+        if self.args.optimizer is not None:
+            from swift.plugin import optimizers_map
+            optimizer_callback = optimizers_map[self.args.optimizer]
+            self.optimizer, self.lr_scheduler = optimizer_callback(self.args, self.model, self.train_dataset)
+            if self.optimizer is None:
+                self.create_optimizer()
+            if self.lr_scheduler is None:
+                self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
+        else:
+            super().create_optimizer_and_scheduler(num_training_steps=num_training_steps)
+
+    def _compute_acc(self, outputs, labels) -> None:
+        args = self.args
+        acc_steps = args.acc_steps
+        preds = outputs.logits.argmax(dim=-1)
+        if self.state.global_step % acc_steps == 0:
+            if use_torchacc():
+                ta_trim_graph()
+                preds = preds.to('cpu')
+                labels = labels.to('cpu')
+            metrics = compute_acc(
+                preds, labels, acc_strategy=args.acc_strategy, is_encoder_decoder=self.template.is_encoder_decoder)
+            for k, v in metrics.items():
+                if k not in self._custom_metrics:
+                    self._custom_metrics[k] = MeanMetric(nan_value=None)
+                self._custom_metrics[k].update(v)
+
+    @torch.no_grad()
+    def _evalscope_eval(self):
+        from ..llm.eval.utils import EvalModel
+        from evalscope import TaskConfig, run_task
+        from evalscope.constants import EvalType
+
+        self.model.eval()
+        max_batch_size = self.args.per_device_eval_batch_size
+        custom_model = EvalModel(
+            self.model, self.template, max_batch_size=max_batch_size, model_name=f'model-step{self.state.global_step}')
+        task_config = TaskConfig(
+            model=custom_model,
+            eval_type=EvalType.CUSTOM,
+            datasets=self.args.eval_datasets,
+            dataset_args=self.args.eval_datasets_args,
+            limit=self.args.eval_limit,
+            work_dir=os.path.join(self.args.output_dir, 'eval'),
+            eval_batch_size=max_batch_size,
+            generation_config=self.args.eval_generation_config or {'max_tokens': 512},
+        )
+        # start evaluation
+        eval_report = run_task(task_config)
+        # convert to dict
+        eval_dict = {f'test_{k}': v.score for k, v in eval_report.items()}
+        self.log(eval_dict)
+
+        self.model.train()
+        return eval_dict
+
+    def get_batch_samples(self, *args, **kwargs):
+        res = super().get_batch_samples(*args, **kwargs)
+        if self.template.sequence_parallel_size == 1:
+            return res
+
+        batch_samples, num_items_in_batch = res
+        if num_items_in_batch is None:
+            num_items_in_batch = torch.tensor(0).to(args[2])
+        from swift.trainers.sequence_parallel import sequence_parallel
+        dist.all_reduce(num_items_in_batch, dist.ReduceOp.SUM, sequence_parallel.sp_group)
+        return batch_samples, num_items_in_batch
+
+
+class DataLoaderMixin:
+
+    def get_train_dataloader(self):
+        dataloader = None
+        if self.template.sequence_parallel_size > 1:
+            from swift.trainers.sequence_parallel import sequence_parallel
+            dataloader = sequence_parallel.get_dataloader(self, self.train_dataset, self._train_batch_size)
+        if dataloader is None:
+            # Higher efficiency
+            if self.train_dataset is None:
+                raise ValueError('Trainer: training requires a train_dataset.')
+            args = self.args
+            train_dataset = self.train_dataset
+
+            dataloader_params = {
+                'collate_fn': self.data_collator,
+                'num_workers': args.dataloader_num_workers,
+                'pin_memory': args.dataloader_pin_memory,
+                'persistent_workers': args.dataloader_persistent_workers,
+                'prefetch_factor': args.dataloader_prefetch_factor
+            }
+            batch_sampler_params = {
+                'drop_last': args.dataloader_drop_last,
+                'shuffle': args.train_dataloader_shuffle,
+                'data_seed': args.data_seed,
+            }
+
+            if hasattr(train_dataset, '__len__'):
+                batch_sampler = BatchSamplerShard(
+                    len(train_dataset), batch_size=self._train_batch_size, **batch_sampler_params)
+                dataloader = DataLoaderShard(train_dataset, batch_sampler, **dataloader_params)
+            else:
+                # IterableDataset
+                if dist.is_initialized() and dataloader_params['prefetch_factor']:
+                    dataloader_params['prefetch_factor'] = dataloader_params['prefetch_factor'] * dist.get_world_size()
+                dataloader = DataLoader(train_dataset, batch_size=self._train_batch_size, **dataloader_params)
+                dataloader = DataLoaderDispatcher(dataloader)
+
+        return dataloader
+
+    def get_eval_dataloader(self, eval_dataset=None):
+        dataloader = None
+        if self.template.sequence_parallel_size > 1:
+            from swift.trainers.sequence_parallel import sequence_parallel
+            if eval_dataset is None and self.eval_dataset is None:
+                raise ValueError('Trainer: evaluation requires an eval_dataset.')
+            eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
+            dataloader = sequence_parallel.get_dataloader(self, eval_dataset, self.args.eval_batch_size)
+        if dataloader is None:
+            return super().get_eval_dataloader(eval_dataset=eval_dataset)
+        return dataloader

ms-swift/swift/trainers/optimizers/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.

ms-swift/swift/trainers/optimizers/galore/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .utils import create_optimizer_and_scheduler, GaLoreConfig
+    from .adafactor import GaLoreAdafactor
+    from .adamw8bit import GaLoreAdamW8bit
+    from .adamw import GaLoreAdamW
+else:
+    _import_structure = {
+        'utils': ['GaLoreConfig', 'create_optimizer_and_scheduler'],
+        'adafactor': ['GaLoreAdafactor'],
+        'adamw8bit': ['GaLoreAdamW8bit'],
+        'adamw': ['GaLoreAdamW'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

ms-swift/swift/trainers/optimizers/galore/adafactor.py

@@ -0,0 +1,272 @@
+# copy dependencies from transformers/optimization.py
+# code borrowed from https://github.com/jiaweizzhao/GaLore
+import math
+
+import torch
+from torch.optim import Optimizer
+from transformers.utils.versions import require_version
+
+from .galore_projector import GaLoreProjector
+
+
+class Adafactor(Optimizer):
+    """
+    AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:
+    https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py
+
+    Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://arxiv.org/abs/1804.04235 Note that
+    this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and
+    `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+
+    Arguments:
+        params (`Iterable[nn.parameter.Parameter]`):
+            Iterable of parameters to optimize or dictionaries defining parameter groups.
+        lr (`float`, *optional*):
+            The external learning rate.
+        eps (`Tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`):
+            Regularization constants for square gradient and parameter scale respectively
+        clip_threshold (`float`, *optional*, defaults to 1.0):
+            Threshold of root mean square of final gradient update
+        decay_rate (`float`, *optional*, defaults to -0.8):
+            Coefficient used to compute running averages of square
+        beta1 (`float`, *optional*):
+            Coefficient used for computing running averages of gradient
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Weight decay (L2 penalty)
+        scale_parameter (`bool`, *optional*, defaults to `True`):
+            If True, learning rate is scaled by root mean square
+        relative_step (`bool`, *optional*, defaults to `True`):
+            If True, time-dependent learning rate is computed instead of external learning rate
+        warmup_init (`bool`, *optional*, defaults to `False`):
+            Time-dependent learning rate computation depends on whether warm-up initialization is being used
+
+    This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.
+
+    Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):
+
+        - Training without LR warmup or clip_threshold is not recommended.
+
+           - use scheduled LR warm-up to fixed LR
+           - use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235)
+        - Disable relative updates
+        - Use scale_parameter=False
+        - Additional optimizer operations like gradient clipping should not be used alongside Adafactor
+
+    Example:
+
+    ```python
+    Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)
+    ```
+
+    Others reported the following combination to work well:
+
+    ```python
+    Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+    ```
+
+    When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]
+    scheduler as following:
+
+    ```python
+    from transformers.optimization import Adafactor, AdafactorSchedule
+
+    optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+    lr_scheduler = AdafactorSchedule(optimizer)
+    trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))
+    ```
+
+    Usage:
+
+    ```python
+    # replace AdamW with Adafactor
+    optimizer = Adafactor(
+        model.parameters(),
+        lr=1e-3,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        relative_step=False,
+        scale_parameter=False,
+        warmup_init=False,
+    )
+    ```"""
+
+    def __init__(
+        self,
+        params,
+        lr=None,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        scale_parameter=True,
+        relative_step=True,
+        warmup_init=False,
+    ):
+        require_version('torch>=1.5.0')  # add_ with alpha
+        if lr is not None and relative_step:
+            raise ValueError('Cannot combine manual `lr` and `relative_step=True` options')
+        if warmup_init and not relative_step:
+            raise ValueError('`warmup_init=True` requires `relative_step=True`')
+
+        defaults = {
+            'lr': lr,
+            'eps': eps,
+            'clip_threshold': clip_threshold,
+            'decay_rate': decay_rate,
+            'beta1': beta1,
+            'weight_decay': weight_decay,
+            'scale_parameter': scale_parameter,
+            'relative_step': relative_step,
+            'warmup_init': warmup_init,
+        }
+        super().__init__(params, defaults)
+
+    @staticmethod
+    def _get_lr(param_group, param_state):
+        rel_step_sz = param_group['lr']
+        if param_group['relative_step']:
+            min_step = 1e-6 * param_state['step'] if param_group['warmup_init'] else 1e-2
+            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state['step']))
+        param_scale = 1.0
+        if param_group['scale_parameter']:
+            param_scale = max(param_group['eps'][1], param_state['RMS'])
+        return param_scale * rel_step_sz
+
+    @staticmethod
+    def _get_options(param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group['beta1'] is not None
+        return factored, use_first_moment
+
+    @staticmethod
+    def _rms(tensor):
+        return tensor.norm(2) / (tensor.numel()**0.5)
+
+    @staticmethod
+    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
+        # copy from fairseq's adafactor implementation:
+        # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
+        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """
+        Performs a single optimization step
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError('Adafactor does not support sparse gradients.')
+
+                state = self.state[p]
+
+                if 'step' not in state:
+                    state['step'] = 0
+
+                # GaLore Projection
+                if 'rank' in group:
+                    if 'projector' not in state:
+                        state['projector'] = GaLoreProjector(
+                            group['rank'],
+                            update_proj_gap=group['update_proj_gap'],
+                            scale=group['scale'],
+                            proj_type=group['proj_type'])
+
+                    grad = state['projector'].project(grad, state['step'])
+
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if 'RMS' not in state:
+                    state['step'] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state['exp_avg'] = torch.zeros_like(grad)
+                    if factored:
+                        state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state['exp_avg_sq_col'] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+                    else:
+                        state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                    state['RMS'] = 0
+                else:
+                    if use_first_moment:
+                        state['exp_avg'] = state['exp_avg'].to(grad)
+                    if factored:
+                        state['exp_avg_sq_row'] = state['exp_avg_sq_row'].to(grad)
+                        state['exp_avg_sq_col'] = state['exp_avg_sq_col'].to(grad)
+                    else:
+                        state['exp_avg_sq'] = state['exp_avg_sq'].to(grad)
+
+                p_data_fp32 = p
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state['step'] += 1
+                state['RMS'] = self._rms(p_data_fp32)
+                lr = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state['step'], group['decay_rate'])
+                update = (grad**2) + group['eps'][0]
+                if factored:
+                    exp_avg_sq_row = state['exp_avg_sq_row']
+                    exp_avg_sq_col = state['exp_avg_sq_col']
+
+                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
+                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state['exp_avg_sq']
+
+                    exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_((self._rms(update) / group['clip_threshold']).clamp_(min=1.0))
+                update.mul_(lr)
+
+                if use_first_moment:
+                    exp_avg = state['exp_avg']
+                    exp_avg.mul_(group['beta1']).add_(update, alpha=(1 - group['beta1']))
+                    update = exp_avg
+
+                # GaLore Projection Back
+                if 'rank' in group:
+                    update = state['projector'].project_back(update)
+
+                if group['weight_decay'] != 0:
+                    p_data_fp32.add_(p_data_fp32, alpha=(-group['weight_decay'] * lr))
+
+                p_data_fp32.add_(-update)
+
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p.copy_(p_data_fp32)
+
+        return loss
+
+
+GaLoreAdafactor = Adafactor

ms-swift/swift/trainers/optimizers/galore/galore_projector.py

@@ -0,0 +1,109 @@
+# code borrowed from https://github.com/jiaweizzhao/GaLore
+
+import torch
+
+
+class GaLoreProjector:
+
+    def __init__(self, rank, verbose=False, update_proj_gap=200, scale=1.0, proj_type='std'):
+        self.rank = rank
+        self.verbose = verbose
+        self.update_proj_gap = update_proj_gap
+        self.scale = scale
+        self.ortho_matrix = None
+        self.proj_type = proj_type
+
+    def project(self, full_rank_grad, iter):
+
+        if self.proj_type == 'std':
+            if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
+                low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t())
+            else:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='left')
+                low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad)
+        elif self.proj_type == 'reverse_std':
+            if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='left')
+                low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad)
+            else:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
+                low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t())
+        elif self.proj_type == 'right':
+            if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
+            low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t())
+        elif self.proj_type == 'left':
+            if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='left')
+            low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad)
+        elif self.proj_type == 'full':
+            if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='full')
+            low_rank_grad = torch.matmul(self.ortho_matrix[0].t(), full_rank_grad) @ self.ortho_matrix[1].t()
+
+        return low_rank_grad
+
+    def project_back(self, low_rank_grad):
+
+        if self.proj_type == 'std':
+            if low_rank_grad.shape[0] >= low_rank_grad.shape[1]:
+                full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix)
+            else:
+                full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
+        elif self.proj_type == 'reverse_std':
+            if low_rank_grad.shape[0] <= low_rank_grad.shape[1]:  # note this is different from std
+                full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
+            else:
+                full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix)
+        elif self.proj_type == 'right':
+            full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix)
+        elif self.proj_type == 'left':
+            full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
+        elif self.proj_type == 'full':
+            full_rank_grad = torch.matmul(self.ortho_matrix[0], low_rank_grad) @ self.ortho_matrix[1]
+
+        return full_rank_grad * self.scale
+
+    # svd decomposition
+    def get_orthogonal_matrix(self, weights, rank, type):
+        module_params = weights
+
+        if module_params.data.dtype != torch.float:
+            float_data = False
+            original_type = module_params.data.dtype
+            original_device = module_params.data.device
+            matrix = module_params.data.float()
+        else:
+            float_data = True
+            matrix = module_params.data
+
+        U, s, Vh = torch.linalg.svd(matrix, full_matrices=False)
+
+        # make the smaller matrix always to be orthogonal matrix
+        if type == 'right':
+            A = U[:, :rank] @ torch.diag(s[:rank])
+            B = Vh[:rank, :]
+
+            if not float_data:
+                B = B.to(original_device).type(original_type)
+            return B
+        elif type == 'left':
+            A = U[:, :rank]
+            B = torch.diag(s[:rank]) @ Vh[:rank, :]
+            if not float_data:
+                A = A.to(original_device).type(original_type)
+            return A
+        elif type == 'full':
+            A = U[:, :rank]
+            B = Vh[:rank, :]
+            if not float_data:
+                A = A.to(original_device).type(original_type)
+                B = B.to(original_device).type(original_type)
+            return [A, B]
+        else:
+            raise ValueError('type should be left, right or full')

ms-swift/swift/trainers/optimizers/galore/utils.py

@@ -0,0 +1,214 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import importlib
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple, Union
+
+import torch
+from torch import nn
+from torch.optim import Optimizer
+from transformers import Trainer, TrainingArguments, get_scheduler
+
+from swift.utils import get_logger
+
+try:
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import LRScheduler
+
+logger = get_logger()
+
+
+@dataclass
+class GaLoreConfig:
+    """
+    The configuration class for the Galore module.
+
+
+    See https://arxiv.org/abs/2403.03507
+
+    Args:
+        rank (`int`): The galore rank
+        target_modules (`Union[str, List[str]]`): The target modules to use, if `None`,
+            will use all attn and mlp linears
+        update_proj_gap(`int`): The projection update interval for galore
+        proj_type(`str`) The project type of Galore, valid values are `std`,
+            `reverse_std`, `right`, `left`, `full`
+        galore_scale(float): the scale of gradient
+        optim_per_parameter(bool): Gives one optimizer per parameter
+    """
+    rank: int = 128
+    target_modules: Union[str, List[str]] = None
+    update_proj_gap: int = 50
+    galore_scale: float = 1.0
+    proj_type: str = 'std'
+    optim_per_parameter: bool = False
+    quantize: bool = False
+    proj_quant: bool = False
+    proj_bits: int = 4
+    proj_group_size: int = 256
+    cos_threshold: float = 0.4
+    gamma_proj: int = 2
+    queue_size: int = 5
+
+
+class GaloreOptimizerWrapper(Optimizer):
+
+    def __init__(self, optimizers: Dict[Any, Optimizer]):
+        self.optimizers = optimizers
+        super().__init__([torch.tensor([1., 2., 3.])], {'lr': 1.})
+
+    def zero_grad(self, *args, **kwargs) -> None:
+        for optim in self.optimizers.values():
+            optim.zero_grad(*args, **kwargs)
+
+    def step(self, *args, **kwargs) -> None:
+        for optim in self.optimizers.values():
+            optim.step(*args, **kwargs)
+
+
+class GaloreSchedulerWrapper(LRScheduler):
+
+    def __init__(self, lr_schedulers: Dict[Any, LRScheduler]):
+        self.lr_schedulers = lr_schedulers
+
+    def step(self, *args, **kwargs) -> None:
+        for lr_scheduler in self.lr_schedulers.values():
+            lr_scheduler.step(*args, **kwargs)
+        self._last_lr = lr_scheduler.get_last_lr()
+
+
+def create_optimizer_and_scheduler(model: nn.Module, args: TrainingArguments, config: GaLoreConfig, max_steps,
+                                   **defaults):
+    galore_params = []
+    for module_name, module in model.named_modules():
+        if not isinstance(module, (nn.Linear, nn.Embedding)) or \
+                not any(target_key in module_name for target_key in config.target_modules):
+            continue
+
+        if not module.weight.requires_grad:
+            continue
+
+        logger.info(f'Enable GaLore for weights in module: {module_name}')
+        galore_params.append(module.weight)
+
+    id_galore_params = [id(p) for p in galore_params]
+    galore_defaults = {
+        'rank': config.rank,
+        'update_proj_gap': config.update_proj_gap,
+        'scale': config.galore_scale,
+        'proj_type': config.proj_type,
+        **defaults
+    }
+    if config.quantize:
+        galore_defaults['quant'] = config.proj_quant
+        galore_defaults['quant_n_bit'] = config.proj_bits
+        galore_defaults['quant_group_size'] = config.proj_group_size
+        galore_defaults['cos_threshold'] = config.cos_threshold
+        galore_defaults['gamma_proj'] = config.gamma_proj
+        galore_defaults['queue_size'] = config.queue_size
+    optim_cls, optim_kwargs = get_optimizer(args, config)
+
+    if config.optim_per_parameter and not config.quantize:
+        # q-galore does not support optim_per_parameter
+        optimizer_dict = {}
+        galore_defaults['update_proj_gap'] = galore_defaults['update_proj_gap'] * 2
+        for p in model.parameters():
+            if p.requires_grad:
+                if id(p) in id_galore_params:
+                    optimizer_dict[p] = optim_cls([{'params': [p], **galore_defaults}], **optim_kwargs)
+                else:
+                    optimizer_dict[p] = optim_cls([{'params': [p], **defaults}], **optim_kwargs)
+
+        # get scheduler dict
+        scheduler_dict = {}
+        for p in model.parameters():
+            if p.requires_grad:
+                scheduler_dict[p] = get_scheduler(
+                    optimizer=optimizer_dict[p],
+                    name=args.lr_scheduler_type,
+                    num_training_steps=max_steps * 2,
+                    num_warmup_steps=args.warmup_steps * 2,
+                    scheduler_specific_kwargs=args.lr_scheduler_kwargs,
+                )
+
+        return GaloreOptimizerWrapper(optimizer_dict), GaloreSchedulerWrapper(scheduler_dict)
+    else:
+        decay_parameters = Trainer.get_decay_parameter_names(Trainer, model)
+        param_groups = [{
+            'params': galore_params,
+            **galore_defaults,
+        }]
+        param_groups.extend([
+            {
+                'params': [
+                    p for n, p in model.named_parameters()
+                    if (n in decay_parameters and id(p) not in id_galore_params and p.requires_grad)
+                ],
+                'weight_decay':
+                defaults['weight_decay'],
+            },
+            {
+                'params': [
+                    p for n, p in model.named_parameters()
+                    if (n not in decay_parameters and id(p) not in id_galore_params and p.requires_grad)
+                ],
+                'weight_decay':
+                0.0,
+            },
+        ])
+        optim = optim_cls(param_groups, **optim_kwargs)
+        scheduler = get_scheduler(
+            optimizer=optim,
+            name=args.lr_scheduler_type,
+            num_training_steps=max_steps,
+            num_warmup_steps=args.warmup_steps,
+            scheduler_specific_kwargs=args.lr_scheduler_kwargs,
+        )
+        return optim, scheduler
+
+
+def get_optimizer(args: TrainingArguments, config: GaLoreConfig) -> Tuple[Any, Any]:
+    # parse args.optim_args
+    optim_args = {}
+    if args.optim_args:
+        for mapping in args.optim_args.replace(' ', '').split(','):
+            key, value = mapping.split('=')
+            optim_args[key] = value
+
+    optimizer_kwargs = {'lr': args.learning_rate}
+
+    adam_kwargs = {
+        'betas': (args.adam_beta1, args.adam_beta2),
+        'eps': args.adam_epsilon,
+    }
+    if args.optim == 'adafactor':
+        from .adafactor import GaLoreAdafactor
+        optimizer_cls = GaLoreAdafactor
+        optimizer_kwargs.update({'scale_parameter': False, 'relative_step': False})
+    elif args.optim in ('adamw_hf', 'adamw_torch'):
+        if config.quantize:
+            assert importlib.util.find_spec('q_galore_torch') is not None, \
+                'Please install q-galore by `pip install q_galore_torch`'
+            logger.info('If you encounter `absmax2` error, please downgrade your bitsandbytes to 0.40.0')
+            from swift.utils import get_dist_setting
+            _, _, world_size, _ = get_dist_setting()
+            if world_size > 1:
+                # from q_galore_torch import QGaLoreAdamW8bit_simulate as GaLoreAdamW
+                from q_galore_torch import QGaLoreAdamW8bit as GaLoreAdamW
+            else:
+                from q_galore_torch import QGaLoreAdamW8bit as GaLoreAdamW
+        else:
+            from .adamw import GaLoreAdamW
+        optimizer_cls = GaLoreAdamW
+        optimizer_kwargs.update(adam_kwargs)
+    elif 'adamw' in args.optim and '8bit' in args.optim:
+        try:
+            from .adamw8bit import GaLoreAdamW8bit
+            optimizer_cls = GaLoreAdamW8bit
+            optimizer_kwargs.update(adam_kwargs)
+            optimizer_kwargs.update({'optim_bits': 8, 'is_paged': 'paged' in args.optim})
+        except ImportError:
+            raise ValueError('Trainer tried to instantiate bnb optimizer but bnb is not installed!')
+    else:
+        raise ValueError(f'Galore not supported for optimizer type: {args.optim}')
+    return optimizer_cls, optimizer_kwargs

ms-swift/swift/trainers/rlhf_arguments.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from trl import CPOConfig as HfCPOConfig
+from trl import DPOConfig as HfDPOConfig
+from trl import GRPOConfig as HfGRPOConfig
+from trl import KTOConfig as HfKTOConfig
+from trl import ORPOConfig as HfORPOConfig
+from trl import PPOConfig as HfPPOConfig
+from trl import RewardConfig as HfRewardConfig
+
+from .arguments import GRPOArgumentsMixin, SwiftArgumentsMixin
+
+
+@dataclass
+class DPOConfig(SwiftArgumentsMixin, HfDPOConfig):
+    pass
+
+
+@dataclass
+class CPOConfig(SwiftArgumentsMixin, HfCPOConfig):
+    pass
+
+
+@dataclass
+class ORPOConfig(SwiftArgumentsMixin, HfORPOConfig):
+    pass
+
+
+@dataclass
+class KTOConfig(SwiftArgumentsMixin, HfKTOConfig):
+    pass
+
+
+@dataclass
+class RewardConfig(SwiftArgumentsMixin, HfRewardConfig):
+    pass
+
+
+@dataclass
+class PPOConfig(SwiftArgumentsMixin, HfPPOConfig):
+    pass
+
+
+@dataclass
+class GRPOConfig(GRPOArgumentsMixin, SwiftArgumentsMixin, HfGRPOConfig):
+    stop_words: List[str] = field(default_factory=list)
+
+    def __post_init__(self):
+        from swift.llm.argument.base_args.model_args import ModelArguments
+        super().__post_init__()
+        if self.cosine_max_len is None:
+            self.cosine_max_len = self.max_completion_length
+        self.vllm_limit_mm_per_prompt = ModelArguments.parse_to_dict(self.vllm_limit_mm_per_prompt)
+
+        if self.deepspeed and 'zero_optimization' in self.deepspeed and self.deepspeed['zero_optimization'][
+                'stage'] == 3:
+            # https://github.com/modelscope/ms-swift/issues/3237
+            self.deepspeed['zero_optimization']['stage3_prefetch_bucket_size'] = 0
+            self.deepspeed_plugin.hf_ds_config.config['zero_optimization']['stage3_prefetch_bucket_size'] = 0
+
+        # https://github.com/modelscope/ms-swift/issues/3863
+        self.dataloader_drop_last = True

ms-swift/swift/trainers/rlhf_trainer/kto_trainer.py

@@ -0,0 +1,69 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from peft import PeftModel
+from transformers import PreTrainedModel
+from trl import KTOTrainer as HFKTOTrainer
+
+from swift.utils import get_logger
+from ..mixin import SwiftMixin
+from .rlhf_mixin import RLHFTrainerMixin
+
+logger = get_logger()
+
+del HFKTOTrainer.__init__
+
+
+class KTOTrainer(RLHFTrainerMixin, SwiftMixin, HFKTOTrainer):
+
+    def __init__(self,
+                 model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
+                 ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
+                 *_args,
+                 **kwargs):
+        args = kwargs['args']
+        args.disable_dropout = True
+        self.desirable_weight = args.desirable_weight
+        self.undesirable_weight = args.undesirable_weight
+        self.precompute_ref_log_probs = args.precompute_ref_log_probs
+        self.is_peft_model = isinstance(model, PeftModel)
+        if hasattr(args, 'loss_type'):
+            self.loss_type = args.loss_type
+        else:
+            self.loss_type = 'kto'
+
+        self.ref_adapter_name = None
+        # Not all losses require a KL calculation
+        self.calculate_KL = True
+        if self.loss_type in ['apo_zero_unpaired']:
+            self.calculate_KL = False
+        super().__init__(model, ref_model, *_args, **kwargs)
+
+    def forward(
+        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        is_kl = True
+
+        def _add_data_hook(model, args, kwargs):
+            nonlocal is_kl
+            if is_kl:
+                kwargs = {k[len('KL_completion_'):]: v for k, v in batch.items() if k.startswith('KL_completion_')}
+            else:
+                kwargs = {k[len('completion_'):]: v for k, v in batch.items() if k.startswith('completion_')}
+            is_kl = not is_kl
+            return (), kwargs
+
+        @contextmanager
+        def _patch_model_call():
+            handle = model.register_forward_pre_hook(_add_data_hook, with_kwargs=True, prepend=True)
+
+            try:
+                yield
+            finally:
+                handle.remove()
+
+        with _patch_model_call():
+            return super().forward(model, batch)

ms-swift/swift/trainers/rlhf_trainer/orpo_trainer.py

@@ -0,0 +1,19 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Optional, Union
+
+import torch.nn as nn
+from transformers import PreTrainedModel
+from trl import ORPOTrainer as HFORPOTrainer
+
+from ..mixin import SwiftMixin
+from .rlhf_mixin import RLHFTrainerMixin
+
+del HFORPOTrainer.__init__
+
+
+class ORPOTrainer(RLHFTrainerMixin, SwiftMixin, HFORPOTrainer):
+
+    def __init__(self, model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, *_args, **kwargs):
+        ref_model = kwargs.get('ref_model')
+        assert ref_model is None, 'ORPO does not require a ref_model.'
+        super().__init__(model, *_args, **kwargs)

ms-swift/swift/trainers/rlhf_trainer/ppo_trainer.py

@@ -0,0 +1,65 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import inspect
+from contextlib import contextmanager
+
+import transformers
+from packaging import version
+from torch.utils.data import DataLoader
+from transformers import PreTrainedModel
+from trl import PPOTrainer as HFPPOTrainer
+
+from swift.utils import patch_getattr
+from ..mixin import SwiftMixin
+
+ppo_trainer_init = HFPPOTrainer.__init__
+del HFPPOTrainer.__init__
+
+
+class PPOTrainer(SwiftMixin, HFPPOTrainer):
+
+    @staticmethod
+    @contextmanager
+    def _patch_dataloader(collate_fn):
+        __init__ = DataLoader.__init__
+
+        def __new_init__(self, *args, **kwargs):
+            kwargs['collate_fn'] = collate_fn
+            __init__(self, *args, **kwargs)
+
+        DataLoader.__init__ = __new_init__
+        try:
+            yield
+        finally:
+            DataLoader.__init__ = __init__
+
+    def __init__(self, model: PreTrainedModel, ref_model: PreTrainedModel, *_args, **kwargs):
+        super().__init__(model, *_args, **{k: v for k, v in kwargs.items() if k not in {'reward_model', 'value_model'}})
+        with self._patch_dataloader(kwargs['data_collator']):
+            new_kwargs = {
+                k: v
+                for k, v in kwargs.items()
+                if k in ['train_dataset', 'data_collator', 'reward_model', 'value_model', 'eval_dataset']
+            }
+            parameters = inspect.signature(ppo_trainer_init).parameters
+            if 'config' in parameters:
+                new_kwargs['config'] = kwargs['args']
+            else:
+                new_kwargs['args'] = kwargs['args']
+            if 'processing_class' in parameters:
+                new_kwargs['processing_class'] = self.tokenizer
+            else:
+                new_kwargs['tokenizer'] = self.tokenizer
+            ppo_trainer_init(self, model=model, ref_model=ref_model, **new_kwargs)
+        unwrap_model = self.accelerator.unwrap_model(self.model)
+        patch_getattr(unwrap_model.__class__, 'policy')
+
+    def train(self, *args, **kwargs):
+        # remove args that are not needed for the HFPPOTrainer
+        super().train()
+
+    def _save_checkpoint(self, *args, **kwargs):
+        if version.parse(transformers.__version__) >= version.parse('4.47'):
+            metrics = kwargs.pop('metrics', None)
+            trial = kwargs.get('trial')
+            self._determine_best_metric(metrics=metrics, trial=trial)
+        return super()._save_checkpoint(*args, **kwargs)

ms-swift/swift/trainers/rlhf_trainer/reward_trainer.py

@@ -0,0 +1,78 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from collections import defaultdict
+from typing import Any, Dict, Tuple, Union
+
+import pandas as pd
+import torch
+import torch.nn as nn
+from accelerate.utils import gather_object
+from transformers import PreTrainedModel
+from trl import RewardTrainer as HFRewardTrainer
+from trl.trainer.utils import print_rich_table
+
+from ..mixin import SwiftMixin
+from .rlhf_mixin import RLHFTrainerMixin
+
+del HFRewardTrainer.__init__
+
+
+class RewardTrainer(RLHFTrainerMixin, SwiftMixin, HFRewardTrainer):
+
+    def compute_loss(self,
+                     model: Union[PreTrainedModel, nn.Module],
+                     inputs: Dict[str, Union[torch.Tensor, Any]],
+                     return_outputs=False,
+                     num_items_in_batch=None) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
+        inputs.pop('labels', None)  # not use
+        attention_mask = inputs['attention_mask']
+        batch_size = attention_mask.shape[0] // 2
+        rewards = model(**inputs).logits
+        rewards_chosen, rewards_rejected = torch.split(rewards, batch_size, dim=0)
+        if 'margin' in inputs:
+            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs['margin']).mean()
+        else:
+            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+        if self.args.center_rewards_coefficient is not None:
+            loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected)**2)
+        # compat transformers>=4.46.*
+        if num_items_in_batch is not None and self.model_accepts_loss_kwargs:
+            loss /= self.args.gradient_accumulation_steps
+        if return_outputs:
+            return loss, {
+                'rewards_chosen': rewards_chosen,
+                'rewards_rejected': rewards_rejected,
+            }
+        return loss
+
+    def visualize_samples(self, num_print_samples: int):
+        """
+        Visualize the reward model logits prediction
+
+        Args:
+            num_print_samples (`int`, defaults to `4`):
+                The number of samples to print. Set to `-1` to print all samples.
+        """
+        eval_dataloader = self.get_eval_dataloader()
+        table = defaultdict(list)
+        for _, inputs in enumerate(eval_dataloader):
+            _, logits, _ = self.prediction_step(self.model, inputs, prediction_loss_only=False)
+            input_ids = inputs['input_ids']
+            attention_mask = inputs['attention_mask']
+            sequence_lengths = ((torch.eq(attention_mask, 0).int().argmax(-1) - 1) % attention_mask.shape[1]).tolist()
+            text = [self.template.safe_decode(tokens[:sequence_lengths[i]]) for i, tokens in enumerate(input_ids)]
+            batch_size = input_ids.shape[0] // 2
+            chosen_text, rejected_text = text[:batch_size], text[batch_size:]
+            table['chosen_text'].extend(gather_object(chosen_text))
+            table['rejected_text'].extend(gather_object(rejected_text))
+            table['logits'].extend(
+                gather_object([[round(inner_item, 4) for inner_item in item] for item in logits.tolist()]))
+            if 0 <= num_print_samples <= len(table['chosen_text']):
+                break
+        df = pd.DataFrame(table)
+        if self.accelerator.process_index == 0:
+            print_rich_table(df[:num_print_samples])
+            if 'wandb' in self.args.report_to:
+                import wandb
+
+                if wandb.run is not None:
+                    wandb.log({'completions': wandb.Table(dataframe=df)})

ms-swift/swift/trainers/rlhf_trainer/rlhf_mixin.py

@@ -0,0 +1,104 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from collections import defaultdict
+from contextlib import contextmanager, nullcontext
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+try:
+    from trl import AutoModelForCausalLMWithValueHead
+except (ImportError, RuntimeError):
+    AutoModelForCausalLMWithValueHead = None
+
+
+class RLHFTrainerMixin:
+
+    def __init__(self,
+                 model: Optional[Union[PreTrainedModel, nn.Module]] = None,
+                 ref_model: Optional[Union[PreTrainedModel, nn.Module]] = None,
+                 *_args,
+                 **kwargs):
+        from trl.trainer import disable_dropout_in_model
+        from swift.llm import HfConfigFactory
+        self.ref_model = ref_model
+        self._stored_metrics = defaultdict(lambda: defaultdict(list))
+        args = kwargs['args']
+        self.beta = getattr(args, 'beta', 0.0)
+        if getattr(args, 'disable_dropout', False):
+            disable_dropout_in_model(model)
+            if self.ref_model is not None:
+                disable_dropout_in_model(self.ref_model)
+
+        self.is_encoder_decoder = kwargs['template'].is_encoder_decoder
+        self.aux_loss_enabled = getattr(model.config, 'output_router_logits', False)
+        self._peft_has_been_casted_to_bf16 = False
+        self.generate_during_eval = getattr(args, 'generate_during_eval', False)
+        if self.is_encoder_decoder:
+            self.decoder_start_token_id = HfConfigFactory.get_config_attr(model.config, 'decoder_start_token_id')
+            self.pad_token_id = HfConfigFactory.get_config_attr(model.config, 'pad_token_id')
+        # not use
+        self.is_vision_model = False
+        self.label_pad_token_id = -100
+        self.use_dpo_data_collator = True
+        super().__init__(model, *_args, **kwargs)
+        if is_deepspeed_zero3_enabled() and ref_model is not None:
+            try:
+                from trl.models.utils import prepare_deepspeed
+            except ImportError as e:
+                raise ImportError('Please install trl>=0.14 via `pip install "trl>=0.14"`') from e
+            prepare_deepspeed(self.ref_model, self.accelerator)  # Does not wrap DeepSpeedEngine
+        self.padding_value = self.tokenizer.pad_token_id
+
+    def concatenated_forward(
+        self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        model_kwargs = batch.copy()
+        labels = model_kwargs.pop('labels', None)
+        if self.is_encoder_decoder:
+            model_kwargs['labels'] = labels
+
+        if self.aux_loss_enabled:
+            model_kwargs['output_router_logits'] = True
+        outputs = model(**model_kwargs, use_cache=False)
+        model_kwargs['labels'] = labels
+        model_kwargs['chosen_labels'] = torch.zeros(model_kwargs['labels'].shape[0] // 2)  # just get shape
+        if outputs.logits.shape[1] != labels.shape[1]:
+            # for llava, the model returns logits for the entire sequence, including the image tokens
+            # (placed before the text tokens)
+            outputs.logits = outputs.logits[:, -labels.shape[1]:]
+        for key in ['input_ids', 'attention_mask', 'labels']:
+            model_kwargs[f'concatenated_{key}'] = model_kwargs.pop(key, None)
+        if self.__class__.__name__ == 'ORPOTrainer':  # Pass-through labels
+            model_kwargs['concatenated_input_ids'] = model_kwargs['concatenated_labels']
+
+        @contextmanager
+        def _patch_concatenated_forward():
+            _old_concatenated_inputs = self.concatenated_inputs
+            _old_model_call = model.__class__.__call__
+            self.concatenated_inputs = lambda *args, **kwargs: model_kwargs
+            model.__class__.__call__ = lambda *args, **kwargs: outputs
+            try:
+                yield
+            finally:
+                self.concatenated_inputs = _old_concatenated_inputs
+                model.__class__.__call__ = _old_model_call
+
+        with _patch_concatenated_forward():
+            return super().concatenated_forward(model, model_kwargs)
+
+    def get_batch_logps(self, logits: torch.FloatTensor, labels: torch.LongTensor, *args, **kwargs):
+        if self.is_encoder_decoder:
+            labels = labels.clone()  # fix trl bug
+        return super().get_batch_logps(logits, labels, *args, **kwargs)
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        res = super().compute_loss(model, inputs, return_outputs=return_outputs)
+        # compat transformers>=4.46.*
+        if num_items_in_batch is not None and self.model_accepts_loss_kwargs:
+            loss = res[0] if return_outputs else res
+            loss /= self.args.gradient_accumulation_steps
+            return (loss, res[1:]) if return_outputs else loss
+        return res

ms-swift/swift/trainers/rlhf_trainer/utils.py

@@ -0,0 +1,132 @@
+from contextlib import contextmanager
+from types import MethodType
+from typing import Any, List, Optional
+
+import torch
+from peft.tuners import lora
+from peft.tuners.lora import LoraLayer
+
+
+def round_robin(num_reqs, num_workers):
+    """Distribute requests evenly across workers using round-robin algorithm.
+
+    Args:
+        num_reqs (int): Total number of requests to distribute
+        num_workers (int): Number of available workers
+
+    Returns:
+        list: A list of lists where each sublist contains the request indices
+                assigned to that particular node
+    """
+    distribution = [[] for _ in range(num_workers)]
+    for idx in range(num_reqs):
+        worker_id = idx % num_workers
+        distribution[worker_id].append(idx)
+    return distribution
+
+
+@contextmanager
+def patch_lora_merge(model, parameter_group=None):
+    """Patch LoraLayer's merge and get_delta_weight methods for controlled merging.
+
+    Args:
+        model: The PEFT model to patch
+        parameter_group: Optional list of parameter names to restrict merging
+
+    Yields:
+        The patched model (context manager ensures cleanup)
+    """
+    from peft.tuners.tuners_utils import check_adapters_to_merge
+
+    def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
+        if parameter_group and all(self.name not in pg for pg in parameter_group):
+            return  # Skip if not in target parameter group
+        adapter_names = check_adapters_to_merge(self, adapter_names)
+        if not adapter_names:
+            return
+
+        for active_adapter in adapter_names:
+            if active_adapter in self.lora_A.keys():
+                base_layer = self.get_base_layer()
+                if self.use_dora.get(active_adapter, False):
+                    self.lora_magnitude_vector[active_adapter].weight.data = \
+                        self.lora_magnitude_vector[active_adapter].weight.data.to(base_layer.weight.device)
+
+        return self.merge_origin(safe_merge, adapter_names)
+
+    def get_delta_weight(self, adapter) -> torch.Tensor:
+        # Ensure tensors are on correct device
+        if isinstance(self, lora.Embedding):
+            self.lora_embedding_A[adapter].data = self.lora_embedding_A[adapter].data.to(self.base_layer.weight.device)
+            self.lora_embedding_B[adapter].data = self.lora_embedding_B[adapter].data.to(self.base_layer.weight.device)
+        else:
+            self.lora_A[adapter].weight.data = self.lora_A[adapter].weight.data.to(self.base_layer.weight.device)
+            self.lora_B[adapter].weight.data = self.lora_B[adapter].weight.data.to(self.base_layer.weight.device)
+        return self.get_delta_weight_origin(adapter).to(self.base_layer.weight.device)
+
+    def _cache_pop(self, key: str) -> Any:
+        value = self._caches.pop(key).to(self.base_layer.weight.device)
+        return value
+
+    # Patch all LoraLayer instances
+    for name, module in model.named_modules():
+        if isinstance(module, LoraLayer):
+            module.name = name
+            if not hasattr(module, 'merge_origin') and hasattr(module, 'base_layer'):
+                module.merge_origin = module.merge
+                module.merge = MethodType(merge, module)
+                module.get_delta_weight_origin = module.get_delta_weight
+                module.get_delta_weight = MethodType(get_delta_weight, module)
+                module._cache_pop_origin = module._cache_pop
+                module._cache_pop = MethodType(_cache_pop, module)
+
+    try:
+        yield model
+    finally:
+        # Cleanup: restore original methods
+        for module in model.modules():
+            if isinstance(module, LoraLayer):
+                if hasattr(module, 'merge_origin'):
+                    module.merge = module.merge_origin
+                    del module.merge_origin
+                    module.get_delta_weight = module.get_delta_weight_origin
+                    del module.get_delta_weight_origin
+                    module._cache_pop = module._cache_pop_origin
+                    del module._cache_pop_origin
+
+
+@contextmanager
+def patch_lora_unmerge(model):
+    """Patch the unmerge method to ensure proper device handling."""
+
+    def _cache_pop_patched(self, key: str) -> Any:
+        value = self._caches.pop(key).to(self.base_layer.weight.device)
+        return value
+
+    def unmerge_patched(self):
+        if not self.merged:
+            return
+        # Move magnitude vectors to correct device first
+        for adapter in list(self.merged_adapters):
+            if self.use_dora.get(adapter, False):
+                self.lora_magnitude_vector[adapter].weight.data = \
+                    self.lora_magnitude_vector[adapter].weight.data.to(self.base_layer.weight.device)
+
+        return self.unmerge_origin()
+
+    for module in model.modules():
+        if isinstance(module, LoraLayer) and not hasattr(module, 'unmerge_origin'):
+            module.unmerge_origin = module.unmerge
+            module.unmerge = MethodType(unmerge_patched, module)
+            module._cache_pop_origin = module._cache_pop
+            module._cache_pop = MethodType(_cache_pop_patched, module)
+
+    try:
+        yield model
+    finally:
+        for module in model.modules():
+            if isinstance(module, LoraLayer) and hasattr(module, 'unmerge_origin'):
+                module.unmerge = module.unmerge_origin
+                del module.unmerge_origin
+                module._cache_pop = module._cache_pop_origin
+                del module._cache_pop_origin

ms-swift/swift/trainers/rlhf_trainer/vllm_client.py

@@ -0,0 +1,212 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+# Code partially sourced from Hugging Face TRL
+
+import atexit
+import logging
+import time
+from typing import List, Optional
+
+import requests
+import torch
+from dacite import from_dict
+from requests import ConnectionError
+from torch import nn
+
+from swift.llm import AdapterRequest, InferRequest, Template
+from swift.llm.infer.protocol import ChatCompletionResponse, RequestConfig
+from swift.plugin import Metric
+from swift.utils import is_vllm_ascend_available, is_vllm_available
+
+if is_vllm_available():
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.distributed.utils import StatelessProcessGroup
+
+    if is_vllm_ascend_available():
+        from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator  # noqa
+
+logger = logging.getLogger(__name__)
+
+
+class VLLMClient:
+    """
+    A client class to interact with a vLLM server.
+
+    This class provides methods to infer completions, initialize and manage weight update groups, and update model
+    weights in a distributed setting. Before using it, start the vLLM server with `trl vllm-serve`.
+
+    Args:
+        host (`str`, *optional*, defaults to `"0.0.0.0"`):
+            IP address of the vLLM server.
+        server_port (`int`, *optional*, defaults to `8000`):
+            Port number of the vLLM server.
+        group_port (`int`, *optional*, defaults to `51216`):
+            Port number for the weight update group.
+        connection_timeout (`float`, *optional*, defaults to `0.0`):
+            Total timeout duration in seconds to wait for the server to be up. If the server is not up after the
+            timeout, a `ConnectionError` is raised.
+    """
+
+    def __init__(self,
+                 host: str = '0.0.0.0',
+                 server_port: int = 8000,
+                 group_port: int = 51216,
+                 connection_timeout: float = 0.0):
+        if not is_vllm_available():
+            raise ImportError('vLLM is not installed. Please install it with `pip install vllm`.')
+
+        self.session = requests.Session()
+        self.host = host
+        self.server_port = server_port
+        self.group_port = group_port
+        self.check_server(connection_timeout)  # check server and fail after timeout
+
+    def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0):
+        """
+        Check server availability with retries on failure, within a total timeout duration. If the server is not up
+        after the total timeout duration, raise a `ConnectionError`.
+
+        Args:
+            retry_interval (`float`, *optional*, defaults to `2.0`):
+                Interval in seconds between retries.
+            total_timeout (`float`, *optional*, defaults to `0.0`):
+                Total timeout duration in seconds.
+        """
+        url = f'http://{self.host}:{self.server_port}/health/'
+        start_time = time.time()  # Record the start time
+
+        while True:
+            try:
+                response = requests.get(url)
+            except requests.exceptions.RequestException as exc:
+                # Check if the total timeout duration has passed
+                elapsed_time = time.time() - start_time
+                if elapsed_time >= total_timeout:
+                    raise ConnectionError(
+                        f"The vLLM server can't be reached at {self.host}:{self.server_port} after {total_timeout} "
+                        'seconds. Make sure the server is running by running `swift deploy`.') from exc
+            else:
+                if response.status_code == 200:
+                    logger.info('Server is up!')
+                    return None
+
+            # Retry logic: wait before trying again
+            logger.info(f'Server is not up yet. Retrying in {retry_interval} seconds...')
+            time.sleep(retry_interval)
+
+    def infer(
+        self,
+        infer_requests: List[InferRequest],
+        request_config: Optional[RequestConfig] = None,
+        metrics: Optional[List[Metric]] = None,
+        *,
+        template: Optional[Template] = None,
+        use_tqdm: Optional[bool] = None,
+        adapter_request: Optional[AdapterRequest] = None,
+    ):
+        url = f'http://{self.host}:{self.server_port}/infer/'
+        response = self.session.post(
+            url,
+            json={
+                'infer_requests': infer_requests,
+                'request_config': request_config,
+                'metrics': metrics,
+                'template': template,
+                'use_tqdm': use_tqdm,
+                'adapter_request': adapter_request,
+            },
+        )
+        if response.status_code == 200:
+            return [from_dict(data_class=ChatCompletionResponse, data=resp) for resp in response.json()]
+        else:
+            raise Exception(f'Request failed: {response.status_code}, {response.text}')
+
+    def init_communicator(self):
+        """
+        Initializes the weight update group in a distributed setup for model synchronization.
+        """
+        # Get the tensor parallel size from the server
+        url = f'http://{self.host}:{self.server_port}/get_world_size/'
+        response = requests.get(url)
+        if response.status_code == 200:
+            vllm_world_size = response.json()['world_size']
+        else:
+            raise Exception(f'Request failed: {response.status_code}, {response.text}')
+
+        world_size = vllm_world_size + 1  # add the client to the world
+        self.rank = vllm_world_size  # the client's rank is the last process
+
+        # Initialize weight update group
+        url = f'http://{self.host}:{self.server_port}/init_communicator/'
+        # In the server side, the host is set to 0.0.0.0
+        response = self.session.post(url, json={'host': '0.0.0.0', 'port': self.group_port, 'world_size': world_size})
+        if response.status_code != 200:
+            raise Exception(f'Request failed: {response.status_code}, {response.text}')
+
+        # Brief delay to allow server initialization. While not strictly required (client socket will retry on
+        # connection failure), this prevents log warnings like:
+        # [W416 23:24:57.460001114 socket.cpp:204] [c10d] The hostname of the client socket cannot be retrieved. err=-3
+        time.sleep(0.1)
+
+        # Set up the communication group for weight broadcasting
+        pg = StatelessProcessGroup.create(host=self.host, port=self.group_port, rank=self.rank, world_size=world_size)
+        self.pynccl_comm = PyNcclCommunicator(pg, device=0)
+
+        # When the client object is deleted, close the weight update group
+        atexit.register(self.close_communicator)
+
+    def update_named_param(self, name: str, weights: torch.Tensor):
+        """
+        Updates a specific named parameter in the model and broadcasts it to other processes.
+
+        Args:
+            name (`str`):
+                Name of the layer whose weights are being updated.
+            weights (`torch.Tensor`):
+                Tensor containing the updated weights.
+        """
+        dtype, shape = str(weights.dtype), tuple(weights.shape)
+        url = f'http://{self.host}:{self.server_port}/update_named_param/'
+        response = self.session.post(url, json={'name': name, 'dtype': dtype, 'shape': shape})
+        if response.status_code != 200:
+            raise Exception(f'Request failed: {response.status_code}, {response.text}')
+
+        # Broadcast the weights to the other processes
+        self.pynccl_comm.broadcast(weights, src=self.rank)
+        self.pynccl_comm.group.barrier()
+
+    def update_model_params(self, model: nn.Module):
+        """
+        Updates all parameters of the given model by calling `update_named_param` for each parameter in the model.
+
+        Args:
+            model (`nn.Module`):
+                Model whose parameters (weights/biases) are to be updated.
+        """
+        for name, param in model.named_parameters():
+            # Update each parameter individually
+            self.update_named_param(name, param.data)
+
+    def reset_prefix_cache(self):
+        """
+        Resets the prefix cache for the model.
+        """
+        url = f'http://{self.host}:{self.server_port}/reset_prefix_cache/'
+        response = self.session.post(url)
+        if response.status_code != 200:
+            raise Exception(f'Request failed: {response.status_code}, {response.text}')
+
+    def close_communicator(self):
+        """
+        Closes the weight update group and cleans up the communication group.
+        """
+        url = f'http://{self.host}:{self.server_port}/close_communicator/'
+
+        try:
+            response = self.session.post(url)
+        except ConnectionError:
+            # The server might be already down, so we don't need to close the communicator
+            pass
+        else:
+            if response.status_code != 200:
+                raise Exception(f'Request failed: {response.status_code}, {response.text}')

ms-swift/swift/trainers/sequence_parallel/base.py

@@ -0,0 +1,45 @@
+import abc
+from abc import abstractmethod
+
+
+class SequenceParallel(abc.ABC):
+
+    @abstractmethod
+    def init_sequence_parallel(self, size):
+        pass
+
+    @abstractmethod
+    def prepare_model(self, model, tokenizer, split_in_forward):
+        pass
+
+    @abstractmethod
+    def pad_and_split_inputs(self,
+                             tokenizer,
+                             input_ids,
+                             input_embeds,
+                             labels,
+                             position_ids,
+                             attention_mask,
+                             loss_scale,
+                             embed_tokens=None):
+        pass
+
+    @abstractmethod
+    def reduce_outputs(self, loss, labels):
+        pass
+
+    @property
+    def sp_group(self):
+        return None
+
+    @abstractmethod
+    def world_size(self):
+        pass
+
+    @abstractmethod
+    def prepare_trainer(self, trainer):
+        pass
+
+    @abstractmethod
+    def get_dataloader(self, trainer, dataset, batch_size):
+        pass

ms-swift/swift/trainers/sequence_parallel/ulysses.py

@@ -0,0 +1,594 @@
+import math
+from functools import partial
+from types import MethodType
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+import datasets
+import numpy as np
+import torch
+import torch.distributed as dist
+from peft import PeftModel
+from torch.distributed.device_mesh import init_device_mesh
+from torch.nn import CrossEntropyLoss
+from torch.utils.data import DataLoader, Sampler
+from transformers.trainer_utils import seed_worker
+
+from swift.llm import DataLoaderDispatcher, get_model_arch
+from swift.tuners import SwiftModel
+from swift.utils import get_current_device, get_device, get_dist_setting
+from .base import SequenceParallel
+
+
+class GatherLoss(torch.autograd.Function):
+    """Gather loss from sequence group"""
+
+    @staticmethod
+    def forward(ctx, loss, labels, process_group, gather_idx=None):
+        """
+        Args:
+            loss: loss tensor after splitting
+            labels: labels tensor after splitting
+            process_group: the sequence parallel group
+            gather_idx: gather the tensors on this dim
+        """
+        ctx.process_group = process_group
+        shape0 = labels.shape[0]
+        ctx.scatter_shape = labels.shape[gather_idx or 0]
+        ctx.gather_idx = gather_idx or 0
+        world_size = dist.get_world_size(group=process_group)  # the sp world size
+        output = torch.empty((shape0 * world_size, *loss.shape[1:]), dtype=loss.dtype, device=loss.device)
+        # gather all from sp group
+        dist.all_gather_into_tensor(output, loss, group=process_group)
+        if gather_idx is not None:
+            output = torch.cat(output.split(shape0, dim=0), dim=gather_idx)
+        labels_output = torch.empty((shape0 * world_size, *labels.shape[1:]), dtype=labels.dtype, device=labels.device)
+        dist.all_gather_into_tensor(labels_output, labels, group=process_group)
+        if gather_idx is not None:
+            labels_output = torch.cat(labels_output.split(shape0, dim=0), dim=gather_idx)
+        return output, labels_output
+
+    @staticmethod
+    def backward(ctx, *grad_output):
+        _grad = grad_output[0] * dist.get_world_size(group=ctx.process_group)
+        return _grad.split(
+            ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(ctx.process_group)].contiguous(), None, None, None
+
+
+# For nll loss
+def loss_scale_sp_func(outputs, labels, loss_scale=None, num_items_in_batch=None, process_group=None) -> torch.Tensor:
+    if hasattr(outputs, 'logits'):
+        logits = outputs.logits
+    else:
+        logits = outputs
+    device = logits.device
+    logits = logits.view(-1, logits.shape[-1])
+    labels = labels.flatten().to(device)
+    # Flatten the tokens
+    loss_fct = CrossEntropyLoss(reduction='none')
+    # flatten loss
+    loss = loss_fct(logits, labels)
+
+    if loss_scale is not None:
+        loss_scale = loss_scale.flatten().to(loss.device)
+        loss = (loss_scale * loss)
+    loss, labels = GatherLoss.apply(loss, labels, process_group)
+    loss = loss[labels != -100].sum()
+    if num_items_in_batch is None:
+        loss = loss / (labels != -100).sum()
+    else:
+        loss = loss / num_items_in_batch
+    return loss
+
+
+# For DPO
+def get_batch_logps(logits: torch.FloatTensor,
+                    labels: torch.LongTensor,
+                    label_pad_token_id: int = -100,
+                    is_encoder_decoder: bool = False,
+                    process_group=None) -> Tuple[torch.FloatTensor, torch.LongTensor]:
+    labels = labels.clone()  # No need to shift, pad and split has shifted the inputs.
+    loss_mask = labels != label_pad_token_id
+    labels[labels == label_pad_token_id] = 0
+    labels = labels.to(logits.device)
+    loss_mask = loss_mask.to(logits.device)
+    per_token_logps = torch.gather(logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)
+    total_per_token_logps, total_loss_mask = GatherLoss.apply(per_token_logps, loss_mask, process_group, 1)
+    return (total_per_token_logps * total_loss_mask).sum(-1), total_loss_mask.sum(-1)
+
+
+class UlyssesSampler(Sampler):
+
+    # Code borrowed from mmengine
+    def __init__(self, ulysses, dataset, shuffle: bool = True, seed=None, round_up: bool = True) -> None:
+        self.ulysses = ulysses
+        rank = dist.get_rank(ulysses.device_mesh['data'].get_group())
+        world_size = ulysses.device_mesh['data'].size()
+        self.rank = rank
+        self.world_size = world_size
+
+        self.dataset = dataset
+        self.shuffle = shuffle
+        assert seed is not None
+        self.seed = seed
+        self.epoch = 0
+        self.round_up = round_up
+
+        if self.round_up:
+            self.num_samples = math.ceil(len(self.dataset) / world_size)
+            self.total_size = self.num_samples * self.world_size
+        else:
+            self.num_samples = math.ceil((len(self.dataset) - rank) / world_size)
+            self.total_size = len(self.dataset)
+
+    def __iter__(self) -> Iterator[int]:
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = torch.arange(len(self.dataset)).tolist()
+
+        if self.round_up:
+            indices = (indices * int(self.total_size / len(indices) + 1))[:self.total_size]
+
+        indices = indices[self.rank:self.total_size:self.world_size]
+
+        return iter(indices)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class UlyssesDispatcher(DataLoaderDispatcher):
+
+    def __init__(self, base_dataloader, ulysses):
+        super().__init__(base_dataloader)
+        self.ulysses = ulysses
+
+    def __iter__(self):
+        base_iter = iter(self.base_dataloader)
+        while True:
+            data = None
+            try:
+                for i in range(self.ulysses.dp_world_size):
+                    data = next(base_iter)
+                    if i == self.ulysses.dp_rank:
+                        break
+            except StopIteration:
+                pass
+            if data is None:
+                break
+            yield data
+
+
+# Code borrowed from deepspeed, here is why:
+# 1. Reduce the dependency
+# 2. The original code is complex
+def _generate_layout_params(scatter_idx, seq_world_size, input):
+    if scatter_idx < 2:
+        bs, global_seq_len, num_local_head, head_dim = input.shape
+        pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim]
+        pre_all2all_permute_idx = (1, 0, 2, 3, 4)
+
+        post_all2all_permute_idx = (1, 2, 0, 3, 4)
+        post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim]
+    else:
+        bs, local_seq_len, num_total_head, head_dim = input.shape
+        assert num_total_head % seq_world_size == 0, (f'Number of heads ({num_total_head}) must be divisible '
+                                                      f'by the sequence parallel size ({seq_world_size})!')
+        pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim]
+        pre_all2all_permute_idx = (2, 0, 1, 3, 4)
+
+        post_all2all_permute_idx = (1, 0, 2, 3, 4)
+        post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim]
+
+    return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape
+
+
+def post_all2all(permute_idx, res_shape):
+    """
+    Post-processing function for `all2all` communication.
+    """
+
+    def post_func(input):
+        if permute_idx is not None:
+            input = input.permute(permute_idx).contiguous()
+        output = input.reshape(res_shape).contiguous()
+
+        return output
+
+    return post_func
+
+
+def pre_all2all_fun(permute_idx, inp_shape, input):
+    """
+    Pre-processing function for `all2all` communication.
+    """
+    input_t = input.reshape(inp_shape).contiguous()
+    if permute_idx is not None:
+        input_t = input_t.permute(permute_idx).contiguous()
+    return input_t
+
+
+def single_all_to_all(input, scatter_idx, gather_idx, group, **kwargs):
+    seq_world_size = dist.get_world_size(group)
+    num_heads = input.shape[2]
+    if num_heads % seq_world_size != 0 and not scatter_idx < 2:
+        raise NotImplementedError
+    pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = (
+        _generate_layout_params(scatter_idx, seq_world_size, input))
+
+    input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input)
+
+    post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape)
+    output = torch.empty_like(input_t)
+    dist.all_to_all_single(output, input_t, group=group)
+
+    res = post_all2all_fun(output)
+    return res
+
+
+class _SeqAllToAll(torch.autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        input: torch.Tensor,
+        scatter_idx: int,
+        gather_idx: int,
+    ) -> torch.Tensor:
+        ctx.group = group
+        ctx.scatter_idx = scatter_idx
+        ctx.gather_idx = gather_idx
+        res = single_all_to_all(input, scatter_idx, gather_idx, group)
+        return res
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]:
+        return None, _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), None, None
+
+
+class DistributedAttention(torch.nn.Module):
+
+    def __init__(
+        self,
+        local_attention,
+        sequence_process_group: dist.ProcessGroup,
+        scatter_idx: int = 2,
+        gather_idx: int = 1,
+    ) -> None:
+        super(DistributedAttention, self).__init__()
+        self.local_attn = local_attention
+        self.spg = sequence_process_group
+        self.scatter_idx = scatter_idx
+        self.gather_idx = gather_idx
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor,
+                *args: Any, **kwargs) -> torch.Tensor:
+        query_layer = _SeqAllToAll.apply(self.spg, query, self.scatter_idx, self.gather_idx)
+        key_layer = _SeqAllToAll.apply(self.spg, key, self.scatter_idx, self.gather_idx)
+        value_layer = _SeqAllToAll.apply(self.spg, value, self.scatter_idx, self.gather_idx)
+        position_ids = kwargs.pop('position_ids', None)
+        if position_ids is not None:
+            shape0 = position_ids.shape[0]
+            position_ids_output = torch.empty((shape0 * dist.get_world_size(self.spg), position_ids.shape[1]),
+                                              dtype=position_ids.dtype,
+                                              device=position_ids.device)
+            dist.all_gather_into_tensor(position_ids_output, position_ids, group=self.spg)
+            position_ids = torch.cat(position_ids_output.split(shape0, dim=0), dim=1)
+        context_layer = self.local_attn(
+            query_layer, key_layer, value_layer, attention_mask, *args, position_ids=position_ids, **kwargs)
+        output = _SeqAllToAll.apply(self.spg, context_layer, self.gather_idx, self.scatter_idx)
+        return output
+
+
+class Ulysses(SequenceParallel):
+
+    def __init__(self):
+        self.split_in_forward = None
+        self.dp_world_size = None
+        self.sp_world_size = None
+        self.model_dtype = None
+        self.causal_mask_func = None
+        self.device_mesh = None
+        self._inited = False
+
+    def init_sequence_parallel(self, size):
+        if self._inited:
+            return
+        self._inited = True
+        self.sp_world_size = size
+        rank, local_rank, world_size, local_world_size = get_dist_setting()
+        self.dp_world_size = world_size // size
+        self.device_mesh = init_device_mesh(
+            get_device().split(':')[0], mesh_shape=(world_size // size, size), mesh_dim_names=['data', 'sequence'])
+
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+        ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2']
+        ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa']
+
+        def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
+                             dist_attn, **kwargs):
+            if dist_attn.local_attn is None:
+
+                def _attention(query, key, value, *args, **kwargs):
+                    query = query.transpose(1, 2)
+                    key = key.transpose(1, 2)
+                    value = value.transpose(1, 2)
+                    return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args,
+                                                                               **kwargs)[0]
+
+                dist_attn.local_attn = _attention
+
+            return dist_attn(
+                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
+                *args, **kwargs), None
+
+        def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
+                            dist_attn, **kwargs):
+            if dist_attn.local_attn is None:
+
+                def _attention(query, key, value, *args, **kwargs):
+                    query = query.transpose(1, 2)
+                    key = key.transpose(1, 2)
+                    value = value.transpose(1, 2)
+                    return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0]
+
+                dist_attn.local_attn = _attention
+            return dist_attn(
+                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
+                *args, **kwargs), None
+
+        ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = partial(
+            local_flash_attn, dist_attn=DistributedAttention(None, self.sp_group))
+        ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self.sp_group))
+
+        from transformers.modeling_flash_attention_utils import is_flash_attn_available
+        if is_flash_attn_available():
+            # TODO this works for multi-modal models like qwen2.5-vl
+            # SDPA is not supported, because we need to copy the code to our project, which will bring
+            # more works for maintaining.
+            from transformers import modeling_flash_attention_utils
+            from transformers.modeling_flash_attention_utils import _flash_attention_forward
+            _distributed_flash_attention = DistributedAttention(_flash_attention_forward, self.sp_group)
+
+            def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor,
+                                        value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len,
+                                        *args, **kwargs):
+                return _distributed_flash_attention(query_states, key_states, value_states, attention_mask,
+                                                    q_len * self.sp_world_size, *args, **kwargs)
+
+            modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward
+
+    def prepare_model(self, model, tokenizer, split_in_forward):
+        self.split_in_forward = split_in_forward
+
+        def forward(_self, **kwargs):
+            # Split embedding here for multi-modal
+            inputs_embeds = kwargs['inputs_embeds']
+            position_ids = kwargs['position_ids']
+            attention_mask = kwargs['attention_mask']
+            _, inputs_embeds, _, position_ids, attention_mask, _ = self.pad_and_split_inputs(
+                tokenizer,
+                None,
+                inputs_embeds,
+                None,
+                position_ids,
+                attention_mask,
+                None,
+                embed_tokens=_self.embed_tokens)
+            kwargs['inputs_embeds'] = inputs_embeds
+            kwargs['position_ids'] = position_ids
+            kwargs['attention_mask'] = attention_mask
+            return _self.forward_origin(**kwargs)
+
+        if isinstance(model, (SwiftModel, PeftModel)):
+            model = model.model
+        model_meta = model.model_meta
+        llm_prefix = getattr(get_model_arch(model_meta.model_arch), 'language_model', None)
+        if llm_prefix:
+            llm_model = getattr(model, llm_prefix[0])
+        else:
+            llm_model = model
+
+        if 'CausalLM' not in llm_model.__class__.__name__:
+            llm_model = model
+
+        base_model = llm_model.model
+        self.causal_mask_func = base_model._update_causal_mask
+        if self.split_in_forward:
+            # for multi modal models
+            base_model.forward_origin = base_model.forward
+            base_model.forward = MethodType(forward, base_model)
+
+        self.model_dtype = next(model.parameters()).dtype
+
+    def _pad_sp(self, tensor, padding_value, dim=-1):
+        # code borrowed from xtuner
+        length = tensor.shape[dim]
+        if length % self.sp_world_size == 0:
+            return tensor
+
+        pad_num = self.sp_world_size - (length % self.sp_world_size)
+        if not isinstance(padding_value, torch.Tensor):
+            # ids
+            pad_shape = ((*tensor.shape[:dim], pad_num, *tensor.shape[dim + 1:]) if dim != -1 else
+                         (*tensor.shape[:dim], pad_num))
+            pad = torch.full(pad_shape, padding_value, dtype=tensor.dtype, device=tensor.device)
+            tensor = torch.cat([tensor, pad], dim=dim)
+        else:
+            # For embeddings
+            tensor = torch.cat([tensor, padding_value.unsqueeze(0).repeat(tensor.shape[0], pad_num, 1)], dim=dim)
+        return tensor
+
+    def world_size(self):
+        return self.sp_world_size
+
+    def _split_sp(self, input, dim: int, sp_group: dist.ProcessGroup):
+        # code borrowed from xtuner
+        if self.sp_world_size == 1:
+            return input
+
+        rank = dist.get_rank(sp_group)
+        dim_size = input.size(dim)
+        assert dim_size % self.sp_world_size == 0, (f'The dimension to split ({dim_size}) is not a multiple of '
+                                                    f'world size ({self.sp_world_size}), cannot split tensor evenly')
+
+        tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim)
+        output = tensor_list[rank].contiguous()
+
+        return output
+
+    def pad_and_split_inputs(self,
+                             tokenizer,
+                             input_ids,
+                             input_embeds,
+                             labels,
+                             position_ids,
+                             attention_mask,
+                             loss_scale,
+                             embed_tokens=None):
+        sp_group = self.sp_group
+        split_inputs = False
+        if (input_ids is not None and not self.split_in_forward) or input_embeds is not None:
+            # Whether split the model inputs
+            # cannot split input_ids for multi-modal models
+            split_inputs = True
+        if input_ids is not None and split_inputs:
+            input_ids = self._pad_sp(input_ids, padding_value=tokenizer.pad_token_id, dim=-1)
+        if input_embeds is not None:
+            pad_emb = embed_tokens(torch.tensor(tokenizer.pad_token_id).to(embed_tokens.weight.device)).unsqueeze(0)
+            input_embeds = self._pad_sp(input_embeds, padding_value=pad_emb, dim=1)
+        if position_ids is not None and split_inputs:
+            position_ids = self._pad_sp(position_ids, padding_value=0, dim=-1)
+        if split_inputs:
+            inputs = input_ids if input_ids is not None else input_embeds
+            attn_shape = inputs.shape[1]  # The sequence length
+            if attention_mask is None:
+                attention_mask = torch.ones_like(position_ids)
+            attention_mask = self._pad_sp(attention_mask, padding_value=0, dim=-1)
+            cache_position = torch.arange(0, attn_shape, device=inputs.device)
+            # pad attention mask to 4d to avoid calculation errors
+            attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), cache_position, None,
+                                                   None)
+        if input_ids is not None and split_inputs:
+            input_ids = self._split_sp(input_ids, dim=1, sp_group=sp_group)
+        if input_embeds is not None:
+            input_embeds = self._split_sp(input_embeds, dim=1, sp_group=sp_group)
+        if position_ids is not None and split_inputs:
+            position_ids = self._split_sp(position_ids, dim=-1, sp_group=sp_group)
+        if labels is not None:
+            labels = self._pad_sp(labels, padding_value=-100, dim=-1)
+            labels[:, 0] = -100  # make the last invalid, so we do not need to cut the loss of last token
+            labels = torch.roll(labels, shifts=-1, dims=1)
+            labels = self._split_sp(labels, dim=1, sp_group=sp_group)
+
+        if loss_scale is not None:
+            loss_scale = self._pad_sp(loss_scale, padding_value=0., dim=-1)
+            loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1)
+            loss_scale = self._split_sp(loss_scale, dim=-1, sp_group=sp_group)
+
+        return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale
+
+    def reduce_outputs(self, loss, labels):
+        return loss
+
+    @property
+    def sp_rank(self):
+        return dist.get_rank(self.device_mesh['sequence'].get_group())
+
+    @property
+    def dp_rank(self):
+        return dist.get_rank(self.device_mesh['data'].get_group())
+
+    @property
+    def sp_group(self):
+        return self.device_mesh['sequence'].get_group()
+
+    @property
+    def dp_group(self):
+        return self.device_mesh['data'].get_group()
+
+    def get_dataloader(self, trainer, dataset, batch_size):
+        data_collator = trainer.data_collator
+        if isinstance(dataset, datasets.Dataset):
+            dataset = trainer._remove_unused_columns(dataset, description='training')
+        else:
+            data_collator = trainer._get_collator_with_removed_columns(data_collator, description='training')
+        if hasattr(dataset, '__len__'):
+            sampler = UlyssesSampler(self, dataset, seed=42)
+            dataloader_params = {
+                'batch_size': batch_size,
+                'collate_fn': data_collator,
+                'num_workers': trainer.args.dataloader_num_workers,
+                'pin_memory': trainer.args.dataloader_pin_memory,
+                'persistent_workers': trainer.args.dataloader_persistent_workers,
+            }
+
+            if not isinstance(dataset, torch.utils.data.IterableDataset):
+                dataloader_params['sampler'] = sampler
+                dataloader_params['drop_last'] = trainer.args.dataloader_drop_last
+                dataloader_params['worker_init_fn'] = seed_worker
+
+            return DataLoader(dataset, **dataloader_params)
+        else:
+            dataloader_params = {
+                'collate_fn': data_collator,
+                'num_workers': trainer.args.dataloader_num_workers,
+                'pin_memory': trainer.args.dataloader_pin_memory,
+                'persistent_workers': trainer.args.dataloader_persistent_workers,
+                'prefetch_factor': trainer.args.dataloader_prefetch_factor
+            }
+            if dist.is_initialized() and dataloader_params['prefetch_factor']:
+                dataloader_params['prefetch_factor'] = dataloader_params['prefetch_factor'] * dist.get_world_size()
+            dataloader = DataLoader(dataset, batch_size=batch_size, **dataloader_params)
+            dataloader = UlyssesDispatcher(dataloader, self)
+            return dataloader
+
+    def prepare_trainer(self, trainer):
+        if trainer.train_dataset is None:
+            raise ValueError('Trainer: training requires a train_dataset.')
+
+        trainer.compute_loss_func = partial(loss_scale_sp_func, process_group=self.sp_group)
+        if hasattr(trainer, 'get_batch_logps'):
+            trainer.get_batch_logps = partial(get_batch_logps, process_group=self.sp_group)
+        if hasattr(trainer, 'get_nll_loss'):
+
+            def rlhf_loss_scale_sp_func(_, *args, **kwargs):
+                return loss_scale_sp_func(*args, process_group=self.sp_group, **kwargs)
+
+            trainer.get_nll_loss = MethodType(rlhf_loss_scale_sp_func, trainer)
+
+        from swift.plugin import metric
+        from swift.trainers import mixin
+        compute_acc_origin = metric.compute_acc
+
+        def compute_acc(preds, labels, *args, **kwargs) -> Dict[str, List[float]]:
+
+            # Gather preds and labels across the sp group
+            if isinstance(preds, np.ndarray):
+                preds = torch.from_numpy(preds).to(get_current_device())
+            if isinstance(labels, np.ndarray):
+                labels = torch.from_numpy(labels).to(get_current_device())
+            shape0 = preds.shape[0]
+            preds_output = torch.empty((shape0 * self.sp_world_size, preds.shape[1]),
+                                       dtype=preds.dtype,
+                                       device=preds.device)
+            dist.all_gather_into_tensor(preds_output, preds, group=self.sp_group)
+            preds_output = torch.cat(preds_output.split(shape0, dim=0), dim=1)
+            shape0 = labels.shape[0]
+            labels_output = torch.empty((shape0 * self.sp_world_size, labels.shape[1]),
+                                        dtype=labels.dtype,
+                                        device=labels.device)
+            dist.all_gather_into_tensor(labels_output, labels, group=self.sp_group)
+            labels_output = torch.cat(labels_output.split(shape0, dim=0), dim=1)
+            # roll back to fit compute_acc
+            labels_output = torch.roll(labels_output, shifts=1, dims=1)
+            return compute_acc_origin(preds_output, labels_output, *args, **kwargs)
+
+        metric.compute_acc = compute_acc
+        mixin.compute_acc = compute_acc

ms-swift/swift/trainers/sequence_parallel/xtuner.py

@@ -0,0 +1,127 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any
+
+import datasets
+import torch
+import torch.distributed as dist
+from datasets import Dataset
+from torch.utils.data import DataLoader
+from transformers.trainer_utils import seed_worker
+
+from .base import SequenceParallel
+
+
+class XTuner(SequenceParallel):
+
+    @staticmethod
+    def assert_xtuner_runtime_condition():
+        from swift.utils import is_xtuner_available
+        assert is_xtuner_available(), \
+            ('Please install XTuner first to pack dataset to `max_length`.'
+             '`pip install -U \'xtuner[deepspeed]\'`')
+        assert dist.is_initialized(), 'pack_to_max_length is only available with distributed training.'
+
+    def pack_dataset_xtuner(self, dataset: Dataset, args: Any) -> Any:
+        self.assert_xtuner_runtime_condition()
+        if dist.get_rank() == 0:
+            ds = [i[0] for i in dataset.data]
+            train_dataset = Dataset.from_list(ds)
+            from xtuner.dataset.huggingface import pack_dataset
+            train_dataset = pack_dataset(
+                train_dataset,
+                max_length=args.max_length,
+                use_varlen_attn=False,
+                shuffle_before_pack=True,
+                map_num_proc=16)
+            objects = [train_dataset]
+            train_dataset.save_to_disk('alpaca_pack')
+        else:
+            objects = [None]
+        dist.broadcast_object_list(objects, src=0)
+        train_dataset = objects[0]
+        return train_dataset
+
+    @property
+    def sp_group(self):
+        from xtuner.parallel.sequence import get_sequence_parallel_group
+        return get_sequence_parallel_group()
+
+    def init_sequence_parallel(self, size):
+        self.assert_xtuner_runtime_condition()
+        from xtuner.parallel.sequence import init_sequence_parallel
+        init_sequence_parallel(size)
+
+    def prepare_model(self, model, tokenizer, split_in_forward):
+        self.assert_xtuner_runtime_condition()
+        from xtuner.model.modules.dispatch import dispatch_modules
+        dispatch_modules(model)
+
+    def pad_and_split_inputs(self,
+                             tokenizer,
+                             input_ids,
+                             input_embeds,
+                             labels,
+                             position_ids,
+                             attention_mask,
+                             loss_scale,
+                             embed_tokens=None):
+        self.assert_xtuner_runtime_condition()
+        from xtuner.parallel.sequence import (pad_for_sequence_parallel, split_for_sequence_parallel,
+                                              get_sequence_parallel_group)
+        input_ids = pad_for_sequence_parallel(input_ids, padding_value=tokenizer.pad_token_id, dim=-1)
+        labels = pad_for_sequence_parallel(labels, padding_value=-100, dim=-1)
+        position_ids = pad_for_sequence_parallel(position_ids, padding_value=0, dim=-1)
+        if attention_mask is not None:
+            attention_mask = pad_for_sequence_parallel(attention_mask, padding_value=0, dim=-1)
+
+        sp_group = get_sequence_parallel_group()
+        input_ids = split_for_sequence_parallel(input_ids, dim=1, sp_group=sp_group)
+        labels = split_for_sequence_parallel(labels, dim=1, sp_group=sp_group)
+        position_ids = split_for_sequence_parallel(position_ids, dim=1, sp_group=sp_group)
+        if attention_mask is not None:
+            attention_mask = split_for_sequence_parallel(attention_mask, dim=-1, sp_group=sp_group)
+        if loss_scale is not None:
+            loss_scale = pad_for_sequence_parallel(loss_scale, padding_value=0., dim=-1)
+            loss_scale = split_for_sequence_parallel(loss_scale, dim=1, sp_group=sp_group)
+
+        return input_ids, None, labels, position_ids, attention_mask, loss_scale
+
+    def reduce_outputs(self, loss, labels):
+        from xtuner.parallel.sequence import (reduce_sequence_parallel_loss, get_sequence_parallel_group)
+        # reduce loss for logging correctly
+        num_tokens = (labels != -100).sum()
+        return reduce_sequence_parallel_loss(loss, num_tokens, get_sequence_parallel_group())
+
+    def world_size(self):
+        self.assert_xtuner_runtime_condition()
+        from xtuner.parallel.sequence import get_sequence_parallel_world_size
+        return get_sequence_parallel_world_size()
+
+    def prepare_trainer(self, trainer):
+        pass
+
+    def get_dataloader(self, trainer, dataset, batch_size):
+        # modified from HFTrainer.get_train_dataloader
+        # RandomSampler -> SequenceParallelSampler
+        self.assert_xtuner_runtime_condition()
+        data_collator = trainer.data_collator
+        if isinstance(dataset, datasets.Dataset):
+            dataset = trainer._remove_unused_columns(dataset, description='training')
+        else:
+            data_collator = trainer._get_collator_with_removed_columns(data_collator, description='training')
+
+        dataloader_params = {
+            'batch_size': batch_size,
+            'collate_fn': data_collator,
+            'num_workers': trainer.args.dataloader_num_workers,
+            'pin_memory': trainer.args.dataloader_pin_memory,
+            'persistent_workers': trainer.args.dataloader_persistent_workers,
+        }
+
+        if not isinstance(dataset, torch.utils.data.IterableDataset):
+            from xtuner.parallel import SequenceParallelSampler
+            dataloader_params['sampler'] = SequenceParallelSampler(dataset, seed=1024)
+            dataloader_params['drop_last'] = trainer.args.dataloader_drop_last
+            dataloader_params['worker_init_fn'] = seed_worker
+
+        return DataLoader(dataset, **dataloader_params)

ms-swift/swift/trainers/torchacc_mixin.py

@@ -0,0 +1,156 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import shutil
+from typing import Optional
+
+from transformers import PreTrainedModel, is_datasets_available
+
+from swift.utils import use_torchacc
+from swift.utils.torchacc_utils import (patch_clip_grad_norm, save_ta_ddp_checkpoint, save_ta_fsdp_checkpoint,
+                                        ta_eval_dataloader, ta_load_optimizer_and_scheduler,
+                                        ta_save_optimizer_and_scheduler, ta_test_dataloader, ta_train_dataloader,
+                                        ta_trim_graph)
+
+
+class TorchAccMixin:
+
+    def __init__(self, *args, **kwargs):
+        if use_torchacc():
+            patch_clip_grad_norm(self.accelerator)
+        super().__init__(*args, **kwargs)
+
+    def get_train_dataloader(self):
+        if not use_torchacc():
+            return super().get_train_dataloader()
+
+        if is_datasets_available():
+            import datasets
+
+        if self.train_dataset is None:
+            raise ValueError('Trainer: training requires a train_dataset.')
+
+        train_dataset = self.train_dataset
+        data_collator = self.data_collator
+
+        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
+            train_dataset = self._remove_unused_columns(train_dataset, description='training')
+        else:
+            data_collator = self._get_collator_with_removed_columns(data_collator, description='training')
+
+        return ta_train_dataloader(train_dataset, data_collator, self._get_train_sampler(), self.args,
+                                   self._train_batch_size)
+
+    def get_eval_dataloader(self, eval_dataset=None):
+
+        if not use_torchacc():
+            return super().get_eval_dataloader(eval_dataset)
+
+        if is_datasets_available():
+            import datasets
+
+        if eval_dataset is None and self.eval_dataset is None:
+            raise ValueError('Trainer: evaluation requires an eval_dataset.')
+        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
+        data_collator = self.data_collator
+
+        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
+            eval_dataset = self._remove_unused_columns(eval_dataset, description='evaluation')
+        else:
+            data_collator = self._get_collator_with_removed_columns(data_collator, description='evaluation')
+
+        return ta_eval_dataloader(eval_dataset, data_collator, self._get_eval_sampler(eval_dataset), self.args)
+
+    def get_test_dataloader(self, test_dataset):
+
+        if not use_torchacc():
+            return super().get_test_dataloader(test_dataset)
+
+        if is_datasets_available():
+            import datasets
+
+        data_collator = self.data_collator
+
+        if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
+            test_dataset = self._remove_unused_columns(test_dataset, description='test')
+        else:
+            data_collator = self._get_collator_with_removed_columns(data_collator, description='test')
+
+        return ta_test_dataloader(test_dataset, data_collator, self._get_eval_sampler(test_dataset), self.args)
+
+    def _save_tpu(self, output_dir: Optional[str] = None):
+
+        if not use_torchacc():
+            return super()._save_tpu(output_dir)
+
+        import torch_xla.core.xla_model as xm
+
+        # Compatible with swift and peft
+        output_dir = output_dir if output_dir is not None else self.args.output_dir
+
+        if xm.is_master_ordinal(local=False):
+            os.makedirs(output_dir, exist_ok=True)
+            # configuration.json
+            model_dir = getattr(self.model, 'model_dir', None)
+            if model_dir is not None:
+                src_path = os.path.join(model_dir, 'configuration.json')
+                dst_path = os.path.join(output_dir, 'configuration.json')
+                if os.path.exists(src_path):
+                    shutil.copy(src_path, dst_path)
+            else:
+                self._create_configuration_file(self.model, output_dir)
+            self._save_sft_args(output_dir)
+            # generation_config
+            generation_config = getattr(self.args, 'generation_config', None)
+            if generation_config is not None:
+                generation_config.save_pretrained(output_dir)
+
+        # model
+        if self.args.fsdp_num > 1:
+            save_ta_fsdp_checkpoint(self.model, self.tokenizer, self.args, output_dir)
+        else:
+            save_ta_ddp_checkpoint(self.model, self.tokenizer, self.args, output_dir)
+
+        # additional files
+        if xm.is_master_ordinal(local=False):
+            if self.args is not None and self.args.sft_type == 'full':
+                additional_files = getattr(self.args, 'additional_saved_files',
+                                           None) or [] + ['preprocessor_config.json']
+                if model_dir is not None:
+                    for file in additional_files:
+                        src_path = os.path.join(model_dir, file)
+                        dst_path = os.path.join(output_dir, file)
+                        if os.path.isfile(src_path):
+                            shutil.copy(src_path, dst_path)
+                        elif os.path.isdir(src_path):
+                            shutil.copytree(src_path, dst_path)
+
+    def _load_optimizer_and_scheduler(self, checkpoint):
+
+        if not use_torchacc() or self.args.fsdp_num == 1:
+            return super()._load_optimizer_and_scheduler(checkpoint)
+
+        self.optimizer, self.lr_scheduler = ta_load_optimizer_and_scheduler(self.optimizer, self.lr_scheduler,
+                                                                            checkpoint, self.args.device)
+
+    def _save_optimizer_and_scheduler(self, output_dir):
+        if not use_torchacc() or not self.args.fsdp_num == 1:
+            return super()._save_optimizer_and_scheduler(output_dir)
+
+        return ta_save_optimizer_and_scheduler(self.optimizer, self.lr_scheduler, output_dir)
+
+    def _maybe_log_save_evaluate(self, tr_loss, *args, **kwargs):
+        if use_torchacc() and self.control.should_log:
+            ta_trim_graph()
+        super()._maybe_log_save_evaluate(tr_loss, *args, **kwargs)
+
+    def _load_from_checkpoint(self, resume_from_checkpoint: str, model=None) -> None:
+        if use_torchacc():
+            if model is None:
+                model = self.model
+            # Loading checkpoint of TorchAcc has been done in tuner.py when
+            # sft_type is 'full'.
+            if self.args.fsdp_num > 1:
+                model = model._get_underlay_model().module.module
+            if isinstance(model, PreTrainedModel):
+                return
+        return super()._load_from_checkpoint(resume_from_checkpoint, model)