# Copyright (c) Alibaba, Inc. and its affiliates. # Copyright 2023-present the HuggingFace Inc. team. import os.path from dataclasses import asdict, dataclass, field from functools import partial, reduce from types import MethodType from typing import Dict, Optional import json import peft import torch import torch.nn import transformers from modelscope import snapshot_download from peft import (AdaLoraConfig, BOFTConfig, BOFTModel, LoftQConfig, LoHaConfig, LoKrConfig, LoraModel, OFTConfig, PeftConfig, PeftModel, PeftModelForCausalLM, PeftModelForSeq2SeqLM, PeftModelForSequenceClassification, PeftModelForTokenClassification, PrefixTuningConfig, PromptEncoderConfig, PromptLearningConfig, PromptTuningConfig, VeraConfig, VeraModel, get_peft_config, get_peft_model, get_peft_model_state_dict) from peft.config import PeftConfigMixin from peft.tuners import lora from peft.tuners.adalora import AdaLoraModel, RankAllocator from peft.tuners.lora import Embedding from transformers import Trainer from swift.utils import get_logger try: from peft import FourierFTModel except ImportError: FourierFTModel = None try: from peft import BoneModel except ImportError: BoneModel = None logger = get_logger() dispatchers = [] @dataclass class LoraConfig(peft.LoraConfig): lora_dtype: Optional[str] = field( default=None, metadata={'help': 'The lora dtype, default None means following the original layer\'s dtype'}) lorap_lr_ratio: Optional[float] = field(default=None, metadata={'help': 'The lr ratio of lora_B in lora+'}) lorap_emb_lr: float = field(default=1e-6, metadata={'help': 'The lr for embedding in lora+'}) def to_peft_config(self) -> peft.LoraConfig: _dict = asdict(self) _dict.pop('lora_dtype') _dict.pop('lorap_lr_ratio') _dict.pop('lorap_emb_lr') return peft.LoraConfig(**_dict) def save_pretrained(self, save_directory: str, **kwargs) -> None: self.to_peft_config().save_pretrained(save_directory, **kwargs) additional_args = { 'lora_dtype': self.lora_dtype, 'lorap_lr_ratio': self.lorap_lr_ratio, 'lorap_emb_lr': self.lorap_emb_lr, } with open(os.path.join(save_directory, 'additional_config.json'), 'w', encoding='utf-8') as f: json.dump(additional_args, f) @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, subfolder: Optional[str] = None, **kwargs): if hasattr(PeftConfigMixin, 'from_pretrained_origin'): self = PeftConfigMixin.from_pretrained_origin(pretrained_model_name_or_path, subfolder, **kwargs) else: self = super(LoraConfig, cls).from_pretrained(pretrained_model_name_or_path, subfolder, **kwargs) if type(self) == peft.LoraConfig: self = LoraConfig(**self.to_dict()) if os.path.isfile(os.path.join(pretrained_model_name_or_path, 'additional_config.json')): with open( os.path.join(pretrained_model_name_or_path, 'additional_config.json'), 'r', encoding='utf-8') as f: _json = json.load(f) for key, value in _json.items(): setattr(self, key, value) return self def _create_and_replace_hook(self, peft_config, adapter_name, target, *args, **kwargs): all_supported_names = ('linear', ) all_supported_types = (torch.nn.Embedding, torch.nn.Conv2d, transformers.pytorch_utils.Conv1D, lora.Linear) target_modules = getattr(peft_config, 'target_modules', None) if target is None: return if isinstance(target_modules, str) and not any( [name in target.__class__.__name__.lower() for name in all_supported_names]) and not any([isinstance(target, type_) for type_ in all_supported_types]): return if target.__class__.__name__ == 'NonDynamicallyQuantizableLinear': return return self._create_and_replace_origin(peft_config, adapter_name, target, *args, **kwargs) def _convert_dtype(target: torch.nn.Module, adapter_name: str, lora_dtype: str): if lora_dtype is not None: torch_dtype = eval(f'torch.{lora_dtype}') if hasattr(target, 'lora_A') and adapter_name in target.lora_A: target.lora_A[adapter_name].to(torch_dtype) target.lora_B[adapter_name].to(torch_dtype) if hasattr(target, 'lora_embedding_A') and adapter_name in target.lora_embedding_A: target.lora_embedding_A[adapter_name].to(torch_dtype) target.lora_embedding_B[adapter_name].to(torch_dtype) def create_optimizer_param_groups(self: PeftModel, **defaults): if not isinstance(self.peft_config[self.active_adapter], LoraConfig) or self.peft_config[self.active_adapter].lorap_lr_ratio is None: return None def get_module(name): parent_idx = 2 if 'lora' in name else 1 module_names = name.split(sep='.')[:-parent_idx] module = reduce(getattr, module_names, self.base_model) return module param_groups = { 'groupA': {}, 'groupB': {}, 'groupB_no_decay': {}, 'embedding': {}, } decay_parameters = Trainer.get_decay_parameter_names(None, self.base_model) for name, param in self.base_model.named_parameters(): if not param.requires_grad: continue module = get_module(name) if isinstance(module, Embedding): param_groups['embedding'][name] = param elif 'lora_B' in name or param.ndim == 1: if name in decay_parameters: param_groups['groupB'][name] = param else: param_groups['groupB_no_decay'][name] = param else: param_groups['groupA'][name] = param lr = defaults['lr'] weight_decay = defaults.get('weight_decay', 0.0) param_groups = [ { 'params': list(param_groups['groupA'].values()), 'weight_decay': weight_decay, 'lr': lr, }, { 'params': list(param_groups['embedding'].values()), 'weight_decay': weight_decay, 'lr': self.peft_config[self.active_adapter].lorap_emb_lr, }, { 'params': list(param_groups['groupB'].values()), 'weight_decay': weight_decay, 'lr': lr * self.peft_config[self.active_adapter].lorap_lr_ratio, }, { 'params': list(param_groups['groupB_no_decay'].values()), 'weight_decay': 0.0, 'lr': lr * self.peft_config[self.active_adapter].lorap_lr_ratio, }, ] return param_groups def adalora_forward(self, *args, **kwargs): from peft.utils.integrations import gather_params_ctx outputs = self.model.forward(*args, **kwargs) if (getattr(outputs, 'loss', None) is not None) and isinstance(outputs.loss, torch.Tensor): # Calculate the orthogonal regularization orth_reg_weight = self.peft_config[self.trainable_adapter_name].orth_reg_weight if orth_reg_weight <= 0: raise ValueError('orth_reg_weight should be greater than 0. ') regu_loss = 0 num_param = 0 for n, p in self.model.named_parameters(): if ('lora_A' in n or 'lora_B' in n) and self.trainable_adapter_name in n: if p.shape == torch.Size([0]): with gather_params_ctx(p, fwd_module=self): para_cov = p @ p.T if 'lora_A' in n else p.T @ p else: para_cov = p @ p.T if 'lora_A' in n else p.T @ p I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov)) # noqa: E741 I.requires_grad = False num_param += 1 if isinstance(regu_loss, torch.Tensor): regu_loss = regu_loss.to(para_cov.device) regu_loss += torch.norm(para_cov - I, p='fro') if num_param > 0: regu_loss = regu_loss / num_param else: regu_loss = 0 if isinstance(regu_loss, torch.Tensor) and isinstance(outputs.loss, torch.Tensor): regu_loss = regu_loss.to(outputs.loss.device) outputs.loss += orth_reg_weight * regu_loss return outputs def adalora_mask_to_budget(self, model, budget): value_ipt = {} vector_ipt = {} triplet_ipt = {} # Get the importance score for A, E, B for n, p in model.named_parameters(): if f'lora_A.{self.adapter_name}' in n: entry_ipt = self._element_score(n) comb_ipt = torch.mean(entry_ipt, dim=1, keepdim=True) name_m = n.replace('lora_A', '%s') if name_m not in vector_ipt: vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) if f'lora_B.{self.adapter_name}' in n: entry_ipt = self._element_score(n) comb_ipt = torch.mean(entry_ipt, dim=0, keepdim=False).view(-1, 1) name_m = n.replace('lora_B', '%s') if name_m not in vector_ipt: vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) if f'lora_E.{self.adapter_name}' in n: entry_ipt = self._element_score(n) name_m = n.replace('lora_E', '%s') value_ipt[name_m] = entry_ipt all_score = [] # Calculate the score for each triplet for name_m in vector_ipt: ipt_E = value_ipt[name_m] ipt_AB = torch.cat(vector_ipt[name_m], dim=1) sum_ipt = self._combine_ipt(ipt_E, ipt_AB) name_E = name_m % 'lora_E' triplet_ipt[name_E] = sum_ipt.view(-1, 1) sum_ipt = sum_ipt.view(-1) if all_score: sum_ipt = sum_ipt.to(all_score[0].device) all_score.append(sum_ipt) # Get the threshold by ranking ipt mask_threshold = torch.kthvalue( torch.cat(all_score), k=self.init_bgt - budget, )[0].item() rank_pattern = {} # Mask the unimportant triplets with torch.no_grad(): for n, p in model.named_parameters(): if f'lora_E.{self.adapter_name}' in n: p.masked_fill_(triplet_ipt[n] <= mask_threshold, 0.0) rank_pattern[n] = (~(triplet_ipt[n] <= mask_threshold)).view(-1).tolist() return rank_pattern def keep_device_forward(self, *args, **kwargs): x = args[0] weight = self.weight if hasattr(self, 'weight') else self.weight0 # compat megatron if weight.device != x.device: return self.forward_origin(x.to(weight.device), *args[1:], **kwargs) else: return self.forward_origin(*args, **kwargs) def hot_patch_peft_module(): from peft.tuners.lora import LoraLayer if hasattr(LoraModel, '_create_and_replace_origin'): return # Fix Lora does not support NonDynamicallyQuantizableLinear LoraModel._create_and_replace_origin = LoraModel._create_and_replace LoraModel._create_and_replace = _create_and_replace_hook AdaLoraModel._create_and_replace_origin = AdaLoraModel._create_and_replace AdaLoraModel._create_and_replace = _create_and_replace_hook VeraModel._create_and_replace_origin = VeraModel._create_and_replace VeraModel._create_and_replace = _create_and_replace_hook BOFTModel._create_and_replace_origin = BOFTModel._create_and_replace BOFTModel._create_and_replace = _create_and_replace_hook if FourierFTModel is not None: FourierFTModel._create_and_replace_origin = FourierFTModel._create_and_replace FourierFTModel._create_and_replace = _create_and_replace_hook if BoneModel is not None: BoneModel._create_and_replace_origin = BoneModel._create_and_replace BoneModel._create_and_replace = _create_and_replace_hook # Support type conversion def __new_init__(self, model: torch.nn.Module, config: Dict[str, LoraConfig], adapter_name: str): self.__init_origin__(model, config, adapter_name) active_adapters = self.active_adapter if isinstance(active_adapters, str): active_adapters = [active_adapters] for active_adapter in active_adapters: active_config = config[active_adapter] if isinstance(config, dict) else config if hasattr(active_config, 'lora_dtype'): for name, module in model.named_modules(): if isinstance(module, LoraLayer): _convert_dtype(module, active_adapter, active_config.lora_dtype) for lora in list(module.lora_A.values()) + list(module.lora_B.values()): if not hasattr(lora, 'forward_origin'): lora.forward_origin = lora.forward lora.forward = MethodType(keep_device_forward, lora) LoraModel.__init_origin__ = LoraModel.__init__ LoraModel.__init__ = __new_init__ # Support LoRA+ PeftModel.create_optimizer_param_groups = create_optimizer_param_groups PeftConfigMixin.from_pretrained_origin = PeftConfigMixin.from_pretrained PeftConfigMixin.from_pretrained = LoraConfig.from_pretrained # Compatible with SwiftModel def dummy_function(*args, **kwargs): logger.warn(f'The function {kwargs["func"]} has no effects, consider using other functions.') PeftModel.activate_adapter = PeftModel.set_adapter PeftModel.deactivate_adapter = partial(dummy_function, func='deactivate_adapter') PeftModel.set_active_adapters = partial(dummy_function, func='set_active_adapters') # Fix adalora does not support device_map AdaLoraModel.forward = adalora_forward RankAllocator.mask_to_budget = adalora_mask_to_budget def get_wrapped_class(module_class): """Get a custom wrapper class for peft classes to download the models from the ModelScope hub Args: module_class: The actual module class Returns: The wrapper """ class PeftWrapper(module_class): @classmethod def from_pretrained(cls, model, model_id, *args, revision: Optional[str] = None, **kwargs): if not os.path.exists(model_id): model_id = snapshot_download(model_id, revision=revision) return module_class.from_pretrained(model, model_id, *args, **kwargs) PeftWrapper.__name__ = module_class.__name__ PeftWrapper.__qualname__ = module_class.__qualname__ return PeftWrapper def wrap_module(module): if not hasattr(module, 'from_pretrained'): return module return get_wrapped_class(module) hot_patch_peft_module() PeftModel = wrap_module(PeftModel) PeftConfig = wrap_module(PeftConfig) PeftModelForSeq2SeqLM = wrap_module(PeftModelForSeq2SeqLM) PeftModelForSequenceClassification = wrap_module(PeftModelForSequenceClassification) PeftModelForTokenClassification = wrap_module(PeftModelForTokenClassification) PeftModelForCausalLM = wrap_module(PeftModelForCausalLM) PromptEncoderConfig = wrap_module(PromptEncoderConfig) PromptTuningConfig = wrap_module(PromptTuningConfig) PrefixTuningConfig = wrap_module(PrefixTuningConfig) PromptLearningConfig = wrap_module(PromptLearningConfig) LoraConfig = wrap_module(LoraConfig) AdaLoraConfig = wrap_module(AdaLoraConfig) LoHaConfig = wrap_module(LoHaConfig) LoKrConfig = wrap_module(LoKrConfig) LoftQConfig = wrap_module(LoftQConfig) OFTConfig = wrap_module(OFTConfig) BOFTConfig = wrap_module(BOFTConfig) VeraConfig = wrap_module(VeraConfig) OFTConfig = wrap_module(OFTConfig) get_peft_config = get_peft_config get_peft_model_state_dict = get_peft_model_state_dict get_peft_model = get_peft_model