chore: import upstream snapshot with attribution

swift/utils/__init__.py

@@ -0,0 +1,91 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from .import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .env import (get_dist_setting, get_hf_endpoint, get_node_setting, get_pai_tensorboard_dir,
+                      is_deepspeed_enabled, is_dist, is_last_rank, is_local_master, is_master, is_mp, is_mp_ddp,
+                      is_pai_training_job, select_device, use_hf_hub)
+    from .hf_config import HfConfigFactory
+    from .hub_utils import download_file, download_ms_file, git_clone_github, patch_kernels, safe_snapshot_download
+    from .import_utils import (is_flash_attn_2_available, is_flash_attn_3_available, is_liger_available,
+                               is_lmdeploy_available, is_megatron_available, is_swanlab_available, is_trl_available,
+                               is_unsloth_available, is_vllm_ascend_available, is_vllm_available,
+                               is_vllm_metax_available, is_wandb_available)
+    from .io_utils import JsonlWriter, append_to_jsonl, get_file_mm_type, read_from_jsonl, write_to_jsonl
+    from .logger import get_logger, ms_logger_context
+    from .np_utils import get_seed, stat_array, transform_jsonl_to_df
+    from .processor_utils import Processor, ProcessorMixin
+    from .shutdown_manager import ShutdownManager
+    from .tb_utils import TB_COLOR, TB_COLOR_SMOOTH, plot_images, read_tensorboard_file, tensorboard_smoothing
+    from .torch_utils import (Serializer, check_shared_disk, disable_safe_ddp_context_use_barrier, empty_cache,
+                              gc_collect, get_current_device, get_device, get_device_count,
+                              get_generative_reranker_logits, get_last_valid_indices, get_max_reserved_memory,
+                              get_physical_device_count, get_torch_device, init_process_group, ipc_collect,
+                              is_torch_rocm, nanstd, safe_ddp_context, set_default_ddp_config, set_device, synchronize,
+                              time_synchronize, to_device, to_float_dtype)
+    from .transformers_utils import (activate_parameters, disable_deepspeed_zero3, find_all_linears, find_embedding,
+                                     find_layers, find_norm, find_sub_module, freeze_parameters,
+                                     get_cu_seqlens_from_position_ids, get_model_parameter_info,
+                                     get_modules_to_not_convert, get_multimodal_target_regex, get_n_params_grads,
+                                     get_packed_seq_params, get_position_ids_from_cu_seqlens, seed_worker, show_layers,
+                                     unwrap_model_for_generation)
+    from .utils import (add_version_to_work_dir, check_json_format, copy_files_by_pattern, deep_getattr, find_free_port,
+                        find_node_ip, format_time, get_env_args, import_external_file, json_parse_to_dict, lower_bound,
+                        parse_args, parse_args_from_dict, patch_getattr, read_multi_line, remove_response,
+                        retry_decorator, seed_everything, shutdown_event_loop_in_daemon, split_list,
+                        start_event_loop_in_daemon, subprocess_run, swanlab_get_run, test_time, to_abspath, upper_bound)
+else:
+    _import_structure = {
+        'env': [
+            'get_dist_setting', 'get_hf_endpoint', 'get_node_setting', 'get_pai_tensorboard_dir',
+            'is_deepspeed_enabled', 'is_dist', 'is_last_rank', 'is_local_master', 'is_master', 'is_mp', 'is_mp_ddp',
+            'is_pai_training_job', 'select_device', 'use_hf_hub'
+        ],
+        'hf_config': ['HfConfigFactory'],
+        'hub_utils':
+        ['download_ms_file', 'git_clone_github', 'patch_kernels', 'safe_snapshot_download', 'download_file'],
+        'import_utils': [
+            'is_flash_attn_2_available', 'is_flash_attn_3_available', 'is_liger_available', 'is_lmdeploy_available',
+            'is_megatron_available', 'is_swanlab_available', 'is_trl_available', 'is_unsloth_available',
+            'is_vllm_ascend_available', 'is_vllm_available', 'is_vllm_metax_available', 'is_wandb_available'
+        ],
+        'io_utils': ['JsonlWriter', 'append_to_jsonl', 'get_file_mm_type', 'read_from_jsonl', 'write_to_jsonl'],
+        'logger': ['get_logger', 'ms_logger_context'],
+        'np_utils': ['get_seed', 'stat_array', 'transform_jsonl_to_df'],
+        'processor_utils': ['Processor', 'ProcessorMixin'],
+        'shutdown_manager': ['ShutdownManager'],
+        'tb_utils': ['TB_COLOR', 'TB_COLOR_SMOOTH', 'plot_images', 'read_tensorboard_file', 'tensorboard_smoothing'],
+        'torch_utils': [
+            'Serializer', 'check_shared_disk', 'disable_safe_ddp_context_use_barrier', 'empty_cache', 'gc_collect',
+            'get_current_device', 'get_device', 'get_device_count', 'get_generative_reranker_logits',
+            'get_last_valid_indices', 'get_max_reserved_memory', 'get_torch_device', 'init_process_group',
+            'ipc_collect', 'safe_ddp_context', 'set_default_ddp_config', 'set_device', 'synchronize',
+            'time_synchronize', 'to_device', 'to_float_dtype', 'nanstd', 'get_physical_device_count', 'is_torch_rocm'
+        ],
+        'transformers_utils': [
+            'activate_parameters', 'disable_deepspeed_zero3', 'find_all_linears', 'find_embedding', 'find_layers',
+            'find_norm', 'find_sub_module', 'freeze_parameters', 'get_cu_seqlens_from_position_ids',
+            'get_model_parameter_info', 'get_modules_to_not_convert', 'get_multimodal_target_regex',
+            'get_n_params_grads', 'get_packed_seq_params', 'get_position_ids_from_cu_seqlens', 'seed_worker',
+            'show_layers', 'unwrap_model_for_generation'
+        ],
+        'utils': [
+            'add_version_to_work_dir', 'check_json_format', 'copy_files_by_pattern', 'deep_getattr', 'find_free_port',
+            'find_node_ip', 'format_time', 'get_env_args', 'import_external_file', 'json_parse_to_dict', 'lower_bound',
+            'parse_args', 'parse_args_from_dict', 'patch_getattr', 'read_multi_line', 'remove_response',
+            'retry_decorator', 'seed_everything', 'shutdown_event_loop_in_daemon', 'split_list',
+            'start_event_loop_in_daemon', 'subprocess_run', 'swanlab_get_run', 'test_time', 'to_abspath', 'upper_bound'
+        ],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/utils/constants.py

@@ -0,0 +1,27 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+BIN_EXTENSIONS = [
+    '.*.bin',
+    '.*.ts',
+    '.*.pt',
+    '.*.data-00000-of-00001',
+    '.*.onnx',
+    '.*.meta',
+    '.*.pb',
+    '.*.index',
+]
+
+PEFT_TYPE_KEY = 'peft_type'
+SWIFT_TYPE_KEY = 'swift_type'
+DEFAULT_ADAPTER = 'default'
+
+
+class Invoke(object):
+    KEY = 'invoked_by'
+    THIRD_PARTY = 'third_party'
+    PRETRAINED = 'from_pretrained'
+    PIPELINE = 'pipeline'
+    TRAINER = 'trainer'
+    LOCAL_TRAINER = 'local_trainer'
+    PREPROCESSOR = 'preprocessor'
+    SWIFT = 'swift'

swift/utils/env.py

@@ -0,0 +1,97 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from transformers.integrations import deepspeed_config
+from transformers.utils import strtobool
+from typing import Optional, Tuple
+
+from .logger import get_logger
+
+logger = get_logger()
+
+
+def use_hf_hub():
+    return strtobool(os.environ.get('USE_HF', '0'))
+
+
+def get_hf_endpoint():
+    hf_endpoint = os.environ.get('HF_ENDPOINT', 'https://huggingface.co/')
+    if hf_endpoint.endswith('/'):
+        hf_endpoint = hf_endpoint[:-1]
+    return hf_endpoint
+
+
+def is_deepspeed_enabled():
+    return deepspeed_config() is not None
+
+
+def get_dist_setting() -> Tuple[int, int, int, int]:
+    """return rank, local_rank, world_size, local_world_size"""
+    rank = int(os.getenv('RANK', -1))
+    local_rank = int(os.getenv('LOCAL_RANK', -1))
+    world_size = int(os.getenv('WORLD_SIZE') or os.getenv('_PATCH_WORLD_SIZE') or 1)
+    # compat deepspeed launch
+    local_world_size = int(os.getenv('LOCAL_WORLD_SIZE', None) or os.getenv('LOCAL_SIZE', 1))
+    return rank, local_rank, world_size, local_world_size
+
+
+def get_node_setting():
+    node_rank = int(os.getenv('NODE_RANK', 0))
+    nnodes = int(os.getenv('NNODES', 1))
+    return node_rank, nnodes
+
+
+def is_local_master():
+    local_rank = get_dist_setting()[1]
+    return local_rank in {-1, 0}
+
+
+def is_master():
+    rank = get_dist_setting()[0]
+    return rank in {-1, 0}
+
+
+def is_last_rank():
+    rank, _, world_size, _ = get_dist_setting()
+    return rank in {-1, world_size - 1}
+
+
+def is_dist():
+    """Determine if the training is distributed"""
+    rank, local_rank, _, _ = get_dist_setting()
+    return rank >= 0 and local_rank >= 0
+
+
+def is_mp() -> bool:
+
+    from swift.utils import get_device_count
+    n_gpu = get_device_count()
+    local_world_size = get_dist_setting()[3]
+    if os.environ.get('SWIFT_SINGLE_DEVICE_MODE', '0') != '1':
+        assert n_gpu % local_world_size == 0, f'n_gpu: {n_gpu}, local_world_size: {local_world_size}'
+        if n_gpu // local_world_size >= 2:
+            return True
+        return False
+    else:
+        return False
+
+
+def is_mp_ddp() -> bool:
+    _, _, world_size, _ = get_dist_setting()
+    disable_mp_ddp = strtobool(os.environ.get('DISABLE_MP_DDP', '0'))
+    if not disable_mp_ddp and is_dist() and is_mp() and world_size > 1:
+        logger.info_once('Using MP(device_map) + DDP')
+        return True
+    return False
+
+
+def select_device(device_ids='0'):
+    os.environ['CUDA_VISIBLE_DEVICES'] = device_ids
+    os.environ['ASCEND_RT_VISIBLE_DEVICES'] = device_ids
+
+
+def is_pai_training_job() -> bool:
+    return 'PAI_TRAINING_JOB_ID' in os.environ
+
+
+def get_pai_tensorboard_dir() -> Optional[str]:
+    return os.environ.get('PAI_OUTPUT_TENSORBOARD')

swift/utils/hf_config.py

@@ -0,0 +1,222 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+import torch
+from transformers import PretrainedConfig
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from .utils import deep_getattr
+
+
+class HfConfigFactory:
+    llm_keys = ['language_config', 'llm_config', 'text_config']
+    vision_keys = ['vit_config', 'vision_config', 'audio_config']
+    """This class is used to read config from config.json(maybe params.json also)"""
+
+    @staticmethod
+    def get_torch_dtype(config: Union[PretrainedConfig, Dict[str, Any]],
+                        quant_info: Dict[str, Any]) -> Optional[torch.dtype]:
+        for key in ['torch_dtype', 'params_dtype']:
+            torch_dtype = HfConfigFactory.get_config_attr(config, key)
+            if torch_dtype is not None:
+                break
+        torch_dtype = HfConfigFactory.to_torch_dtype(torch_dtype)
+        if torch_dtype is None:
+            torch_dtype = quant_info.get('torch_dtype')
+        return torch_dtype
+
+    @staticmethod
+    def get_text_config(config):
+        for key in HfConfigFactory.llm_keys:
+            value = getattr(config, key, None)
+            if value is not None:
+                return value
+        return config
+
+    @staticmethod
+    def _get_config_attrs(config: Union[PretrainedConfig, Dict[str, Any]],
+                          attr_name: str,
+                          include_vit: bool = False,
+                          parent_key: Optional[str] = None) -> List[Tuple[PretrainedConfig, Any]]:
+        res = []
+        if isinstance(config, dict):
+            keys = config.keys()
+        elif isinstance(config, PretrainedConfig):
+            keys = dir(config)
+        else:
+            return []
+        config_keys = [None] + HfConfigFactory.llm_keys
+        if include_vit:
+            config_keys += HfConfigFactory.vision_keys
+        if attr_name in keys and parent_key in config_keys:
+            res.append((config, deep_getattr(config, attr_name)))
+
+        for k in keys:
+            if k.endswith('_config') and k != 'talker_config':
+                if isinstance(config, dict):
+                    v = config[k]
+                else:
+                    v = getattr(config, k)
+                res += HfConfigFactory._get_config_attrs(v, attr_name, include_vit, k)
+        return res
+
+    @staticmethod
+    def is_moe_model(config) -> bool:
+        if 'Moe' in config.__class__.__name__:
+            return True
+        for key in ['num_experts', 'num_experts_per_tok', 'moe_intermediate_size']:
+            if HfConfigFactory.get_config_attr(config, key):
+                return True
+        return False
+
+    @staticmethod
+    def is_multimodal(config) -> bool:
+        if isinstance(config, dict):
+            keys = config.keys()
+        elif isinstance(config, PretrainedConfig):
+            keys = dir(config)
+        else:
+            keys = []
+        keys = set(keys)
+        for key in (HfConfigFactory.llm_keys + HfConfigFactory.vision_keys + ['thinker_config']):
+            if key in keys:
+                return True
+        return False
+
+    @staticmethod
+    def get_config_attr(config: Union[PretrainedConfig, Dict[str, Any]],
+                        attr_name: str,
+                        include_vit: bool = False) -> Optional[Any]:
+        """Get the value of the attribute named attr_name."""
+        attrs = HfConfigFactory._get_config_attrs(config, attr_name, include_vit)
+        if len(attrs) == 0:
+            return None
+        else:
+            return attrs[0][1]
+
+    @staticmethod
+    def set_config_attr(config: Union[PretrainedConfig, Dict[str, Any]],
+                        attr_name: str,
+                        value: Any,
+                        include_vit: bool = False,
+                        ensure_set: bool = True) -> int:
+        """Set all the attr_name attributes to value."""
+        attrs = HfConfigFactory._get_config_attrs(config, attr_name, include_vit)
+        if ensure_set and len(attrs) == 0:
+            attrs.append((config, None))
+        for config, _ in attrs:
+            if isinstance(config, dict):
+                config[attr_name] = value
+            else:
+                setattr(config, attr_name, value)
+        return len(attrs)
+
+    @staticmethod
+    def del_config_attr(config: Union[PretrainedConfig, Dict[str, Any]],
+                        attr_name: str,
+                        include_vit: bool = False) -> int:
+        """Remove all the attr_name attributes."""
+        attrs = HfConfigFactory._get_config_attrs(config, attr_name, include_vit)
+        for config, _ in attrs:
+            if isinstance(config, dict):
+                config.pop(attr_name, None)
+            elif hasattr(config, attr_name):
+                delattr(config, attr_name)
+        return len(attrs)
+
+    @staticmethod
+    def get_max_model_len(config: Union[PretrainedConfig, Dict[str, Any]]) -> Optional[int]:
+        """Get the max length supported by the model"""
+        INF = int(1e9)
+        max_model_len = INF
+
+        possible_keys = [
+            'seq_length',  # qwen, chatglm
+            'max_position_embeddings',  # qwen1.5, llama2
+            'n_positions',  # polylm, phi-2
+            'model_max_length',  # baichuan2
+            # others
+            'seq_len',
+            'max_seq_len',
+            'max_sequence_length',
+            'max_seq_length',
+        ]
+        for key in possible_keys:
+            max_len_key = HfConfigFactory.get_config_attr(config, key)
+            if max_len_key is not None:
+                max_model_len = min(max_model_len, max_len_key)
+        if max_model_len == INF:
+            max_model_len = None
+        return max_model_len
+
+    @staticmethod
+    def set_max_model_len(config: Union[PretrainedConfig, Dict[str, Any]], value: int):
+        """Set the max length supported by the model"""
+
+        possible_keys = [
+            'seq_length',  # qwen, chatglm
+            'max_position_embeddings',  # qwen1.5, llama2
+            'n_positions',  # polylm, phi-2
+            'model_max_length',  # baichuan2
+            # others
+            'seq_len',
+            'max_seq_len',
+            'max_sequence_length',
+            'max_seq_length',
+        ]
+        for key in possible_keys:
+            max_len_value = HfConfigFactory.get_config_attr(config, key)
+            if max_len_value is not None:
+                HfConfigFactory.set_config_attr(config, key, value)
+
+    @staticmethod
+    def compat_zero3(config: PretrainedConfig) -> None:
+        value = HfConfigFactory.get_config_attr(config, 'hidden_size')
+        try:
+            # AttributeError: can't set attribute 'hidden_size'
+            config.hidden_size = value
+        except AttributeError:
+            pass
+
+    @staticmethod
+    def to_torch_dtype(torch_dtype: Union[str, torch.dtype, None]) -> Optional[torch.dtype]:
+        if torch_dtype is None:
+            return None
+        if isinstance(torch_dtype, str):
+            torch_dtype = torch_dtype.replace('torch.', '')
+            torch_dtype = getattr(torch, torch_dtype)
+        return torch_dtype
+
+    @staticmethod
+    def get_quant_info(config: Union[PretrainedConfig, Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+        """Get quant_method, quant_bits, dtype. not support hqq/eetq now, support awq/gptq/bnb/aqlm"""
+        if isinstance(config, dict):
+            quantization_config = config.get('quantization_config')
+        else:
+            quantization_config = getattr(config, 'quantization_config', None)
+        if quantization_config is None:
+            return
+        quantization_config = dict(quantization_config)
+        quant_method = quantization_config.get('quant_method')
+        res = {}
+        if quant_method in {'gptq', 'awq', 'aqlm'}:
+            res['quant_method'] = quant_method
+            res['torch_dtype'] = torch.float16
+            quant_bits = quantization_config.get('bits')
+            if quant_bits is not None:
+                res['quant_bits'] = quant_bits
+        elif quant_method == 'bitsandbytes':
+            res['quant_method'] = 'bnb'
+            load_in_4bit = quantization_config.get('_load_in_4bit')
+            load_in_8bit = quantization_config.get('_load_in_8bit')
+            bnb_4bit_compute_dtype = quantization_config.get('bnb_4bit_compute_dtype')
+            if load_in_4bit:
+                res['quant_bits'] = 4
+            elif load_in_8bit:
+                res['quant_bits'] = 8
+            res['torch_dtype'] = HfConfigFactory.to_torch_dtype(bnb_4bit_compute_dtype)
+        elif quant_method == 'hqq':
+            res['quant_method'] = quant_method
+            res['quant_bits'] = quantization_config['quant_config']['weight_quant_params']['nbits']
+        elif quant_method is not None:
+            res['quant_method'] = quant_method
+        return res or None

swift/utils/hub_utils.py

@@ -0,0 +1,248 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib.util
+import os
+import requests
+from modelscope.hub.api import HubApi, ModelScopeConfig
+from modelscope.hub.utils.utils import get_cache_dir
+from pathlib import Path
+from tqdm import tqdm
+from typing import List, Optional
+
+from .env import use_hf_hub
+from .logger import get_logger
+from .torch_utils import is_local_master, safe_ddp_context
+from .utils import subprocess_run
+
+logger = get_logger()
+
+
+def safe_snapshot_download(model_id_or_path: str,
+                           revision: Optional[str] = None,
+                           download_model: bool = True,
+                           use_hf: Optional[bool] = None,
+                           hub_token: Optional[str] = None,
+                           ignore_patterns: Optional[List[str]] = None,
+                           check_local: bool = False,
+                           **kwargs) -> str:
+    """Download model snapshot safely with DDP context protection.
+
+    This function attempts to download a model from HuggingFace or ModelScope hub,
+    with support for local paths, subfolder specification, and distributed training
+    context protection. It handles various path formats and provides flexible
+    file filtering options.
+
+    Args:
+        model_id_or_path (str): The model identifier on the hub (e.g., 'Qwen/Qwen2.5-7B-Instruct')
+            or a local path to the model directory. Supports subfolder specification
+            using colon syntax (e.g., 'model_id:subfolder').
+        revision (Optional[str], optional): Specific model version/revision to download
+            (branch name, tag, or commit hash). Defaults to None (latest version).
+        download_model (bool, optional): Whether to download model weight files
+            (.bin, .safetensors). If False, only config and tokenizer files are
+            downloaded. Defaults to True.
+        use_hf (Optional[bool], optional): Force using HuggingFace Hub (True) or ModelScope (False).
+            If None, it is controlled by the environment variable `USE_HF`, which defaults to '0'.
+            Default: None.
+        hub_token (Optional[str], optional): Authentication token for accessing private
+            or gated models. Defaults to None.
+        ignore_patterns (Optional[List[str]], optional): List of glob patterns for files
+            to exclude from download. If None, uses default patterns to exclude zip,
+            gguf, pth, pt, and other auxiliary files. Defaults to None.
+        check_local (bool, optional): Whether to check for a local directory matching
+            the last component of model_id_or_path before attempting download.
+            Defaults to False.
+        **kwargs: Additional keyword arguments passed to the underlying hub download function.
+
+    Returns:
+        str: Absolute path to the model directory where files are stored.
+
+    Raises:
+        ValueError: If model_id_or_path starts with '/' (absolute path) and the path
+            does not exist.
+    Examples:
+        >>> # Download from hub
+        >>> model_dir = safe_snapshot_download('Qwen/Qwen2.5-7B-Instruct')
+
+        >>> # Download config only (no weights)
+        >>> model_dir = safe_snapshot_download('Qwen/Qwen2.5-7B-Instruct', download_model=False)
+    """
+    from swift.hub import get_hub
+    if check_local:
+        model_suffix = model_id_or_path.rsplit('/', 1)[-1]
+        if os.path.exists(model_suffix):
+            model_dir = os.path.abspath(os.path.expanduser(model_suffix))
+            logger.info(f'Loading the model using local model_dir: {model_dir}')
+            return model_dir
+    if ignore_patterns is None:
+        ignore_patterns = [
+            '*.zip', '*.gguf', '*.pth', '*.pt', 'consolidated*', 'onnx/*', '*.safetensors.md', '*.msgpack', '*.onnx',
+            '*.ot', '*.h5'
+        ]
+    if not download_model:
+        ignore_patterns += ['*.bin', '*.safetensors']
+    hub = get_hub(use_hf)
+    if model_id_or_path.startswith('~'):
+        model_id_or_path = os.path.abspath(os.path.expanduser(model_id_or_path))
+    model_path_to_check = '/'.join(model_id_or_path.split(':', 1))
+    if os.path.exists(model_id_or_path):
+        model_dir = model_id_or_path
+        sub_folder = None
+    elif os.path.exists(model_path_to_check):
+        model_dir = model_path_to_check
+        sub_folder = None
+    else:
+        if model_id_or_path.startswith('/'):  # startswith
+            raise ValueError(f"path: '{model_id_or_path}' not found")
+        model_id_or_path = model_id_or_path.split(':', 1)  # get sub_folder
+        if len(model_id_or_path) == 1:
+            model_id_or_path = [model_id_or_path[0], None]
+        model_id_or_path, sub_folder = model_id_or_path
+        if sub_folder is not None:
+            kwargs['allow_patterns'] = [f"{sub_folder.rstrip('/')}/*"]
+        with safe_ddp_context(hash_id=model_id_or_path):
+            model_dir = hub.download_model(model_id_or_path, revision, ignore_patterns, token=hub_token, **kwargs)
+
+        logger.info(f'Loading the model using model_dir: {model_dir}')
+
+    model_dir = os.path.abspath(os.path.expanduser(model_dir))
+    if sub_folder:
+        model_dir = os.path.join(model_dir, sub_folder)
+    assert os.path.isdir(model_dir), f'model_dir: {model_dir}'
+    return model_dir
+
+
+def git_clone_github(github_url: str,
+                     *,
+                     local_repo_name: Optional[str] = None,
+                     branch: Optional[str] = None,
+                     commit_hash: Optional[str] = None) -> str:
+    if github_url.endswith('.git'):
+        github_url = github_url[:-4]
+    git_cache_dir = os.path.join(get_cache_dir(), '_github')
+    os.makedirs(git_cache_dir, exist_ok=True)
+    if local_repo_name is None:
+        github_url = github_url.rstrip('/')
+        local_repo_name = github_url.rsplit('/', 1)[1]
+    github_url = f'{github_url}.git'
+    local_repo_path = os.path.join(git_cache_dir, local_repo_name)
+    with safe_ddp_context('git_clone', use_barrier=True):
+        repo_existed = os.path.exists(local_repo_path)
+        if not is_local_master() and repo_existed:
+            return local_repo_path
+        if repo_existed:
+            command = ['git', '-C', local_repo_path, 'fetch']
+            subprocess_run(command)
+            if branch is not None:
+                command = ['git', '-C', local_repo_path, 'checkout', branch]
+                subprocess_run(command)
+        else:
+            command = ['git', '-C', git_cache_dir, 'clone', github_url, local_repo_name]
+            if branch is not None:
+                command += ['--branch', branch]
+            subprocess_run(command)
+
+        if commit_hash is not None:
+            command = ['git', '-C', local_repo_path, 'reset', '--hard', commit_hash]
+            subprocess_run(command)
+        elif repo_existed:
+            command = ['git', '-C', local_repo_path, 'pull']
+            subprocess_run(command)
+    logger.info(f'local_repo_path: {local_repo_path}')
+    return local_repo_path
+
+
+def download_ms_file(url: str, local_path: str, cookies=None) -> None:
+    if cookies is None:
+        cookies = ModelScopeConfig.get_cookies()
+    resp = requests.get(url, cookies=cookies, stream=True)
+    with open(local_path, 'wb') as f:
+        for data in tqdm(resp.iter_lines()):
+            f.write(data)
+
+
+def _resolve_kernel_variant_str(repo_id: str) -> Optional[str]:
+    """Resolve the kernel build variant matching the current torch/cuda/platform
+    by listing the ``build/`` folder of the ModelScope kernel repository. Returns
+    ``None`` if listing or parsing fails (caller should fall back to downloading
+    the whole repo).
+    """
+    try:
+        from kernels.variants import parse_variant, resolve_variant
+        files = HubApi().get_model_files(repo_id, root='build', recursive=False)
+        variants = []
+        for f in files:
+            name = f.get('Name') or f.get('Path', '').rsplit('/', 1)[-1]
+            if not name:
+                continue
+            try:
+                variants.append(parse_variant(name))
+            except ValueError:
+                continue
+        variant = resolve_variant(variants)
+        return variant.variant_str if variant else None
+    except Exception:
+        return None
+
+
+def patch_kernels() -> bool:
+    """Install a process-wide monkey patch on
+    ``transformers.integrations.hub_kernels.get_kernel`` so that kernel
+    repositories are downloaded from ModelScope and loaded via
+    ``kernels.get_local_kernel``.
+
+    The runtime behavior is controlled by the ``USE_HF`` env (read on each
+    ``get_kernel`` call):
+        - ``USE_HF=1``: fall back to the original HuggingFace-based loading.
+        - otherwise (default): use ModelScope.
+
+    Returns True if the patch was installed, False if skipped (``kernels`` not
+    installed, or the ``transformers`` integration is unavailable). Callers are
+    expected to guarantee idempotency (e.g. via a module-level flag).
+    """
+    if importlib.util.find_spec('kernels') is None:
+        return False
+    try:
+        from kernels import get_local_kernel
+        from transformers.integrations import hub_kernels
+    except ImportError:
+        return False
+
+    origin_get_kernel = hub_kernels.get_kernel
+
+    def patched_get_kernel(repo_id, *args, **kwargs):
+        if use_hf_hub():
+            return origin_get_kernel(repo_id, *args, **kwargs)
+        try:
+            variant_str = _resolve_kernel_variant_str(repo_id)
+            allow_patterns = [f'build/{variant_str}/*'] if variant_str else None
+            model_dir = safe_snapshot_download(repo_id, use_hf=False, allow_patterns=allow_patterns)
+            package_name = repo_id.split('/')[-1].replace('-', '_')
+            # kernels < 0.14
+            kernel = get_local_kernel(Path(model_dir), package_name)
+            logger.info(f'Loaded kernel `{repo_id}` from ModelScope: {model_dir}')
+            return kernel
+        except Exception as e:
+            logger.warning(f'Failed to load kernel `{repo_id}` from ModelScope ({e}), fallback to HuggingFace.')
+            return origin_get_kernel(repo_id, *args, **kwargs)
+
+    hub_kernels.get_kernel = patched_get_kernel
+    return True
+
+
+def download_file(url: str) -> str:
+    url = url.rstrip('/')
+    file_name = url.rsplit('/', 1)[-1]
+    cache_dir = os.path.join(get_cache_dir(), 'files')
+    os.makedirs(cache_dir, exist_ok=True)
+    file_path = os.path.join(cache_dir, file_name)
+    if os.path.exists(file_path):
+        return file_path
+    resp = requests.get(url, stream=True)
+    resp.raise_for_status()
+    total_size = int(resp.headers.get('content-length', 0))
+    with open(file_path, 'wb') as f, tqdm(
+            total=total_size, unit='B', unit_scale=True, unit_divisor=1024, desc=file_name) as pbar:
+        for chunk in resp.iter_content(chunk_size=8192):
+            f.write(chunk)
+            pbar.update(len(chunk))
+    return file_path

swift/utils/import_utils.py

@@ -0,0 +1,115 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Copyright 2023-present the HuggingFace Inc. team.
+
+import importlib.util
+import os
+from itertools import chain
+from types import ModuleType
+from typing import Any
+
+
+def is_vllm_available():
+    return importlib.util.find_spec('vllm') is not None
+
+
+def is_vllm_ascend_available():
+    return importlib.util.find_spec('vllm_ascend') is not None
+
+
+def is_vllm_metax_available():
+    return importlib.util.find_spec('vllm_metax') is not None
+
+
+def is_lmdeploy_available():
+    return importlib.util.find_spec('lmdeploy') is not None
+
+
+def is_liger_available():
+    return importlib.util.find_spec('liger_kernel') is not None
+
+
+def is_swanlab_available():
+    return importlib.util.find_spec('swanlab') is not None
+
+
+def is_megatron_available():
+    return importlib.util.find_spec('megatron') is not None
+
+
+def is_flash_attn_3_available():
+    return (importlib.util.find_spec('flash_attn_3') is not None
+            and importlib.util.find_spec('flash_attn_interface') is not None)
+
+
+def is_flash_attn_2_available():
+    return importlib.util.find_spec('flash_attn') is not None
+
+
+def is_unsloth_available() -> bool:
+    return importlib.util.find_spec('unsloth') is not None
+
+
+def is_pyreft_available() -> bool:
+    return importlib.util.find_spec('pyreft') is not None
+
+
+def is_wandb_available() -> bool:
+    return importlib.util.find_spec('wandb') is not None
+
+
+def is_trl_available() -> bool:
+    return importlib.util.find_spec('trl') is not None
+
+
+class _LazyModule(ModuleType):
+    """
+    Module class that surfaces all objects but only performs associated imports when the objects are requested.
+    """
+
+    # Very heavily inspired by optuna.integration._IntegrationModule
+    # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py
+    def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None):
+        super().__init__(name)
+        self._modules = set(import_structure.keys())
+        self._class_to_module = {}
+        for key, values in import_structure.items():
+            for value in values:
+                self._class_to_module[value] = key
+        # Needed for autocompletion in an IDE
+        self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values()))
+        self.__file__ = module_file
+        self.__spec__ = module_spec
+        self.__path__ = [os.path.dirname(module_file)]
+        self._objects = {} if extra_objects is None else extra_objects
+        self._name = name
+        self._import_structure = import_structure
+
+    # Needed for autocompletion in an IDE
+    def __dir__(self):
+        result = super().__dir__()
+        # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether
+        # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir.
+        for attr in self.__all__:
+            if attr not in result:
+                result.append(attr)
+        return result
+
+    def __getattr__(self, name: str) -> Any:
+        if name in self._objects:
+            return self._objects[name]
+        if name in self._modules:
+            value = self._get_module(name)
+        elif name in self._class_to_module.keys():
+            module = self._get_module(self._class_to_module[name])
+            value = getattr(module, name)
+        else:
+            raise AttributeError(f'module {self.__name__} has no attribute {name}')
+
+        setattr(self, name, value)
+        return value
+
+    def _get_module(self, module_name: str):
+        return importlib.import_module('.' + module_name, self.__name__)
+
+    def __reduce__(self):
+        return self.__class__, (self._name, self.__file__, self._import_structure)

swift/utils/io_utils.py

@@ -0,0 +1,122 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+import torch.distributed as dist
+from accelerate.utils import gather_object
+from queue import Queue
+from threading import Thread
+from typing import Any, Dict, List, Literal, Union
+
+from .env import is_last_rank, is_master
+from .logger import get_logger
+from .utils import check_json_format
+
+logger = get_logger()
+
+
+def read_from_jsonl(fpath: str, encoding: str = 'utf-8') -> List[Any]:
+    res: List[Any] = []
+    with open(fpath, 'r', encoding=encoding) as f:
+        for line in f:
+            res.append(json.loads(line))
+    return res
+
+
+def write_to_jsonl(fpath: str, obj_list: List[Any], encoding: str = 'utf-8') -> None:
+    res: List[str] = []
+    for obj in obj_list:
+        res.append(json.dumps(obj, ensure_ascii=False))
+    with open(fpath, 'w', encoding=encoding) as f:
+        text = '\n'.join(res)
+        f.write(f'{text}\n')
+
+
+class JsonlWriter:
+
+    def __init__(self,
+                 fpath: str,
+                 *,
+                 encoding: str = 'utf-8',
+                 strict: bool = True,
+                 enable_async: bool = False,
+                 write_on_rank: Literal['master', 'last'] = 'master'):
+        if write_on_rank == 'master':
+            self.is_write_rank = is_master()
+        elif write_on_rank == 'last':
+            self.is_write_rank = is_last_rank()
+        else:
+            raise ValueError(f"Invalid `write_on_rank`: {write_on_rank}, should be 'master' or 'last'")
+        self.fpath = os.path.abspath(os.path.expanduser(fpath)) if self.is_write_rank else None
+        self.encoding = encoding
+        self.strict = strict
+        self.enable_async = enable_async
+        self._queue = Queue()
+        self._thread = None
+
+    def _append_worker(self):
+        while True:
+            item = self._queue.get()
+            self._append(**item)
+
+    def _append(self, obj: Union[Dict, List[Dict]], gather_obj: bool = False):
+        if isinstance(obj, (list, tuple)) and all(isinstance(item, dict) for item in obj):
+            obj_list = obj
+        else:
+            obj_list = [obj]
+        if gather_obj and dist.is_initialized():
+            obj_list = gather_object(obj_list)
+        if not self.is_write_rank:
+            return
+        obj_list = check_json_format(obj_list)
+        for i, _obj in enumerate(obj_list):
+            obj_list[i] = json.dumps(_obj, ensure_ascii=False) + '\n'
+        self._write_buffer(''.join(obj_list))
+
+    def append(self, obj: Union[Dict, List[Dict]], gather_obj: bool = False):
+        if self.enable_async:
+            if self._thread is None:
+                self._thread = Thread(target=self._append_worker, daemon=True)
+                self._thread.start()
+            self._queue.put({'obj': obj, 'gather_obj': gather_obj})
+        else:
+            self._append(obj, gather_obj=gather_obj)
+
+    def _write_buffer(self, text: str):
+        if not text:
+            return
+        assert self.is_write_rank, f'self.is_write_rank: {self.is_write_rank}'
+        try:
+            os.makedirs(os.path.dirname(self.fpath), exist_ok=True)
+            with open(self.fpath, 'a', encoding=self.encoding) as f:
+                f.write(text)
+        except Exception:
+            if self.strict:
+                raise
+            logger.error(f'Cannot write content to jsonl file. text: {text}')
+
+
+def append_to_jsonl(fpath: str,
+                    obj: Union[Dict, List[Dict]],
+                    *,
+                    encoding: str = 'utf-8',
+                    strict: bool = True,
+                    write_on_rank: Literal['master', 'last'] = 'master') -> None:
+    jsonl_writer = JsonlWriter(fpath, encoding=encoding, strict=strict, write_on_rank=write_on_rank)
+    jsonl_writer.append(obj)
+
+
+def get_file_mm_type(file_name: str) -> Literal['image', 'video', 'audio']:
+    video_extensions = {'.mp4', '.mkv', '.mov', '.avi', '.wmv', '.flv', '.webm'}
+    audio_extensions = {'.mp3', '.wav', '.aac', '.flac', '.ogg', '.m4a'}
+    image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
+
+    _, ext = os.path.splitext(file_name)
+
+    if ext.lower() in video_extensions:
+        return 'video'
+    elif ext.lower() in audio_extensions:
+        return 'audio'
+    elif ext.lower() in image_extensions:
+        return 'image'
+    else:
+        raise ValueError(f'file_name: {file_name}, ext: {ext}')

swift/utils/logger.py

@@ -0,0 +1,156 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib.util
+import logging
+import os
+from contextlib import contextmanager
+from modelscope.utils.logger import get_logger as get_ms_logger
+from types import MethodType
+from typing import Optional
+
+
+# Avoid circular reference
+def _is_local_master():
+    local_rank = int(os.getenv('LOCAL_RANK', -1))
+    return local_rank in {-1, 0}
+
+
+init_loggers = {}
+
+# old format
+# formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger_format = logging.Formatter('[%(levelname)s:%(name)s] %(message)s')
+
+info_set = set()
+warning_set = set()
+
+
+def info_if(self, msg, cond, *args, **kwargs):
+    if cond:
+        with logger_context(self, logging.INFO):
+            self.info(msg)
+
+
+def warning_if(self, msg, cond, *args, **kwargs):
+    if cond:
+        with logger_context(self, logging.INFO):
+            self.warning(msg)
+
+
+def info_once(self, msg, *args, **kwargs):
+    hash_id = kwargs.get('hash_id') or msg
+    if hash_id in info_set:
+        return
+    info_set.add(hash_id)
+    self.info(msg)
+
+
+def warning_once(self, msg, *args, **kwargs):
+    hash_id = kwargs.get('hash_id') or msg
+    if hash_id in warning_set:
+        return
+    warning_set.add(hash_id)
+    self.warning(msg)
+
+
+def get_logger(log_file: Optional[str] = None, log_level: Optional[int] = None, file_mode: str = 'w'):
+    """ Get logging logger
+
+    Args:
+        log_file: Log filename, if specified, file handler will be added to
+            logger
+        log_level: Logging level.
+        file_mode: Specifies the mode to open the file, if filename is
+            specified (if filemode is unspecified, it defaults to 'w').
+    """
+    if log_level is None:
+        log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
+        log_level = getattr(logging, log_level, logging.INFO)
+    logger_name = __name__.split('.')[0]
+    logger = logging.getLogger(logger_name)
+    logger.propagate = False
+    if logger_name in init_loggers:
+        add_file_handler_if_needed(logger, log_file, file_mode, log_level)
+        return logger
+
+    # handle duplicate logs to the console
+    # Starting in 1.8.0, PyTorch DDP attaches a StreamHandler <stderr> (NOTSET)
+    # to the root logger. As logger.propagate is True by default, this root
+    # level handler causes logging messages from rank>0 processes to
+    # unexpectedly show up on the console, creating much unwanted clutter.
+    # To fix this issue, we set the root logger's StreamHandler, if any, to log
+    # at the ERROR level.
+    for handler in logger.root.handlers:
+        if type(handler) is logging.StreamHandler:
+            handler.setLevel(logging.ERROR)
+
+    stream_handler = logging.StreamHandler()
+    handlers = [stream_handler]
+
+    is_worker0 = _is_local_master()
+
+    if is_worker0 and log_file is not None:
+        file_handler = logging.FileHandler(log_file, file_mode)
+        handlers.append(file_handler)
+
+    for handler in handlers:
+        handler.setFormatter(logger_format)
+        handler.setLevel(log_level)
+        logger.addHandler(handler)
+
+    if is_worker0:
+        logger.setLevel(log_level)
+    else:
+        logger.setLevel(logging.ERROR)
+
+    init_loggers[logger_name] = True
+
+    logger.info_once = MethodType(info_once, logger)
+    logger.warning_once = MethodType(warning_once, logger)
+    logger.info_if = MethodType(info_if, logger)
+    logger.warning_if = MethodType(warning_if, logger)
+    return logger
+
+
+logger = get_logger()
+ms_logger = get_ms_logger()
+
+logger.handlers[0].setFormatter(logger_format)
+ms_logger.handlers[0].setFormatter(logger_format)
+log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
+if _is_local_master():
+    ms_logger.setLevel(log_level)
+else:
+    ms_logger.setLevel(logging.ERROR)
+
+
+@contextmanager
+def logger_context(logger, log_leval):
+    origin_log_level = logger.level
+    logger.setLevel(log_leval)
+    try:
+        yield
+    finally:
+        logger.setLevel(origin_log_level)
+
+
+@contextmanager
+def ms_logger_context(log_leval):
+    with logger_context(get_ms_logger(), log_leval):
+        yield
+
+
+def add_file_handler_if_needed(logger, log_file, file_mode, log_level):
+    for handler in logger.handlers:
+        if isinstance(handler, logging.FileHandler):
+            return
+
+    if importlib.util.find_spec('torch') is not None:
+        is_worker0 = int(os.getenv('LOCAL_RANK', -1)) in {-1, 0}
+    else:
+        is_worker0 = True
+
+    if is_worker0 and log_file is not None:
+        file_handler = logging.FileHandler(log_file, file_mode)
+        file_handler.setFormatter(logger_format)
+        file_handler.setLevel(log_level)
+        logger.addHandler(file_handler)

swift/utils/np_utils.py

@@ -0,0 +1,39 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+import pandas as pd
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+def transform_jsonl_to_df(dict_list: List[Dict[str, Any]]) -> pd.DataFrame:
+    """Relevant function: `io_utils.read_from_jsonl()`"""
+    data_dict: Dict[str, List[Any]] = {}
+    for i, obj in enumerate(dict_list):
+        for k, v in obj.items():
+            if k not in data_dict:
+                data_dict[k] = [None] * i
+            data_dict[k].append(v)
+        for k in set(data_dict.keys()) - set(obj.keys()):
+            data_dict[k].append(None)
+    return pd.DataFrame.from_dict(data_dict)
+
+
+def get_seed(random_state: Optional[np.random.RandomState] = None) -> int:
+    if random_state is None:
+        random_state = np.random.RandomState()
+    seed_max = np.iinfo(np.int32).max
+    seed = random_state.randint(0, seed_max)
+    return seed
+
+
+def stat_array(array: Union[np.ndarray, List[int], 'torch.Tensor']) -> Tuple[Dict[str, float], str]:
+    if isinstance(array, list):
+        if array and isinstance(array[0], list):
+            array = np.array([sum(sublist) for sublist in array])
+        array = np.array(array)
+    mean = array.mean().item()
+    std = array.std().item()
+    min_ = array.min().item()
+    max_ = array.max().item()
+    size = array.shape[0]
+    string = f'{mean:.6f}±{std:.6f}, min={min_:.6f}, max={max_:.6f}, size={size}'
+    return {'mean': mean, 'std': std, 'min': min_, 'max': max_, 'size': size}, string

swift/utils/processor_utils.py

@@ -0,0 +1,30 @@
+import os
+from transformers import FeatureExtractionMixin, PreTrainedTokenizerBase
+from transformers import ProcessorMixin as HfProcessorMixin
+from typing import Union
+
+try:
+    from transformers import BaseImageProcessor
+    Processor = Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, HfProcessorMixin]
+except ImportError:
+    Processor = Union[PreTrainedTokenizerBase, FeatureExtractionMixin, HfProcessorMixin]
+
+if 'TOKENIZERS_PARALLELISM' not in os.environ:
+    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+
+class ProcessorMixin:
+
+    @property
+    def tokenizer(self):
+        tokenizer = self.processor
+        if not isinstance(tokenizer, PreTrainedTokenizerBase) and hasattr(tokenizer, 'tokenizer'):
+            tokenizer = tokenizer.tokenizer
+        return tokenizer
+
+    @tokenizer.setter
+    def tokenizer(self, value):
+        if self.processor is self.tokenizer:
+            self.processor = value
+        elif self.tokenizer is not value:
+            raise AttributeError('Please use `self.processor` for assignment.')

swift/utils/shutdown_manager.py

@@ -0,0 +1,35 @@
+import os
+import signal
+
+
+class ShutdownManager:
+
+    def __init__(self, signals=None, stop_file=None):
+        if signals is None:
+            signals = [signal.SIGTERM, signal.SIGINT, signal.SIGUSR1, signal.SIGUSR2]
+        self._signals = signals
+        self._stop_file = stop_file or '/tmp/stop'
+
+        self._shutdown_requested = False
+        self._old_handlers = {}
+
+    def _handler(self, signum, frame):
+        self._shutdown_requested = True
+
+    def register(self):
+        for s in self._signals:
+            self._old_handlers[s] = signal.getsignal(s)
+            signal.signal(s, self._handler)
+
+    def unregister(self):
+        for s, handler in self._old_handlers.items():
+            signal.signal(s, handler)
+        self._old_handlers = {}
+
+    def should_shutdown(self) -> bool:
+        if self._shutdown_requested:
+            return True
+        return os.path.exists(self._stop_file)
+
+    def reset(self):
+        self._shutdown_requested = False

swift/utils/tb_utils.py

@@ -0,0 +1,82 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+Item = Dict[str, float]
+TB_COLOR, TB_COLOR_SMOOTH = '#FFE2D9', '#FF7043'
+
+
+def read_tensorboard_file(fpath: str) -> Dict[str, List[Item]]:
+    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
+    if not os.path.isfile(fpath):
+        raise FileNotFoundError(f'fpath: {fpath}')
+    ea = EventAccumulator(fpath)
+    ea.Reload()
+    res: Dict[str, List[Item]] = {}
+    tags = ea.Tags()['scalars']
+    for tag in tags:
+        values = ea.Scalars(tag)
+        r: List[Item] = []
+        for v in values:
+            r.append({'step': v.step, 'value': v.value})
+        res[tag] = r
+    return res
+
+
+def tensorboard_smoothing(values: List[float], smooth: float = 0.9) -> List[float]:
+    norm_factor = 0
+    x = 0
+    res: List[float] = []
+    for i in range(len(values)):
+        x = x * smooth + values[i]  # Exponential decay
+        norm_factor *= smooth
+        norm_factor += 1
+        res.append(x / norm_factor)
+    return res
+
+
+def plot_images(images_dir: str,
+                tb_dir: str,
+                smooth_key: Optional[List[str]] = None,
+                smooth_val: float = 0.9,
+                figsize: Tuple[int, int] = (8, 5),
+                dpi: int = 100) -> None:
+    """Using tensorboard's data content to plot images"""
+    import matplotlib.pyplot as plt
+    if not os.path.exists(tb_dir):
+        return
+    smooth_key = smooth_key or []
+    os.makedirs(images_dir, exist_ok=True)
+
+    matches = []
+    for root, dirs, files in os.walk(tb_dir):
+        for f in files:
+            if f.startswith('events.out.tfevents.'):
+                matches.append(os.path.join(root, f))
+    if not matches:
+        return
+    fname = matches[0]
+    tb_path = os.path.join(tb_dir, fname)
+    data = read_tensorboard_file(tb_path)
+
+    for k in data.keys():
+        _data = data[k]
+        steps = [d['step'] for d in _data]
+        values = [d['value'] for d in _data]
+        if len(values) == 0:
+            continue
+        _, ax = plt.subplots(1, 1, squeeze=True, figsize=figsize, dpi=dpi)
+        ax.set_title(k)
+        if len(values) == 1:
+            ax.scatter(steps, values, color=TB_COLOR_SMOOTH)
+        elif k in smooth_key:
+            ax.plot(steps, values, color=TB_COLOR, label='original')
+            values_s = tensorboard_smoothing(values, smooth_val)
+            ax.plot(steps, values_s, color=TB_COLOR_SMOOTH, label='smoothed')
+            ax.legend()
+        else:
+            ax.plot(steps, values, color=TB_COLOR_SMOOTH)
+        fpath = os.path.join(images_dir, k.replace('/', '_').replace('.', '_'))
+        plt.savefig(fpath, dpi=dpi, bbox_inches='tight')
+        plt.close()

swift/utils/torch_utils.py

@@ -0,0 +1,419 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import gc
+import hashlib
+import numpy as np
+import os
+import pickle
+import time
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import uuid
+from contextlib import contextmanager
+from datasets.utils.filelock import FileLock
+from datetime import timedelta
+from modelscope.hub.utils.utils import get_cache_dir
+from transformers.utils import is_torch_cuda_available, is_torch_mps_available, is_torch_npu_available
+from typing import Any, Mapping, Optional, Union
+
+from .env import get_dist_setting, get_node_setting, is_dist, is_local_master, is_master, is_mp
+
+
+def nanstd(tensor: torch.Tensor, dim: Optional[Union[int, tuple]] = None, keepdim: bool = False) -> torch.Tensor:
+    """Compute the standard deviation of a tensor, ignoring NaNs.
+
+    Refer: trl/trainer/utils.py
+
+    Args:
+        tensor (`torch.Tensor`):
+            Input tensor.
+        dim (`int` or `tuple[int, ...]`, *optional*):
+            Dimension to reduce. Defaults to all dimensions.
+        keepdim (`bool`, *optional*, defaults to `False`):
+            Whether to keep reduced dimensions.
+
+    Returns:
+        `torch.Tensor`:
+            Standard deviation of the tensor, ignoring NaNs.
+    """
+    mean = torch.nanmean(tensor, dim=dim, keepdim=True)
+    variance = torch.nanmean((tensor - mean)**2, dim=dim, keepdim=True)
+    count = torch.sum(~torch.isnan(tensor), dim=dim, keepdim=True)
+    correction = count / (count - 1)
+    correction = torch.where(count > 1, correction, torch.full_like(correction, float('nan')))
+    variance *= correction  # Bessel's correction
+    std = torch.sqrt(variance)
+    if keepdim:
+        return std
+    if dim is None:
+        return std.squeeze()
+    if isinstance(dim, int):
+        return std.squeeze(dim)
+    dims = [(d if d >= 0 else d + std.ndim) for d in dim]
+    for d in sorted(dims, reverse=True):
+        std = std.squeeze(d)
+    return std
+
+
+def _find_local_mac() -> str:
+    mac = uuid.getnode()
+    mac_address = ':'.join(('%012x' % mac)[i:i + 2] for i in range(0, 12, 2))
+    return mac_address
+
+
+def synchronize(device: Union[torch.device, str, int, None] = None):
+    if is_torch_npu_available():
+        torch.npu.synchronize(device)
+    elif is_torch_cuda_available():
+        torch.cuda.synchronize(device)
+    else:
+        torch.cuda.synchronize(device)
+
+
+def time_synchronize() -> float:
+    synchronize()
+    return time.perf_counter()  # second
+
+
+_DISABLE_USE_BARRIER = False
+
+
+@contextmanager
+def disable_safe_ddp_context_use_barrier():
+    global _DISABLE_USE_BARRIER
+    _DISABLE_USE_BARRIER = True
+    try:
+        yield
+    finally:
+        _DISABLE_USE_BARRIER = False
+
+
+@contextmanager
+def safe_ddp_context(hash_id: Optional[str], use_barrier: bool = True):
+    if _DISABLE_USE_BARRIER:
+        use_barrier = False
+    if use_barrier and dist.is_initialized():
+        if is_dist():
+            if not is_master():
+                dist.barrier()
+            if not is_local_master():
+                # Compatible with multi-machine scenarios,
+                # where each machine uses different storage hardware.
+                dist.barrier()
+        yield
+        if is_dist():
+            if is_master():
+                dist.barrier()
+            if is_local_master():
+                dist.barrier()
+    elif hash_id is not None:
+        lock_dir = os.path.join(get_cache_dir(), 'lockers')
+        os.makedirs(lock_dir, exist_ok=True)
+        file_path = hashlib.sha256(hash_id.encode('utf-8')).hexdigest() + '.lock'
+        file_path = os.path.join(lock_dir, file_path)
+        with FileLock(file_path):
+            yield
+    else:
+        yield
+
+
+def get_device(local_rank: Optional[Union[str, int]] = None) -> str:
+    if local_rank is None:
+        local_rank = max(0, get_dist_setting()[1])
+    local_rank = str(local_rank)
+    if is_torch_npu_available():
+        device = 'npu:{}'.format(local_rank)
+    elif is_torch_mps_available():
+        device = 'mps:{}'.format(local_rank)
+    elif is_torch_cuda_available():
+        device = 'cuda:{}'.format(local_rank)
+    else:
+        device = 'cpu'
+
+    return device
+
+
+def get_current_device():
+    if is_torch_npu_available():
+        current_device = torch.npu.current_device()
+    elif is_torch_cuda_available():
+        current_device = torch.cuda.current_device()
+    elif is_torch_mps_available():
+        current_device = 'mps'
+    else:
+        current_device = 'cpu'
+    return current_device
+
+
+def get_torch_device():
+    if is_torch_cuda_available():
+        return torch.cuda
+    elif is_torch_npu_available():
+        return torch.npu
+    elif is_torch_mps_available():
+        return torch.mps
+    else:
+        return torch.cpu
+
+
+def set_device(local_rank: Optional[Union[str, int]] = None):
+    if local_rank is None:
+        local_rank = max(0, get_dist_setting()[1])
+    if is_torch_npu_available():
+        torch.npu.set_device(local_rank)
+    elif is_torch_cuda_available():
+        torch.cuda.set_device(local_rank)
+
+
+def get_device_count() -> int:
+    if is_torch_npu_available():
+        return torch.npu.device_count()
+    elif is_torch_cuda_available():
+        return torch.cuda.device_count()
+    else:
+        return 0
+
+
+def is_torch_rocm() -> bool:
+    """True on an AMD ROCm/HIP torch build.
+
+    PyTorch on ROCm exposes devices through the ``torch.cuda.*`` API (hipify), so
+    ``is_torch_cuda_available()`` is True on ROCm too; the distinguishing signal is
+    a non-None ``torch.version.hip``.
+    """
+    return is_torch_cuda_available() and getattr(torch.version, 'hip', None) is not None
+
+
+def get_physical_device_count() -> int:
+    """Number of physical accelerators, ignoring ``*_VISIBLE_DEVICES`` masks.
+
+    Unlike :func:`get_device_count` (which honors CUDA/HIP_VISIBLE_DEVICES via the
+    runtime), this queries the driver/topology layer, so a process that inherited a
+    restricted mask can still discover the full device set. It still respects the
+    container's device exposure (``--gpus`` / ``NVIDIA_VISIBLE_DEVICES`` / passed
+    ``/dev/dri`` + ``/dev/kfd``). Falls back to the mask-dependent runtime count if
+    the driver query is unavailable.
+    """
+    import glob
+
+    if is_torch_rocm():
+        # ROCm: the kfd topology lists one node per device; GPU nodes have
+        # ``simd_count > 0`` (CPU/host nodes have 0). Reading sysfs neither
+        # initializes HIP nor is affected by the visibility mask.
+        count = 0
+        for prop in glob.glob('/sys/class/kfd/kfd/topology/nodes/*/properties'):
+            try:
+                with open(prop) as f:
+                    for line in f:
+                        if line.startswith('simd_count'):
+                            if int(line.split()[1]) > 0:
+                                count += 1
+                            break
+            except Exception:
+                continue
+        if count > 0:
+            return count
+    elif is_torch_cuda_available():
+        # NVIDIA: NVML/procfs/nvidia-smi enumerate at the driver level and are not
+        # filtered by CUDA_VISIBLE_DEVICES (only by container-level exposure).
+        try:
+            import pynvml
+            pynvml.nvmlInit()
+            try:
+                count = pynvml.nvmlDeviceGetCount()
+            finally:
+                pynvml.nvmlShutdown()
+            if count > 0:
+                return count
+        except Exception:
+            pass
+        count = len(glob.glob('/proc/driver/nvidia/gpus/*'))
+        if count > 0:
+            return count
+        try:
+            import subprocess
+            out = subprocess.check_output(['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'],
+                                          stderr=subprocess.DEVNULL,
+                                          timeout=5)
+            count = len([line for line in out.decode().splitlines() if line.strip()])
+            if count > 0:
+                return count
+        except Exception:
+            pass
+
+    return get_device_count()
+
+
+def empty_cache():
+    if is_torch_npu_available():
+        torch.npu.empty_cache()
+    elif is_torch_mps_available():
+        torch.mps.empty_cache()
+    elif is_torch_cuda_available():
+        torch.cuda.empty_cache()
+
+
+def ipc_collect():
+    if is_torch_cuda_available():
+        torch.cuda.ipc_collect()
+    elif is_torch_npu_available():
+        torch.npu.ipc_collect()
+
+
+def gc_collect() -> None:
+    gc.collect()
+    empty_cache()
+
+
+def get_last_valid_indices(attention_mask: torch.Tensor) -> torch.Tensor:
+    """
+    Get the last valid (non-padding) token position indices for each sample.
+
+    This function correctly handles sequences with different padding directions (left/right/none)
+    within the same batch by computing the last valid index for each sequence individually.
+
+    Args:
+        attention_mask: Attention mask [batch_size, seq_len] where 1=valid, 0=padding
+
+    Returns:
+        torch.Tensor: Indices of last valid positions [batch_size]
+
+    Examples:
+        >>> # Right padding
+        >>> attention_mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 0]])
+        >>> get_last_valid_indices(attention_mask)
+        tensor([2, 3])
+
+        >>> # Left padding
+        >>> attention_mask = torch.tensor([[0, 0, 1, 1, 1], [0, 1, 1, 1, 1]])
+        >>> get_last_valid_indices(attention_mask)
+        tensor([4, 4])
+    """
+    seq_len = attention_mask.shape[1]
+
+    # Flip the mask horizontally to bring the last elements to the front.
+    # `argmax` will then find the index of the first '1', which corresponds to the last valid token.
+    last_valid_indices = torch.fliplr(attention_mask).argmax(dim=1)
+
+    # Convert the index from the right-to-left frame to the original left-to-right frame.
+    indices = seq_len - 1 - last_valid_indices
+
+    return indices
+
+
+class Serializer:
+
+    @staticmethod
+    def to_tensor(obj):
+        res = pickle.dumps(obj)
+        res = np.array([len(res)], dtype=np.int64).tobytes() + res
+        res = np.frombuffer(res, dtype=np.uint8).copy()
+        res = torch.from_numpy(res)
+        return res
+
+    @staticmethod
+    def from_tensor(obj):
+        if isinstance(obj, torch.Tensor):
+            obj = obj.cpu().numpy()
+        res = obj.tobytes()
+        buffer_size = np.frombuffer(res[:8], dtype=np.int64)[0]
+        res = res[8:]
+        return pickle.loads(res[:buffer_size])
+
+
+def set_default_ddp_config():
+    # It runs normally with Python as well.
+    rank, local_rank, _, _ = get_dist_setting()
+    if rank == -1 or local_rank == -1:
+        os.environ['NPROC_PER_NODE'] = '1'
+        os.environ['RANK'] = '0'
+        os.environ['LOCAL_RANK'] = '0'
+        os.environ['WORLD_SIZE'] = '1'
+        os.environ['LOCAL_WORLD_SIZE'] = '1'
+        os.environ['MASTER_ADDR'] = '127.0.0.1'
+        os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '29500')
+
+
+def init_process_group(backend: Optional[str] = None, timeout: int = 18000000):
+    if dist.is_initialized():
+        return
+    set_device()
+    if backend is None:
+        if is_torch_npu_available():
+            backend = 'hccl'
+        elif torch.cuda.is_available():
+            backend = 'nccl'
+        else:
+            backend = 'gloo'
+    timeout = timedelta(seconds=timeout)
+    dist.init_process_group(backend=backend, timeout=timeout)
+
+
+def check_shared_disk(error, cache_dir: Optional[str] = None):
+    nnodes = get_node_setting()[1]
+    if nnodes <= 1:
+        return True
+    assert dist.is_initialized()
+    if cache_dir is None:
+        cache_dir = os.path.join(get_cache_dir(), 'tmp')
+    os.makedirs(cache_dir, exist_ok=True)
+    tmp_path = os.path.join(cache_dir, 'check_shared_disk.tmp')
+    is_shared_disk = True
+
+    try:
+        with safe_ddp_context(None, True):
+            if is_master():
+                with open(tmp_path, 'w'):
+                    pass
+            if not os.path.exists(tmp_path):
+                is_shared_disk = False
+        shared_state = [None] * dist.get_world_size()
+        dist.all_gather_object(shared_state, is_shared_disk)
+    finally:
+        if is_master() and os.path.exists(tmp_path):
+            os.remove(tmp_path)
+    if not all(shared_state):
+        raise error
+
+
+def to_float_dtype(data: Any, dtype: torch.dtype) -> Any:
+    """Change the float inputs to a dtype"""
+    if isinstance(data, Mapping):
+        return type(data)({k: to_float_dtype(v, dtype) for k, v in data.items()})
+    elif isinstance(data, (tuple, list)):
+        return type(data)(to_float_dtype(v, dtype) for v in data)
+    elif isinstance(data, torch.Tensor) and torch.is_floating_point(data):
+        return data.to(dtype=dtype)
+    else:
+        return data
+
+
+def to_device(data: Any, device: Union[str, torch.device, int], non_blocking: bool = False) -> Any:
+    """Move inputs to a device"""
+    if isinstance(data, Mapping):
+        return type(data)({k: to_device(v, device, non_blocking) for k, v in data.items()})
+    elif isinstance(data, (tuple, list)):
+        return type(data)(to_device(v, device, non_blocking) for v in data)
+    elif isinstance(data, torch.Tensor):
+        return data.to(device=device, non_blocking=non_blocking)
+    else:
+        return data
+
+
+def get_generative_reranker_logits(lm_head_weight, tokenizer, hidden_states):
+    positive_token = os.environ.get('GENERATIVE_RERANKER_POSITIVE_TOKEN', 'yes')
+    negative_token = os.environ.get('GENERATIVE_RERANKER_NEGATIVE_TOKEN', 'no')
+    positive_token_id = tokenizer.convert_tokens_to_ids(positive_token)
+    negative_token_id = tokenizer.convert_tokens_to_ids(negative_token)
+    weight = lm_head_weight[[positive_token_id, negative_token_id]]
+    logits = F.linear(hidden_states, weight)
+    return logits[..., 0:1] - logits[..., 1:2]
+
+
+def get_max_reserved_memory() -> float:
+    devices = list(range(get_device_count())) if is_mp() else [None]
+    try:
+        mems = [get_torch_device().max_memory_reserved(device=device) for device in devices]
+    except AttributeError:
+        return 0  # fix mps
+    return sum(mems) / 1024**3

swift/utils/transformers_utils.py

@@ -0,0 +1,360 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import numpy as np
+import re
+import torch
+import torch.nn as nn
+from bisect import bisect_right
+from contextlib import contextmanager, nullcontext
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.trainer_utils import set_seed
+from typing import Callable, List, Optional, Tuple
+
+from .logger import get_logger
+from .utils import deep_getattr
+
+logger = get_logger()
+
+
+def get_n_params_grads(model) -> Tuple[List[int], List[int]]:
+    n_params, n_grads = [], []
+    for p in model.parameters():
+        if is_deepspeed_zero3_enabled():
+            import deepspeed
+            context = deepspeed.zero.GatheredParameters(p)
+        else:
+            context = nullcontext()
+        with context:
+            n_params.append(p.numel())
+            n_grads.append(p.numel() if p.requires_grad else 0)
+    return n_params, n_grads
+
+
+def get_model_parameter_info(model: nn.Module, name: Optional[str] = None) -> str:
+    n_params, n_grads = get_n_params_grads(model)
+    n_params = sum(n_params)
+    n_grads = sum(n_grads)
+    n_buffers = sum(p.numel() for p in model.buffers())
+
+    if name is None:
+        name = model.__class__.__name__
+
+    n_params /= 1e6
+    n_grads /= 1e6
+    n_buffers /= 1e6
+    s = (f'{name}: '
+         f'{n_params:.4f}M Params ({n_grads:.4f}M Trainable '
+         f'[{100 * n_grads / n_params:.4f}%]), '
+         f'{n_buffers:.4f}M Buffers.')
+    return s
+
+
+def find_sub_module(module: torch.nn.Module, module_name: str) -> List[torch.nn.Module]:
+    _modules = list()
+    for name, sub_module in module.named_modules():
+        if not name:
+            continue
+        if name.endswith(module_name):
+            _modules.append(sub_module)
+    return _modules
+
+
+def show_layers(model: nn.Module, max_lines: Optional[int] = 20) -> None:
+    named_p = list(model.named_parameters())
+    for i, (n, p) in enumerate(named_p):
+        if max_lines is not None and i >= max_lines:
+            logger.info('...')
+            break
+        logger.info(f'[{n}]: requires_grad={p.requires_grad}, dtype={p.dtype}, device={p.device}')
+
+
+def freeze_parameters(model: nn.Module,
+                      freeze_parameters_ratio: float,
+                      freeze_parameters: List[str],
+                      freeze_parameters_regex: Optional[str] = None) -> None:
+    if freeze_parameters_ratio > 0:
+        n_parameters = get_n_params_grads(model)[0]
+        n_parameters = np.array(n_parameters, dtype=np.int64)
+        n_freeze_parameters = int(np.sum(n_parameters) * freeze_parameters_ratio)
+        n_parameters_cs = np.cumsum(n_parameters)
+        idx = bisect_right(n_parameters_cs, n_freeze_parameters)
+        for _, p in zip(range(idx), model.parameters()):
+            p.requires_grad = False
+
+    if freeze_parameters:
+        for n, p in model.named_parameters():
+            for freeze_p in freeze_parameters:
+                if n.startswith(freeze_p):
+                    p.requires_grad = False
+
+    if freeze_parameters_regex is not None:
+        try:
+            pattern = re.compile(freeze_parameters_regex)
+        except re.error as e:
+            logger.warning(f"Invalid freeze_parameters_regex '{freeze_parameters_regex}': {e}")
+            return
+
+        for n, p in model.named_parameters():
+            if pattern.search(n):
+                p.requires_grad = False
+
+
+def activate_parameters(model: nn.Module,
+                        additional_trainable_parameters: List[str],
+                        trainable_parameters_regex: Optional[str] = None) -> None:
+    has_activate = False
+    if len(additional_trainable_parameters) > 0:
+        for n, p in model.named_parameters():
+            for additional_tp in additional_trainable_parameters:
+                if n.startswith(additional_tp):
+                    p.requires_grad = True
+                    has_activate = True
+        if not has_activate:
+            logger.warning('len(additional_trainable_parameters) > 0 but no parameters are activated. '
+                           f'additional_trainable_parameters: {additional_trainable_parameters}')
+
+    has_activate = False
+    if trainable_parameters_regex is not None:
+        try:
+            pattern = re.compile(trainable_parameters_regex)
+        except re.error as e:
+            logger.warning(f"Invalid trainable_parameters_regex '{trainable_parameters_regex}': {e}")
+            return
+
+        for n, p in model.named_parameters():
+            if pattern.search(n):
+                p.requires_grad = True
+                has_activate = True
+
+        if not has_activate:
+            logger.warning('trainable_parameters_regex is provided but no parameters are activated. '
+                           f'trainable_parameters_regex: {trainable_parameters_regex}')
+
+
+def find_layers(
+    model: nn.Module,
+    cond: Callable[[str, nn.Module], bool],
+    sub_module: Optional[str] = None,
+    min_name_len: Optional[int] = None,
+) -> List[str]:
+    # The content of target_module_names cannot exist in inner_nodes.
+    sub_module_str = sub_module
+    if sub_module is None:
+        sub_module = model
+    else:
+        sub_module = deep_getattr(model, sub_module)
+    inner_nodes = set()
+    for name, module in model.named_modules():
+        name = re.sub(r'\d+\.', '{}.', name)
+        if not cond(name, module):
+            inner_nodes.add(name)
+    target_module_names = set()
+    for name, module in sub_module.named_modules():
+        if sub_module_str:
+            name = f'{sub_module_str}.{name}' if name else sub_module_str
+        if cond(name, module):
+            module_name_list = name.split('.')
+            module_name = module_name_list.pop()
+            i = 1
+            for inner_node in inner_nodes:
+                while module_name_list and inner_node.endswith(re.sub(
+                        r'\d+\.', '{}.', module_name)) or min_name_len and i < min_name_len:
+                    module_name = f'{module_name_list.pop()}.{module_name}'
+                    i += 1
+            target_module_names.add(module_name)
+    return list(target_module_names)
+
+
+def find_norm(model: nn.Module) -> List[str]:
+    # find_layer_norm
+    return find_layers(
+        model,
+        lambda name, module: isinstance(module, torch.nn.LayerNorm) or 'rmsnorm' in module.__class__.__name__.lower())
+
+
+def find_embedding(model: nn.Module) -> List[str]:
+    return find_layers(model, lambda name, module: isinstance(module, torch.nn.Embedding))
+
+
+def find_all_linears(model, model_arch=None, extra_layers=None, sub_module=None):
+    if model_arch is None:
+        model_arch = model.model_meta.model_arch
+    # lm_head
+    if model_arch and model_arch.lm_head:
+        output = model_arch.lm_head
+        idx = output.rfind('.')
+        lm_head_name = output[idx + 1:]
+    else:
+        lm_head_name = 'lm_head'
+    # 'score', 'classifier': classification model
+    # 'v_head': reward model
+    ignore_layers = [lm_head_name, 'score', 'v_head', 'classifier'] + ['lora_A', 'lora_B', 'base_layer']
+    ignore_linear_cls = [
+        'glulinear',  # phi4-mm
+        'gemma4clippablelinear',  # gemma4
+    ]
+
+    def _cond(name, module):
+        module_name = module.__class__.__name__.lower()
+        if (extra_layers and isinstance(module, tuple(extra_layers)) or
+            ('linear' in module_name and all(linear_cls not in module_name
+                                             for linear_cls in ignore_linear_cls))) and all(layer not in name
+                                                                                            for layer in ignore_layers):
+            return True
+        return False
+
+    return find_layers(model, _cond, sub_module=sub_module)
+
+
+def get_multimodal_target_regex(
+    model,
+    *,
+    freeze_llm: bool = False,
+    freeze_vit: bool = True,
+    freeze_aligner: bool = True,
+    include_embedding: bool = False,
+    exclude_router: bool = False,
+) -> str:
+    model_arch = model.model_meta.model_arch
+    modules = []
+    if not freeze_llm:
+        modules += model_arch.language_model
+    if not freeze_vit:
+        modules += model_arch.vision_tower
+    if not freeze_aligner:
+        modules += model_arch.aligner
+    assert len(modules) > 0, f'modules: {modules}'
+
+    extra_layers = []
+    if include_embedding:
+        extra_layers.append(nn.Embedding)
+    res = []
+    for module in modules:
+        rejected_modules = []
+        if not freeze_vit or not freeze_llm:
+            for aligner in model_arch.aligner:
+                if aligner.startswith(f'{module}.'):
+                    rejected_modules.append(aligner)
+
+        sub_module = deep_getattr(model, module)
+        if sub_module is None:
+            logger.warning(f'module: {module} is None')
+            continue
+        if isinstance(sub_module, nn.Linear) and module.endswith('lm_head'):
+            target_modules = []
+        else:
+            target_modules = find_all_linears(sub_module, model_arch, extra_layers)
+        if exclude_router and model.model_info.is_moe_model:
+            target_modules = [tm for tm in target_modules if tm not in {'gate'}]
+        if not target_modules:
+            continue
+        target_modules = [tm for tm in target_modules if tm]
+        target_pattern = rf'.*\.({"|".join(target_modules)})' if target_modules else ''
+        rejected_pattern = rf'(?!({"|".join(rejected_modules)}))' if rejected_modules else ''
+        res.append(rf'{rejected_pattern}{re.escape(module)}(?=\.){target_pattern}')
+
+    return rf'^({"|".join(res)})$'
+
+
+def get_cu_seqlens_from_position_ids(position_ids: torch.LongTensor):
+    position_ids = position_ids[0]
+    seq_start_indices = torch.where(position_ids == 0)[0]
+    seq_end_indices = torch.cat([seq_start_indices[1:], torch.tensor([len(position_ids)], device=position_ids.device)])
+    seq_lengths = seq_end_indices - seq_start_indices
+    cu_seqlens = torch.cumsum(torch.cat([torch.tensor([0], device=position_ids.device), seq_lengths]), dim=0)
+    return cu_seqlens.to(torch.int32)
+
+
+def get_position_ids_from_cu_seqlens(cu_seqlens: torch.LongTensor):
+    seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+    position_ids = torch.cat([torch.arange(seq_len, device=cu_seqlens.device) for seq_len in seq_lengths], dim=0)
+    return position_ids.unsqueeze(0)
+
+
+def seed_worker(worker_id: int, num_workers: int, rank: int):
+    """
+    Helper function to set worker seed during Dataloader initialization.
+    """
+    init_seed = torch.initial_seed() % 2**32
+    worker_seed = num_workers * rank + init_seed
+    set_seed(worker_seed)
+
+
+@contextmanager
+def unwrap_model_for_generation(
+    model,
+    accelerator,
+    gather_deepspeed3_params=True,
+    gather_parameters: List[nn.Parameter] = None,
+):
+    unwrapped_model = accelerator.unwrap_model(model)
+    if accelerator.state.deepspeed_plugin is not None and accelerator.state.deepspeed_plugin.zero_stage == 3:
+        if not gather_deepspeed3_params:
+            yield accelerator.unwrap_model(model)
+        else:
+            import deepspeed
+            parameters = [
+                parameter for name, parameter in model.named_parameters()
+                if not gather_parameters or name in gather_parameters
+            ]
+            with deepspeed.zero.GatheredParameters(parameters):
+                from trl.models.utils import add_hooks, remove_hooks
+                remove_hooks(model)
+                yield accelerator.unwrap_model(model)
+                add_hooks(model)
+    else:
+        yield unwrapped_model
+
+
+@contextmanager
+def disable_deepspeed_zero3():
+    import transformers.integrations.deepspeed as ds_module
+    orig_weak_ref = ds_module._hf_deepspeed_config_weak_ref
+    ds_module._hf_deepspeed_config_weak_ref = None
+    try:
+        yield
+    finally:
+        ds_module._hf_deepspeed_config_weak_ref = orig_weak_ref
+
+
+def get_modules_to_not_convert(model):
+    if not hasattr(model, 'model_meta') or not hasattr(model, 'model_info'):
+        return
+    model_arch = model.model_meta.model_arch
+    model_type = model.model_meta.model_type
+    prefix_list = []
+    suffix_list = []
+    if model.model_info.is_moe_model:
+        suffix_list += ['mlp.gate', 'mlp.shared_expert_gate']
+    if model_type in {'qwen3_next', 'qwen3_5', 'qwen3_5_moe'}:
+        suffix_list += ['in_proj_a', 'in_proj_b']
+    if model_arch is not None:
+        for key in ['vision_tower', 'aligner']:
+            value = getattr(model_arch, key, None)
+            if value:
+                prefix_list += value
+    suffix_list.append('lm_head')
+    res = []
+    for n, m in model.named_modules():
+        if 'linear' in m.__class__.__name__.lower() and (any(n.endswith(suffix) for suffix in suffix_list)
+                                                         or any(n.startswith(prefix) for prefix in prefix_list)):
+            res.append(n)
+    return res if res else None
+
+
+def get_packed_seq_params(position_ids: torch.Tensor):
+    assert position_ids.shape[0] == 1, f'position_ids.shape: {position_ids.shape}'
+    position_ids_f = position_ids.flatten()
+    indices_q = torch.arange(position_ids_f.shape[0], device=position_ids_f.device, dtype=torch.int32)
+
+    cu_seqlens = torch.cat([
+        indices_q[position_ids_f == 0],
+        torch.tensor(position_ids_f.shape, device=position_ids_f.device, dtype=torch.int32),
+    ])
+
+    max_length = cu_seqlens.diff().max()  # position_ids_f.max() + 1
+    return {
+        'cu_seq_lens_q': cu_seqlens,
+        'cu_seq_lens_k': cu_seqlens,
+        'max_length_q': max_length,
+        'max_length_k': max_length,
+    }

swift/utils/utils.py

@@ -0,0 +1,544 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import asyncio
+import datetime as dt
+import fnmatch
+import glob
+import importlib
+import json
+import json_repair
+import numpy as np
+import os
+import random
+import re
+import shutil
+import socket
+import subprocess
+import sys
+import threading
+import time
+import torch
+import torch.distributed as dist
+from contextlib import contextmanager
+from functools import wraps
+from transformers import HfArgumentParser, enable_full_determinism, set_seed
+from transformers.utils import strtobool
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
+
+from .env import is_dist, is_master
+from .logger import get_logger
+from .np_utils import stat_array
+
+logger = get_logger()
+
+
+def check_json_format(obj: Any, token_safe: bool = True) -> Any:
+    if obj is None or isinstance(obj, (int, float, str, complex)):  # bool is a subclass of int
+        return obj
+    if isinstance(obj, bytes):
+        return '<<<bytes>>>'
+    if isinstance(obj, (torch.dtype, torch.device)):
+        obj = str(obj)
+        return obj[len('torch.'):] if obj.startswith('torch.') else obj
+
+    if isinstance(obj, Sequence):
+        res = []
+        for x in obj:
+            res.append(check_json_format(x, token_safe))
+    elif isinstance(obj, Mapping):
+        res = {}
+        for k, v in obj.items():
+            if token_safe and isinstance(k, str) and '_token' in k and isinstance(v, str):
+                res[k] = None
+            else:
+                res[k] = check_json_format(v, token_safe)
+    else:
+        if token_safe:
+            unsafe_items = {}
+            for k, v in obj.__dict__.items():
+                if '_token' in k:
+                    unsafe_items[k] = v
+                    setattr(obj, k, None)
+            res = repr(obj)
+            # recover
+            for k, v in unsafe_items.items():
+                setattr(obj, k, v)
+        else:
+            res = repr(obj)  # e.g. function, object
+    return res
+
+
+def _get_version(work_dir: str) -> int:
+    if os.path.isdir(work_dir):
+        fnames = os.listdir(work_dir)
+    else:
+        fnames = []
+    v_list = [-1]
+    for fname in fnames:
+        m = re.match(r'v(\d+)', fname)
+        if m is None:
+            continue
+        v = m.group(1)
+        v_list.append(int(v))
+    return max(v_list) + 1
+
+
+def format_time(seconds):
+    days = int(seconds // (24 * 3600))
+    hours = int((seconds % (24 * 3600)) // 3600)
+    minutes = int((seconds % 3600) // 60)
+    seconds = round(seconds % 60)
+
+    if days > 0:
+        time_str = f'{days}d {hours}h {minutes}m {seconds}s'
+    elif hours > 0:
+        time_str = f'{hours}h {minutes}m {seconds}s'
+    elif minutes > 0:
+        time_str = f'{minutes}m {seconds}s'
+    else:
+        time_str = f'{seconds}s'
+
+    return time_str
+
+
+def deep_getattr(obj, attr: str, default=None):
+    attrs = attr.split('.')
+    for a in attrs:
+        if obj is None:
+            break
+        if isinstance(obj, dict):
+            obj = obj.get(a, default)
+        else:
+            obj = getattr(obj, a, default)
+    return obj
+
+
+def seed_everything(seed: Optional[int] = None, full_determinism: bool = False) -> int:
+
+    if seed is None:
+        seed_max = np.iinfo(np.int32).max
+        seed = random.randint(0, seed_max)
+
+    if full_determinism:
+        enable_full_determinism(seed)
+    else:
+        set_seed(seed)
+    return seed
+
+
+def add_version_to_work_dir(work_dir: str) -> str:
+    """add version"""
+    version = _get_version(work_dir)
+    time = dt.datetime.now().strftime('%Y%m%d-%H%M%S')
+    sub_folder = f'v{version}-{time}'
+    if dist.is_initialized() and is_dist():
+        obj_list = [sub_folder]
+        dist.broadcast_object_list(obj_list)
+        sub_folder = obj_list[0]
+
+    work_dir = os.path.join(work_dir, sub_folder)
+    return work_dir
+
+
+_T = TypeVar('_T')
+
+
+def _patch_args(class_type):
+    try:
+        for k, v in class_type.__annotations__.items():
+            if v == Union[str, dict, type(None)]:
+                class_type.__annotations__[k] = Union[dict, str, type(None)]
+    except Exception:
+        logger.warning('patch args failed')
+
+
+@contextmanager
+def _patch_get_type_hints():
+    # Fix parsing string arguments into dicts
+    from transformers import hf_argparser
+    origin_get_type_hints = hf_argparser.get_type_hints
+
+    def get_type_hints(*args, **kwargs):
+        kwargs = origin_get_type_hints(*args, **kwargs)
+        for k, v in kwargs.items():
+            if v == Union[str, dict, type(None)]:
+                kwargs[k] = Union[dict, str, type(None)]
+        return kwargs
+
+    hf_argparser.get_type_hints = get_type_hints
+    try:
+        yield
+    finally:
+        hf_argparser.get_type_hints = origin_get_type_hints
+
+
+def parse_args(class_type: Type[_T], argv: Optional[List[str]] = None) -> Tuple[_T, List[str]]:
+    with _patch_get_type_hints():
+        parser = HfArgumentParser([class_type])
+    _ray_args = os.environ.get('RAY_SWIFT_ARGS')
+    if _ray_args:
+        argv = json.loads(_ray_args)
+    elif argv is None:
+        argv = sys.argv[1:]
+    args, remaining_args = parser.parse_args_into_dataclasses(argv, return_remaining_strings=True)
+    return args, remaining_args
+
+
+def parse_args_from_dict(class_type: Type[_T], args: Dict[str, Any]) -> _T:
+    with _patch_get_type_hints():
+        parser = HfArgumentParser([class_type])
+    return parser.parse_dict(args, allow_extra_keys=True)[0]
+
+
+def lower_bound(lo: int, hi: int, cond: Callable[[int], bool]) -> int:
+    # The lower bound satisfying the condition "cond".
+    while lo < hi:
+        mid = (lo + hi) >> 1
+        if cond(mid):
+            hi = mid
+        else:
+            lo = mid + 1
+    return lo
+
+
+def upper_bound(lo: int, hi: int, cond: Callable[[int], bool]) -> int:
+    # The upper bound satisfying the condition "cond".
+    while lo < hi:
+        mid = (lo + hi + 1) >> 1  # lo + (hi-lo+1)>>1
+        if cond(mid):
+            lo = mid
+        else:
+            hi = mid - 1
+    return lo
+
+
+def test_time(func: Callable[[], _T],
+              number: int = 1,
+              warmup: int = 0,
+              timer: Optional[Callable[[], float]] = None) -> _T:
+    # timer: e.g. time_synchronize
+    timer = timer if timer is not None else time.perf_counter
+
+    ts = []
+    res = None
+    # warmup
+    for _ in range(warmup):
+        res = func()
+
+    for _ in range(number):
+        t1 = timer()
+        res = func()
+        t2 = timer()
+        ts.append(t2 - t1)
+
+    ts = np.array(ts)
+    _, stat_str = stat_array(ts)
+    # print
+    logger.info(f'time[number={number}]: {stat_str}')
+    return res
+
+
+def read_multi_line(addi_prompt: str = '') -> str:
+    res = []
+    prompt = f'<<<{addi_prompt} '
+    while True:
+        text = input(prompt) + '\n'
+        prompt = ''
+        res.append(text)
+        if text.endswith('#\n'):
+            res[-1] = text[:-2]
+            break
+    return ''.join(res)
+
+
+def subprocess_run(command: List[str], env: Optional[Dict[str, str]] = None, stdout=None, stderr=None):
+    # stdout stderr: e.g. subprocess.PIPE.
+    import shlex
+    command_str = ' '.join(shlex.quote(a) for a in command)
+    logger.info_if(f'Run the command: `{command_str}`', is_master())
+    resp = subprocess.run(command, env=env, stdout=stdout, stderr=stderr)
+    resp.check_returncode()
+    return resp
+
+
+def get_env_args(args_name: str, type_func: Callable[[str], _T], default_value: Optional[_T]) -> Optional[_T]:
+    args_name_upper = args_name.upper()
+    value = os.getenv(args_name_upper)
+    if value is None:
+        value = default_value
+        log_info = (f'Setting {args_name}: {default_value}. '
+                    f'You can adjust this hyperparameter through the environment variable: `{args_name_upper}`.')
+    else:
+        if type_func is bool:
+            value = strtobool(value)
+        value = type_func(value)
+        log_info = f'Using environment variable `{args_name_upper}`, Setting {args_name}: {value}.'
+    logger.info_once(log_info)
+    return value
+
+
+def find_node_ip() -> Optional[str]:
+    import psutil
+    main_ip, virtual_ip = None, None
+    for name, addrs in sorted(psutil.net_if_addrs().items()):
+        for addr in addrs:
+            if addr.family.name == 'AF_INET' and not addr.address.startswith('127.'):
+                # Heuristic to prefer non-virtual interfaces
+                if any(s in name for s in ['lo', 'docker', 'veth', 'vmnet']):
+                    if virtual_ip is None:
+                        virtual_ip = addr.address
+                else:
+                    if main_ip is None:
+                        main_ip = addr.address
+    return main_ip or virtual_ip
+
+
+def find_free_port(start_port: Optional[int] = None, retry: int = 100) -> int:
+    if start_port is None:
+        start_port = 0
+    for port in range(start_port, start_port + retry):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            try:
+                sock.bind(('', port))
+                port = sock.getsockname()[1]
+                break
+            except OSError:
+                pass
+    return port
+
+
+def copy_files_by_pattern(source_dir, dest_dir, patterns, exclude_patterns=None):
+    if not os.path.exists(dest_dir):
+        os.makedirs(dest_dir)
+
+    if isinstance(patterns, str):
+        patterns = [patterns]
+
+    if exclude_patterns is None:
+        exclude_patterns = []
+    elif isinstance(exclude_patterns, str):
+        exclude_patterns = [exclude_patterns]
+
+    def should_exclude_file(file_path, file_name):
+        for exclude_pattern in exclude_patterns:
+            if fnmatch.fnmatch(file_name, exclude_pattern):
+                return True
+            rel_file_path = os.path.relpath(file_path, source_dir)
+            if fnmatch.fnmatch(rel_file_path, exclude_pattern):
+                return True
+        return False
+
+    for pattern in patterns:
+        pattern_parts = pattern.split(os.path.sep)
+        if len(pattern_parts) > 1:
+            subdir_pattern = os.path.sep.join(pattern_parts[:-1])
+            file_pattern = pattern_parts[-1]
+
+            for root, dirs, files in os.walk(source_dir):
+                rel_path = os.path.relpath(root, source_dir)
+                if rel_path == '.' or (rel_path != '.' and not fnmatch.fnmatch(rel_path, subdir_pattern)):
+                    continue
+
+                for file in files:
+                    if fnmatch.fnmatch(file, file_pattern):
+                        file_path = os.path.join(root, file)
+
+                        if should_exclude_file(file_path, file):
+                            continue
+
+                        target_dir = os.path.join(dest_dir, rel_path)
+                        if not os.path.exists(target_dir):
+                            os.makedirs(target_dir)
+                        dest_file = os.path.join(target_dir, file)
+
+                        if not os.path.exists(dest_file):
+                            shutil.copy2(file_path, dest_file)
+        else:
+            search_path = os.path.join(source_dir, pattern)
+            matched_files = glob.glob(search_path)
+
+            for file_path in matched_files:
+                if os.path.isfile(file_path):
+                    file_name = os.path.basename(file_path)
+
+                    if should_exclude_file(file_path, file_name):
+                        continue
+
+                    destination = os.path.join(dest_dir, file_name)
+                    if not os.path.exists(destination):
+                        shutil.copy2(file_path, destination)
+
+
+def split_list(ori_list: List[_T], num_shards: int, contiguous=True) -> List[List[_T]]:
+    shard = []
+    if contiguous:
+        idx_list = np.linspace(0, len(ori_list), num_shards + 1, dtype=np.int64)
+        for i in range(len(idx_list) - 1):
+            shard.append(ori_list[idx_list[i]:idx_list[i + 1]])
+    else:
+        ori_list = np.array(ori_list)
+        for i in range(num_shards):
+            shard.append(ori_list[np.arange(i, len(ori_list), num_shards)].tolist())
+    return shard
+
+
+def patch_getattr(obj_cls, item_name: str):
+    if hasattr(obj_cls, '_patch'):  # avoid double patch
+        return
+
+    def __new_getattr__(self, key: str):
+        try:
+            return super(self.__class__, self).__getattr__(key)
+        except AttributeError:
+            if item_name in dir(self):
+                item = getattr(self, item_name)
+                return getattr(item, key)
+            raise
+
+    obj_cls.__getattr__ = __new_getattr__
+    obj_cls._patch = True
+
+
+def import_external_file(file_path: str):
+    file_path = os.path.abspath(os.path.expanduser(file_path))
+    py_dir, py_file = os.path.split(file_path)
+    assert os.path.isdir(py_dir), f'py_dir: {py_dir}'
+    sys.path.insert(0, py_dir)
+    return importlib.import_module(py_file.split('.', 1)[0])
+
+
+def json_parse_to_dict(value: Union[str, Dict, None], strict: bool = True) -> Union[str, Dict]:
+    """Convert a JSON string or JSON file into a dict"""
+    # If the value could potentially be a string, it is generally advisable to set strict to False.
+    if value is None:
+        value = {}
+    elif isinstance(value, str):
+        if os.path.exists(value):  # local path
+            with open(value, 'r', encoding='utf-8') as f:
+                value = json.load(f)
+        else:  # json str
+            try:
+                value = json.loads(value)
+            except json.JSONDecodeError:
+                if strict:
+                    try:
+                        # fix malformed json string, e.g., incorrect quotation marks
+                        old_value = value
+                        value = json_repair.repair_json(value)
+                        logger.warning(f'Unable to parse json string, try to repair it, '
+                                       f"the string before and after repair are '{old_value}' | '{value}'")
+                        value = json.loads(value)
+                    except Exception:
+                        logger.error(f"Unable to parse json string: '{value}', and try to repair failed")
+                        raise
+    return value
+
+
+def retry_decorator(retry=3):
+
+    def _retry(func):
+
+        @wraps(func)
+        def new_func(*args, **kwargs):
+            i = 1
+            while True:
+                try:
+                    return func(*args, **kwargs)
+                except Exception:
+                    if i == retry:
+                        raise
+                    i += 1
+
+        return new_func
+
+    return _retry
+
+
+def start_event_loop_in_daemon(name: str = None) -> Tuple[threading.Thread, asyncio.AbstractEventLoop, threading.Event]:
+    """Create a new daemon thread that runs an asyncio event loop.
+
+    Args:
+        name: Name of the thread. If None, the default thread naming will be used.
+
+    Returns:
+        tuple: (thread, loop, loop_ready_event)
+            - thread: The thread running the event loop.
+            - loop: The event loop being run in the thread.
+            - loop_ready_event: An event that is set when the loop is ready.
+    """
+    loop = asyncio.new_event_loop()
+    loop_ready_event = threading.Event()
+
+    def run_loop():
+        asyncio.set_event_loop(loop)
+        loop_ready_event.set()
+        loop.run_forever()
+
+    thread = threading.Thread(target=run_loop, name=name, daemon=True)
+    thread.start()
+    return thread, loop, loop_ready_event
+
+
+def shutdown_event_loop_in_daemon(thread: threading.Thread = None, loop: asyncio.AbstractEventLoop = None) -> None:
+    """Shutdown an asyncio event loop running in a separate thread.
+
+    This function stops the event loop and waits for the associated thread to finish execution.
+
+    Args:
+        thread: The thread running the event loop.
+        loop: The asyncio event loop to shut down.
+    """
+    if loop is None or thread is None:
+        return
+    loop.call_soon_threadsafe(loop.stop)
+    thread.join(timeout=5)
+
+
+def remove_response(messages) -> Optional[str]:
+    """
+    Removes and returns the content of the last message if its role is 'assistant'.
+
+    Args:
+        messages (List[Dict]):
+            A list of message dictionaries, each typically containing a 'role' and 'content' key.
+
+    Returns:
+        Optional[str]:
+            The content of the removed 'assistant' message if present;
+            otherwise, returns None. The original messages list is modified in place.
+    """
+    last_role = messages[-1]['role'] if messages else None
+    if last_role == 'assistant':
+        return messages.pop()['content']
+
+
+def to_abspath(path: Union[str, List[str], None], check_path_exist: bool = False) -> Union[str, List[str], None]:
+    """Check the path for validity and convert it to an absolute path.
+
+    Args:
+        path: The path to be checked/converted
+        check_path_exist: Whether to check if the path exists
+
+    Returns:
+        Absolute path
+    """
+    if path is None:
+        return
+    elif isinstance(path, str):
+        # Remove user path prefix and convert to absolute path.
+        path = os.path.abspath(os.path.expanduser(path))
+        if check_path_exist and not os.path.exists(path):
+            raise FileNotFoundError(f"path: '{path}'")
+        return path
+    assert isinstance(path, list), f'path: {path}'
+    res = []
+    for v in path:
+        res.append(to_abspath(v, check_path_exist))
+    return res
+
+
+def swanlab_get_run():
+    try:
+        import swanlab
+        return swanlab.get_run()
+    except (RuntimeError, ImportError):
+        return None