# Copyright (c) Alibaba, Inc. and its affiliates. # Copyright 2023-present the HuggingFace Inc. team. import os import re import shutil import tempfile from contextlib import contextmanager from copy import copy from functools import partial from inspect import Parameter, Signature, signature from types import MethodType from typing import Dict, List, Literal, Optional, Union import json import torch from modelscope import snapshot_download from peft.utils import CONFIG_NAME from peft.utils.other import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME from torch import nn from transformers import Trainer from swift.utils.constants import DEFAULT_ADAPTER, SWIFT_TYPE_KEY from swift.utils.logger import get_logger from ..utils.torch_utils import get_device_count from .mapping import SwiftTuners from .peft import PeftConfig, PeftModel, get_peft_model from .utils import SwiftConfig, SwiftOutput logger = get_logger() class SwiftModel(nn.Module): """The Swift wrapper model. Args: model (`Union[nn.Module, 'SwiftModel']`) A module to be tuned by Swift. config (`Union[SwiftConfig, Dict[str, SwiftConfig]]`) A config or a dict of {adapter_name: SwiftConfig}. If it's a config class, the adapter_name will be `default` extra_state_keys (`List[str]`, `optional`) A list of regex to match the extra state keys to be saved. inference_mode (bool, `optional`): Load model at inference mode, default False. """ EXTRA_STATE_DIR = 'extra_states' def __init__(self, model: Union[nn.Module, 'SwiftModel'], config: Union[SwiftConfig, Dict[str, SwiftConfig]], extra_state_keys: List[str] = None, inference_mode: bool = False, **kwargs): super().__init__() self.adapters = {} self.active_adapters = set() if isinstance(model, SwiftModel): self.adapters = model.adapters extra_state_keys = extra_state_keys or [] extra_state_keys.extend(model.extra_state_keys) self.active_adapters = model.active_adapters model = model.base_model self.base_model = model new_adapters = [] if isinstance(config, SwiftConfig): if DEFAULT_ADAPTER not in self.adapters: all_parts = self._deactivate_all_parts() self.adapters[DEFAULT_ADAPTER] = self._prepare_model(model, config, DEFAULT_ADAPTER) for part in all_parts: self.activate_adapter(part) new_adapters.append(DEFAULT_ADAPTER) if self.adapters[DEFAULT_ADAPTER].model is not None: self.base_model = self.adapters[DEFAULT_ADAPTER].model else: logger.warn(f'Adapter {DEFAULT_ADAPTER} has been patched, skip.') elif isinstance(config, dict): assert (all(isinstance(c, SwiftConfig) for c in config.values())) for adapter_name, _config in config.items(): if adapter_name not in self.adapters: all_parts = self._deactivate_all_parts() self.adapters[adapter_name] = self._prepare_model(model, _config, adapter_name) for part in all_parts: self.activate_adapter(part) new_adapters.append(adapter_name) if self.adapters[adapter_name].model is not None: self.base_model = self.adapters[adapter_name].model else: logger.warn(f'Adapter {adapter_name} has been patched, skip.') self.extra_state_keys = extra_state_keys or [] self.has_additional_modules = any([c.config.has_additional_modules for c in self.adapters.values()]) def forward(self, *args, **kwargs): return self.base_model(*args, **kwargs) _parameters = [Parameter('self', Parameter.POSITIONAL_ONLY)] _parameters += list(signature(self.base_model.forward).parameters.values()) forward.__signature__ = Signature(_parameters) self.forward = MethodType(forward, self) for adapter_name in new_adapters: self.activate_adapter(adapter_name) if inference_mode: self.eval() else: for key, output in self.adapters.items(): if key in new_adapters: output.mark_trainable_callback(model) if self.extra_state_keys: for n, p in model.named_parameters(): if any(re.fullmatch(extra_key, n) for extra_key in self.extra_state_keys): p.requires_grad = True @property def model(self): return self.base_model def _deactivate_all_parts(self): deactivated = [] for adapter in self.active_adapters: output = self.adapters[adapter] if output.config.swift_type == SwiftTuners.PART: deactivated.append(adapter) self.deactivate_adapter(adapter) return deactivated def load_state_dict(self, state_dict, strict=True, adapter_name: str = None): if adapter_name is not None: output: SwiftOutput = self.adapters[adapter_name] if getattr(output.config, 'modules_to_save', None): for key, value in copy(state_dict).items(): for module_name in output.config.modules_to_save: if module_name in key: state_dict.pop(key) key = key.replace(module_name, f'{module_name}.modules_to_save.{adapter_name}') break state_dict[key] = value for key, value in copy(state_dict).items(): if key.startswith('base_model.model.'): state_dict.pop(key, None) key = key[len('base_model.model.'):] if f'lora_A.{adapter_name}.' not in key and 'lora_A' in key: state_dict.pop(key, None) key = key.replace('lora_A.', f'lora_A.{adapter_name}.') if f'lora_B.{adapter_name}.' not in key and 'lora_B' in key: state_dict.pop(key, None) key = key.replace('lora_B.', f'lora_B.{adapter_name}.') if f'lora_embedding_A.{adapter_name}.' not in key and 'lora_embedding_A' in key: state_dict.pop(key, None) key = key.replace('lora_embedding_A.', f'lora_embedding_A.{adapter_name}.') if f'lora_embedding_B.{adapter_name}.' not in key and 'lora_embedding_B' in key: state_dict.pop(key, None) key = key.replace('lora_embedding_B.', f'lora_embedding_B.{adapter_name}.') state_dict[key] = value if output.load_state_dict_callback: state_dict = output.load_state_dict_callback(self.base_model, adapter_name, state_dict) incompatible_keys = self.base_model.load_state_dict(state_dict, False) if incompatible_keys and len(incompatible_keys[1]) > 0: logger.error(f'Load state dict with unexpected keys: {incompatible_keys[1]}') def state_dict(self, *args, destination=None, prefix='', keep_vars=False, adapter_name: str = None, peft_format: bool = False, **kwargs): """ Args: destination (`dict`, `optional`): If provided, the state of module will be updated into the dict and the same object is returned. Otherwise, an ``OrderedDict`` will be created and returned. Default: ``None``. prefix (`str`, `optional`): a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ``''``. keep_vars (`bool`, `optional`): by default the :class:`~torch.Tensor` s returned in the state dict are detached from autograd. If it's set to ``True``, detaching will not be performed. Default: ``False``. adapter_name (`str`, `optional`): The name of the adapter's parameters to be saved, `None` input will save all adapters. peft_format (`bool`, `optional`): Save with peft format (extra `base_model.model.` prefix) **kwargs: save_adapter(`bool`): Save adapters or not, default True save_extra_states(`bool`): Save extra states or not, default True Returns: The state dict to be saved. """ state_dict = kwargs.get('state_dict') if state_dict is None: state_dict = self.base_model.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) state_dict = { key[len('base_model.'):] if key.startswith('base_model.') else key: value for key, value in state_dict.items() } if not self.has_additional_modules: return state_dict state_dicts = {} if kwargs.get('save_adapter', True): for name, output in self.adapters.items(): if (adapter_name == name or adapter_name is None) and output.config.has_additional_modules: # noqa state_dicts.update(output.state_dict_callback(state_dict, name)) modules_to_save_names = [ sub_name for sub_name, _ in self.base_model.named_parameters() if f'modules_to_save.{name}' in sub_name ] for module_name in modules_to_save_names: if f'modules_to_save.{name}' in module_name: state_dicts[module_name.replace(f'modules_to_save.{name}.', '')] = state_dict[module_name] if kwargs.get('save_extra_states', True): state_dicts.update({ k: v for k, v in state_dict.items() if any( re.fullmatch(extra_key, k) for extra_key in self.extra_state_keys) }) if peft_format: new_state_dict = {} for key, value in state_dicts.items(): if not key.startswith('base_model.model.'): key = 'base_model.model.' + key key = key.replace(f'lora_A.{adapter_name}.', 'lora_A.') key = key.replace(f'lora_B.{adapter_name}.', 'lora_B.') key = key.replace(f'lora_embedding_A.{adapter_name}.', 'lora_embedding_A.') key = key.replace(f'lora_embedding_B.{adapter_name}.', 'lora_embedding_B.') new_state_dict[key] = value state_dicts = new_state_dict return state_dicts def __getattr__(self, key: str): """Forward missing attributes to the wrapped module.""" try: return super().__getattr__(key) except AttributeError: if 'base_model' in dir(self): return getattr(self.base_model, key) raise @staticmethod def load_state_file(path, device: Optional[str] = None): """Load a state dict file by the input path. Args: path: The local dir to load the state file. Returns: The state dict. """ if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' if os.path.exists(os.path.join(path, SAFETENSORS_WEIGHTS_NAME)): filename = os.path.join(path, SAFETENSORS_WEIGHTS_NAME) from safetensors.torch import load_file as safe_load_file return safe_load_file(filename, device=device) elif os.path.exists(os.path.join(path, WEIGHTS_NAME)): filename = os.path.join(path, WEIGHTS_NAME) return torch.load(filename, map_location=device) return None def create_optimizer_param_groups(self, **defaults): all_param_names = set() param_groups = [] for output in self.adapters.values(): if output.optimizer_group_callback: param_names, param_group = output.optimizer_group_callback(self.model, **defaults) if param_names and all_param_names & param_names: raise ValueError('Cannot set one parameter to different param groups') if param_names and param_group: all_param_names.update(param_names) param_groups.extend(param_group) decay_parameters = Trainer.get_decay_parameter_names(None, self.model) param_groups.extend([ { 'params': [ p for n, p in self.model.named_parameters() if (n in decay_parameters and n not in all_param_names and p.requires_grad) ], 'weight_decay': defaults['weight_decay'], }, { 'params': [ p for n, p in self.model.named_parameters() if (n not in decay_parameters and n not in all_param_names and p.requires_grad) ], 'weight_decay': 0.0, }, ]) return param_groups @classmethod def from_pretrained(cls, model: Union[nn.Module, 'SwiftModel'], model_id: str = None, adapter_name: Union[str, List[str], Dict[str, str]] = None, inference_mode: bool = True, revision: str = None, **kwargs): """Load a set of tuners and corresponding weights by a model_id. Args: model (`Union[torch.nn.Module, 'SwiftModel']`): The model to be tuned, if the model is already a `SwiftModel` it will be un-wrapped and re-wrapped.. model_id (`str`): The model_id or a local model dir of tuners to use to tune the model. adapter_name (`Union[str, List[str], Dict[str, str]]`): The adapter_names saved in the model repo to load. Default `None`, means load all tuners saved in the model_id inference_mode (`bool`): Use in the inference mode or not. revision (`str`): The model revision to use. **kwargs: extra_state_keys (`List[str]`, `optional`) A list of regex to match the extra state keys to be saved. Other parameters will be passed to the device_map. Returns: The `SwiftModel` instance. """ adapters = {} model_dir = model_id if not os.path.exists(model_dir): model_dir = snapshot_download(model_dir, revision=revision) if os.path.isfile(model_dir): raise ValueError(f'Please pass in a local dir or a model id, not a local file: {model_dir}') extra_state_keys = kwargs.pop('extra_state_keys', None) if extra_state_keys is None and os.path.isfile(os.path.join(model_dir, cls.EXTRA_STATE_DIR, CONFIG_NAME)): with open(os.path.join(model_dir, cls.EXTRA_STATE_DIR, CONFIG_NAME), 'r', encoding='utf-8') as file: _json = json.load(file) extra_state_keys = _json.get('extra_state_keys') if adapter_name is None: adapter_name = [ sub_dir for sub_dir in os.listdir(model_dir) if os.path.isfile(os.path.join(model_dir, sub_dir, CONFIG_NAME)) and sub_dir != cls.EXTRA_STATE_DIR ] for _name in adapter_name if isinstance(adapter_name, list) else [adapter_name] \ if isinstance(adapter_name, str) else adapter_name.keys(): sub_folder = os.path.join(model_dir, _name) config_file = os.path.join(sub_folder, CONFIG_NAME) if not os.path.isfile(config_file): logger.warning(f'{_name} is not a valid tuner') continue with open(config_file, 'r', encoding='utf-8') as file: json_object = json.load(file) if SWIFT_TYPE_KEY not in json_object: raise ValueError('Mixed using with peft is not allowed now.') else: key = _name if not isinstance(adapter_name, dict) else adapter_name[_name] adapters[key] = SwiftConfig.from_pretrained(sub_folder) self = SwiftModel(model, adapters, extra_state_keys, inference_mode, **kwargs) for _name in adapter_name if isinstance(adapter_name, list) else [adapter_name] \ if isinstance(adapter_name, str) else adapter_name.keys(): _adapter = _name if not isinstance(adapter_name, dict) else adapter_name[_name] output: SwiftOutput = self.adapters[_adapter] sub_folder = os.path.join(model_dir, _name) if output.load_callback: output.load_callback(self, sub_folder, _adapter) continue state_dict = cls.load_state_file(sub_folder) if state_dict is not None: if isinstance(adapter_name, dict): # TODO this logic is fragile! replace `_name` may cause other parts replaced state_dict = {key.replace(_name, adapter_name[_name]): value for key, value in state_dict.items()} self.load_state_dict(state_dict, adapter_name=_adapter) state_dict = cls.load_state_file(os.path.join(model_dir, self.EXTRA_STATE_DIR)) if state_dict is not None: self.load_state_dict(state_dict) return self @classmethod def _prepare_model( cls, model: nn.Module, config: SwiftConfig, adapter_name: str, ): assert (hasattr(config, SWIFT_TYPE_KEY)) from .mapping import SWIFT_MAPPING adapter_cls = SWIFT_MAPPING[config.swift_type][1] if adapter_cls.has_additional_modules() and not getattr(model, 'model_frozen', False): for _, p in model.named_parameters(): p.requires_grad = False model.model_frozen = True config.has_additional_modules = adapter_cls.has_additional_modules() return adapter_cls.prepare_model(model, config, adapter_name) def create_or_update_model_card(self, output_dir: str): """ Updates or create the model card. """ if not os.path.exists(os.path.join(output_dir, 'README.md')): lines = [] else: with open(os.path.join(output_dir, 'README.md'), 'r', encoding='utf-8') as f: lines = f.readlines() quantization_config = None if hasattr(self.base_model, 'config') and hasattr(self.base_model.config, 'quantization_config'): if hasattr(self.base_model.config.quantization_config, 'to_dict'): quantization_config = self.base_model.config.quantization_config.to_dict() training_config_text = '' # Adds quantization information if it was used if quantization_config is not None: training_config_text += '\nThe following `bitsandbytes` quantization config was used during training:\n' training_config_text += '\n'.join([f'- {name}: {value}' for name, value in quantization_config.items()]) training_config_text += '\n' training_procedure_heading = '## Training procedure\n' if training_procedure_heading in lines: lines.insert(lines.index(training_procedure_heading) + 2, training_config_text) else: lines.append(f'{training_procedure_heading}\n{training_config_text}') framework_block_heading = '### Framework versions\n' from swift.version import __version__ if framework_block_heading in lines: lines.insert(lines.index(framework_block_heading) + 2, f'- SWIFT {__version__}\n') else: lines.append(f'{framework_block_heading}\n\n- SWIFT {__version__}\n') base_model_heading = '### Base model information\n' lines.append(f'{base_model_heading}\n\n- BaseModel Class {self.base_model.__class__.__name__}\n') # write the lines back to README.md with open(os.path.join(output_dir, 'README.md'), 'w', encoding='utf-8') as f: f.writelines(lines) def add_weighted_adapter( self, adapters, weights, adapter_name, combination_type='svd', svd_rank=None, svd_clamp=None, svd_full_matrices=True, svd_driver=None, density=None, majority_sign_method: Literal['total', 'frequency'] = 'total', ): """ This method adds a new adapter by merging the given adapters with the given weights. When using the `cat` combination_type you should be aware that rank of the resulting adapter will be equal to the sum of all adapters ranks. So it's possible that the mixed adapter may become too big and result in OOM errors. Args: adapters (`list`): List of adapter names to be merged. weights (`list`): List of weights for each adapter. adapter_name (`str`): Name of the new adapter. combination_type (`str`): The merging type can be one of [`svd`, `linear`, `cat`, `ties`, `ties_svd`, `dare_ties`, `dare_linear`, `dare_ties_svd`, `dare_linear_svd`, `magnitude_prune`, `magnitude_prune_svd`]. When using the `cat` combination_type, the rank of the resulting adapter is equal to the sum of all adapters ranks (the mixed adapter may be too big and result in OOM errors). svd_rank (`int`, *optional*): Rank of output adapter for svd. If None provided, will use max rank of merging adapters. svd_clamp (`float`, *optional*): A quantile threshold for clamping SVD decomposition output. If None is provided, do not perform clamping. Defaults to None. svd_full_matrices (`bool`, *optional*): Controls whether to compute the full or reduced SVD, and consequently, the shape of the returned tensors U and Vh. Defaults to True. svd_driver (`str`, *optional*): Name of the cuSOLVER method to be used. This keyword argument only works when merging on CUDA. Can be one of [None, `gesvd`, `gesvdj`, `gesvda`]. For more info please refer to `torch.linalg.svd` documentation. Defaults to None. density (`float`, *optional*): Value between 0 and 1. 0 means all values are pruned and 1 means no values are pruned. Should be used with [`ties`, `ties_svd`, `dare_ties`, `dare_linear`, `dare_ties_svd`, `dare_linear_svd`, `magnintude_prune`, `magnitude_prune_svd`] majority_sign_method (`str`): The method, should be one of ["total", "frequency"], to use to get the magnitude of the sign values. Should be used with [`ties`, `ties_svd`, `dare_ties`, `dare_ties_svd`] """ from swift.tuners.lora import LoraModel lora_model = LoraModel(self.model, None, '') lora_model.peft_config = {key: value.config for key, value in self.adapters.items()} from peft.tuners.lora import LoraLayer lora_model.targeted_module_names = [ key for key, value in self.model.named_modules() if isinstance(value, LoraLayer) ] lora_model.active_adapter = self.active_adapters lora_model.add_weighted_adapter( adapters=adapters, weights=weights, adapter_name=adapter_name, combination_type=combination_type, svd_rank=svd_rank, svd_clamp=svd_clamp, svd_full_matrices=svd_full_matrices, svd_driver=svd_driver, density=density, majority_sign_method=majority_sign_method, ) def state_dict_callback(state_dict, adapter_name, cfg): from swift.tuners.lora_layers import lora_state_dict return lora_state_dict(state_dict, adapter_name, cfg.bias) def mark_trainable_callback(model, cfg): from swift.tuners.lora_layers import mark_lora_as_trainable mark_lora_as_trainable(model, adapter_name, cfg.bias) cfg = lora_model.peft_config[adapter_name] cfg.has_additional_modules = True self.adapters[adapter_name] = SwiftOutput( config=cfg, state_dict_callback=partial(state_dict_callback, cfg=cfg), mark_trainable_callback=partial(mark_trainable_callback, cfg=cfg), optimizer_group_callback=None, ) self.set_active_adapters(adapter_name) def save_pretrained(self, save_directory: str, safe_serialization: bool = False, adapter_name: Union[str, List[str]] = None, **kwargs): """Save the adapters to a local directory. Args: save_directory (`str`): The directory to use. safe_serialization (`bool`): Use safe tensors to save the weights, default False. adapter_name(`Union[str, List[str]]`): The adapters to be saved, default is `None` to save all. """ peft_format = kwargs.pop('peft_format', False) if os.path.isfile(save_directory): raise ValueError(f'Provided path ({save_directory}) should be a directory, not a file') os.makedirs(save_directory, exist_ok=True) if not self.has_additional_modules: if hasattr(self.base_model, 'save_pretrained'): self.base_model.save_pretrained(save_directory, safe_serialization=safe_serialization) else: self._save_state_dict(self.base_model.state_dict(), save_directory, safe_serialization) self.create_or_update_model_card(save_directory) else: self.create_or_update_model_card(save_directory) adapter_names = adapter_name if isinstance(adapter_name, list) or adapter_name is None else [adapter_name] state_dict_kwargs = {} state_dict = kwargs.get('state_dict') if state_dict is not None: state_dict_kwargs['state_dict'] = kwargs['state_dict'] for adapter_name, output in self.adapters.items(): if adapter_names is not None and adapter_name not in adapter_names: continue save_to_peft = peft_format and output.config.swift_type == SwiftTuners.LORA save_to_peft = save_to_peft and output.config.can_be_saved_to_peft() if peft_format and not save_to_peft: logger.error('You are using additional lora parameters, which is not compatible with peft,' 'which is unable to save to peft format.') output_dir = os.path.join(save_directory, adapter_name) if adapter_name != 'default' or not save_to_peft else save_directory if save_to_peft: config = output.config.to_peft_config() config.save_pretrained(output_dir) else: output.config.save_pretrained(output_dir) if output.save_callback: output.save_callback(self, output_dir, adapter_name) continue # save only the trainable weights output_state_dict = self.state_dict( adapter_name=adapter_name, save_extra_states=False, peft_format=save_to_peft, **state_dict_kwargs) os.makedirs(output_dir, exist_ok=True) if output_state_dict and output.config.has_additional_modules: self._save_state_dict(output_state_dict, output_dir, safe_serialization) output_state_dict = self.state_dict(save_extra_states=True, save_adapter=False, **state_dict_kwargs) if len(output_state_dict) > 0: if self.has_additional_modules: os.makedirs(os.path.join(save_directory, self.EXTRA_STATE_DIR), exist_ok=True) self._save_state_dict(output_state_dict, os.path.join(save_directory, self.EXTRA_STATE_DIR), safe_serialization) with open( os.path.join(save_directory, self.EXTRA_STATE_DIR, CONFIG_NAME), 'w', encoding='utf-8') as file: json.dump({'extra_state_keys': self.extra_state_keys}, file) else: logger.error('Full parameter training, save_extra_states will be ignored') if not os.path.exists(os.path.join(save_directory, 'configuration.json')): with open(os.path.join(save_directory, 'configuration.json'), 'w', encoding='utf-8') as f: f.write('{}') @staticmethod def _save_state_dict(output_state_dict, save_directory, safe_serialization): if safe_serialization: from safetensors.torch import save_file as safe_save_file safe_save_file( output_state_dict, os.path.join(save_directory, SAFETENSORS_WEIGHTS_NAME), metadata={'format': 'pt'}) else: torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME)) @contextmanager def disable_adapter(self): try: self.set_active_adapters(adapter_names=[]) yield finally: self.set_active_adapters(adapter_names=self.adapters.keys()) def set_active_adapters(self, adapter_names: Union[List[str], str], offload: str = None): """Set activated adapters Args: adapter_names(`Union[List[str], str]`): The adapters needed to be activated offload(`str`): Whether to offload the deactivated ones to `cpu` or `meta` device """ if not adapter_names: adapter_names = [] if isinstance(adapter_names, str): adapter_names = [adapter_names] adapter_names = set(adapter_names) for adapter_name in (adapter_names & set(self.adapters.keys())): self.activate_adapter(adapter_name) for adapter_name in (set(self.adapters.keys()) - adapter_names): self.deactivate_adapter(adapter_name, offload) self.active_adapters = (adapter_names & set(self.adapters.keys())) def activate_adapter(self, adapter_name: str): """Activate one adapter Args: adapter_name(`str`): The adapter needed to be activated """ if adapter_name not in self.adapters: logger.warning(f'{adapter_name} not in adapters: {self.adapters.keys()}') return from .mapping import SWIFT_MAPPING SWIFT_MAPPING[self.adapters[adapter_name].config.swift_type][1]\ .activate_adapter(self.base_model, adapter_name, True) self.active_adapters = self.active_adapters | {adapter_name} def deactivate_adapter(self, adapter_name: str, offload: str = None): """Deactivate one adapter Args: adapter_name(`str`): The adapter needed to be activated offload(`str`): Whether to offload to `cpu` or `meta` device """ if adapter_name not in self.adapters: logger.warning(f'{adapter_name} not in adapters: {self.adapters.keys()}') return from .mapping import SWIFT_MAPPING SWIFT_MAPPING[self.adapters[adapter_name].config.swift_type][1]\ .activate_adapter(self.base_model, adapter_name, False, offload=offload) self.active_adapters = self.active_adapters - {adapter_name} def get_trainable_parameters(self): """ Get the content of trainable parameters in the model. """ trainable_params = 0 all_param = 0 for _, param in self.base_model.named_parameters(): num_params = param.numel() # if using DS Zero 3 and the weights are initialized empty if num_params == 0 and hasattr(param, 'ds_numel'): num_params = param.ds_numel all_param += num_params if param.requires_grad: trainable_params += num_params return f'trainable params: {trainable_params:,d} || all params: {all_param:,d} ' \ f'|| trainable%: {100 * trainable_params / all_param:.4f}' \ '|| cuda memory: ' \ f'{sum([torch.cuda.memory_allocated(i) for i in range(get_device_count())])/1024/1024/1024:.2f}' \ 'GiB.' class Swift: """The Wrapper to use both Peft and Swift tuners.""" @staticmethod def prepare_model(model: Union[nn.Module, SwiftModel], config: Union[SwiftConfig, PeftConfig, Dict[str, SwiftConfig]], **kwargs): """Prepare a model by the input config. Args: model(`Union[nn.Module, 'SwiftModel']`): The model to be tuned. config(`Union[SwiftConfig, PeftConfig, Dict[str, SwiftConfig]]`): The config or config dict, can be either SwiftConfigs or PeftConfigs **kwargs: Extra kwargs needed by SwiftModel or PeftModel. Returns: The model wrapped by SwiftModel or PeftModel. """ if isinstance(config, (SwiftConfig, dict)): return SwiftModel(model, config, **kwargs) else: return get_peft_model(model, config, **kwargs) @staticmethod def merge_and_unload(model: Union[PeftModel, SwiftModel], **kwargs): """Merge tuners into the base model and unload them. Args: model(`Union[PeftModel, SwiftModel]`): The model instance with tuners kwargs: adapter_name(`Union[str, List[str]]`): The adapter_name to unload, only supported in swift tuners. """ from peft import PeftModel as _PeftModel if isinstance(model, _PeftModel): model.merge_and_unload() elif isinstance(model, SwiftModel): from swift import LoRAConfig from swift.tuners import LoRA adapter_name = kwargs.get('adapter_name', None) if isinstance(adapter_name, str): adapter_name = [adapter_name] for adapter, output in model.adapters.items(): if isinstance(output.config, LoRAConfig) and (adapter_name is None or adapter in adapter_name): LoRA.unpatch_lora(model, output.config, adapter) @staticmethod @contextmanager def grpo_context(model: Union[SwiftModel, torch.nn.Module], processor): # Save the model and temporarily modify model.model_dir. if not isinstance(model, SwiftModel): yield return else: assert len(model.adapters) == 1 adapter = list(model.adapters.values())[0] if adapter.config.swift_type == SwiftTuners.LLAMAPRO: from modelscope.hub.utils.utils import get_cache_dir temp_dir = tempfile.mkdtemp(dir=get_cache_dir()) model_dir = model.model_dir from transformers.integrations import is_deepspeed_zero3_enabled if is_deepspeed_zero3_enabled(): raise ValueError('DeepSpeed ZeRO3 not supported for LLaMAPro&GRPO currently.') model.base_model.save_pretrained(temp_dir) processor.save_pretrained(temp_dir) model.model_dir = temp_dir yield if adapter.config.swift_type == SwiftTuners.LLAMAPRO: model.model_dir = model_dir shutil.rmtree(temp_dir) @staticmethod def merge(model: Union[PeftModel, SwiftModel], **kwargs): """Merge tuners into the base model, will not unload them. Args: model(`Union[PeftModel, SwiftModel]`): The model instance with tuners """ from .lora_layers import LoraLayer, LoRALayer for sub_module in model.modules(): if isinstance(sub_module, (LoraLayer, LoRALayer)): sub_module.merge(**kwargs) @staticmethod def unmerge(model: Union[PeftModel, SwiftModel], **kwargs): """Unmerge tuners from the base model Args: model(`Union[PeftModel, SwiftModel]`): The model instance with tuners """ from .lora_layers import LoraLayer, LoRALayer for sub_module in model.modules(): if isinstance(sub_module, (LoraLayer, LoRALayer)): sub_module.unmerge(**kwargs) @staticmethod def save_to_peft_format(ckpt_dir: str, output_dir: str) -> None: """Save swift format to peft format Args: ckpt_dir(`str`): Original swift output dir output_dir(`str`): Converted peft format dir """ assert ckpt_dir and output_dir, 'Please pass in valid ckpt_dir and output_dir.' assert os.path.exists(ckpt_dir), f'ckpt_dir: {ckpt_dir} must exists in local disk.' if os.path.exists(os.path.join(ckpt_dir, SwiftModel.EXTRA_STATE_DIR)): raise AssertionError('Cannot transfer to peft format, because you are additional state dicts.') adapter_names = [ sub_dir for sub_dir in os.listdir(ckpt_dir) if os.path.isfile(os.path.join(ckpt_dir, sub_dir, CONFIG_NAME)) ] def has_custom_content(_json): if _json.get('swift_type', _json.get('peft_type')) != SwiftTuners.LORA: logger.warn('Only LoRA can be converted to peft format') return True from swift import LoRAConfig return not LoRAConfig(**_json).can_be_saved_to_peft() for adapter in adapter_names: with open(os.path.join(ckpt_dir, adapter, CONFIG_NAME), encoding='utf-8') as f: _json = json.load(f) if has_custom_content(_json): raise AssertionError('Cannot transfer to peft format, ' 'because you have special parameters or adapter types.') os.makedirs(output_dir, exist_ok=True) if ckpt_dir != output_dir: shutil.copytree(ckpt_dir, output_dir, dirs_exist_ok=True) for adapter in adapter_names: safe_serialization = os.path.isfile(os.path.join(output_dir, adapter, SAFETENSORS_WEIGHTS_NAME)) state_dict = SwiftModel.load_state_file(os.path.join(output_dir, adapter)) new_state_dict = {} for key, value in state_dict.items(): if not key.startswith('base_model.model.'): key = 'base_model.model.' + key key = key.replace(f'lora_A.{adapter}.', 'lora_A.') key = key.replace(f'lora_B.{adapter}.', 'lora_B.') key = key.replace(f'lora_embedding_A.{adapter}.', 'lora_embedding_A.') key = key.replace(f'lora_embedding_B.{adapter}.', 'lora_embedding_B.') key = key.replace(f'lora_magnitude_vector.{adapter}', 'lora_magnitude_vector') new_state_dict[key] = value state_dict = new_state_dict SwiftModel._save_state_dict(state_dict, os.path.join(output_dir, adapter), safe_serialization) from swift import LoRAConfig with open(os.path.join(output_dir, adapter, CONFIG_NAME), encoding='utf-8') as f: _json = json.load(f) peft_config = LoRAConfig(**_json).to_peft_config() peft_config.save_pretrained(os.path.join(output_dir, adapter)) if 'default' in adapter_names: shutil.move(os.path.join(output_dir, 'default', CONFIG_NAME), os.path.join(output_dir, CONFIG_NAME)) state_dict = SwiftModel.load_state_file(os.path.join(output_dir, 'default')) safe_serialization = os.path.isfile(os.path.join(output_dir, 'default', SAFETENSORS_WEIGHTS_NAME)) SwiftModel._save_state_dict(state_dict, output_dir, safe_serialization) shutil.rmtree(os.path.join(output_dir, 'default')) @staticmethod def from_pretrained(model: Union[nn.Module, SwiftModel, PeftModel], model_id: str = None, adapter_name: Union[str, List[str], Dict[str, str]] = None, revision: str = None, **kwargs): """Prepare a model by a model_id in the ModelScope hub or a local dir. Args: model(`Union[nn.Module, 'SwiftModel']`): The model to be tuned. model_id(`str`): The model id of the modelhub or a local dir containing the configs/weights. adapter_name(`str`, `optional`): The adapter_name to use. revision(`str`, `optional`): The model revision if the model_id is a model id of the modelhub. **kwargs: Extra kwargs needed by ``SwiftModel.from_pretrained`` or ``PeftModel.from_pretrained``. Returns: The model wrapped by SwiftModel or PeftModel. """ if not os.path.exists(model_id): model_id = snapshot_download(model_id, revision=revision) is_peft_model = False if os.path.exists(os.path.join(model_id, CONFIG_NAME)): with open(os.path.join(model_id, CONFIG_NAME), 'r', encoding='utf-8') as f: _json = json.load(f) is_peft_model = SWIFT_TYPE_KEY not in _json _name = adapter_name if isinstance( adapter_name, str) or adapter_name is None else adapter_name[0] \ if isinstance(adapter_name, list) else list(adapter_name.keys())[0] _name = _name or '' if os.path.exists(os.path.join(model_id, _name, CONFIG_NAME)): with open(os.path.join(model_id, _name, CONFIG_NAME), 'r', encoding='utf-8') as f: _json = json.load(f) is_peft_model = SWIFT_TYPE_KEY not in _json and 'extra_state_keys' not in _json if is_peft_model: def load_peft_model(_model, _adapter_name, _new_name=None): if not _new_name: _new_name = _adapter_name import peft if not isinstance(_model, peft.PeftModel): return PeftModel.from_pretrained( _model, os.path.join(model_id, _adapter_name) if _adapter_name != 'default' and os.path.exists(os.path.join(model_id, _adapter_name)) else model_id, revision=revision, adapter_name=_new_name, **kwargs) else: _model.load_adapter( os.path.join(model_id, _adapter_name) if _adapter_name != 'default' and os.path.exists(os.path.join(model_id, _adapter_name)) else model_id, _new_name) return _model if not adapter_name: peft_model = load_peft_model(model, 'default') for _dir in os.listdir(model_id): if os.path.isdir(os.path.join(model_id, _dir)) and \ os.path.exists(os.path.join(model_id, _dir, CONFIG_NAME)): peft_model = load_peft_model(peft_model, _dir) elif isinstance(adapter_name, str): return load_peft_model(model, adapter_name) elif isinstance(adapter_name, list): peft_model = model for name in adapter_name: peft_model = load_peft_model(peft_model, name) else: peft_model = model for key, value in adapter_name.items(): peft_model = load_peft_model(peft_model, key, value) return peft_model else: return SwiftModel.from_pretrained(model, model_id, revision=revision, adapter_name=adapter_name, **kwargs)