From c21afbe868734c0af8bd4577c4c7acdf366b96d1 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 18:56:24 +0530 Subject: [PATCH 01/35] multi adapter for training and inference Might have breaking changes --- src/peft/mapping.py | 43 +-- src/peft/peft_model.py | 284 ++++++++++------ src/peft/tuners/__init__.py | 2 +- src/peft/tuners/lora.py | 563 ++++++++++++------------------- src/peft/utils/__init__.py | 6 +- src/peft/utils/adapters_utils.py | 18 - src/peft/utils/config.py | 16 +- src/peft/utils/other.py | 83 ++++- src/peft/utils/save_and_load.py | 57 +++- 9 files changed, 524 insertions(+), 548 deletions(-) delete mode 100644 src/peft/utils/adapters_utils.py diff --git a/src/peft/mapping.py b/src/peft/mapping.py index dbb9f36..c814655 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -38,27 +38,6 @@ PEFT_TYPE_TO_CONFIG_MAPPING = { "LORA": LoraConfig, } -TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { - "t5": ["q", "v"], - "mt5": ["q", "v"], - "bart": ["q_proj", "v_proj"], - "gpt2": ["c_attn"], - "bloom": ["query_key_value"], - "opt": ["q_proj", "v_proj"], - "gptj": ["q_proj", "v_proj"], - "gpt_neox": ["query_key_value"], - "gpt_neo": ["q_proj", "v_proj"], - "bert": ["query", "value"], - "roberta": ["query", "value"], - "xlm-roberta": ["query", "value"], - "electra": ["query", "value"], - "deberta-v2": ["query_proj", "value_proj"], - "deberta": ["in_proj"], - "layoutlm": ["query", "value"], - "llama": ["q_proj", "v_proj"], - "chatglm": ["query_key_value"], -} - def get_peft_config(config_dict): """ @@ -113,19 +92,6 @@ def _prepare_prompt_learning_config(peft_config, model_config): return peft_config -def _prepare_lora_config(peft_config, model_config): - if peft_config.target_modules is None: - if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: - raise ValueError("Please specify `target_modules` in `peft_config`") - peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] - if len(peft_config.target_modules) == 1: - peft_config.fan_in_fan_out = True - peft_config.enable_lora = [True, False, True] - if peft_config.inference_mode: - peft_config.merge_weights = True - return peft_config - - def get_peft_model(model, peft_config): """ Returns a Peft model object from a model and a config. @@ -137,11 +103,10 @@ def get_peft_model(model, peft_config): model_config = model.config.to_dict() peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) - if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): - peft_config = _prepare_lora_config(peft_config, model_config) + if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys() and not isinstance( + peft_config, PromptLearningConfig + ): return PeftModel(model, peft_config) - if not isinstance(peft_config, PromptLearningConfig): - peft_config = _prepare_lora_config(peft_config, model_config) - else: + if isinstance(peft_config, PromptLearningConfig): peft_config = _prepare_prompt_learning_config(peft_config, model_config) return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index f73a66a..cbc8a89 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -36,6 +36,7 @@ from .utils import ( PeftType, PromptLearningConfig, TaskType, + _set_adapter, _set_trainable, get_peft_model_state_dict, set_peft_model_state_dict, @@ -43,6 +44,14 @@ from .utils import ( ) +PEFT_TYPE_TO_MODEL_MAPPING = { + PeftType.LORA: LoraModel, + PeftType.PROMPT_TUNING: PromptEmbedding, + PeftType.P_TUNING: PromptEncoder, + PeftType.PREFIX_TUNING: PrefixEncoder, +} + + class PeftModel(PushToHubMixin, torch.nn.Module): """ Parameter-Efficient Fine-Tuning Model. Base model encompassing various Peft methods. @@ -67,20 +76,19 @@ class PeftModel(PushToHubMixin, torch.nn.Module): in the base model if `isinstance(self.peft_config, PromptLearningConfig)`. """ - def __init__(self, model, peft_config: PeftConfig): + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): super().__init__() - self.peft_config = peft_config self.base_model = model self.config = self.base_model.config self.modules_to_save = None - if isinstance(self.peft_config, PromptLearningConfig): - self._setup_prompt_encoder() - else: - self.base_model = LoraModel(peft_config, model) - if getattr(self.peft_config, "modules_to_save", None) is not None: - self.modules_to_save = self.peft_config.modules_to_save - _set_trainable(self) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.peft_config = {} + self.active_adapter = adapter_name + if not isinstance(peft_config, PromptLearningConfig): + self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( + self.base_model, peft_config, adapter_name + ) + self.add_adapter(adapter_name, peft_config) def save_pretrained(self, save_directory, **kwargs): r""" @@ -98,27 +106,30 @@ class PeftModel(PushToHubMixin, torch.nn.Module): raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file") os.makedirs(save_directory, exist_ok=True) - # save only the trainable weights - output_state_dict = get_peft_model_state_dict(self, kwargs.get("state_dict", None)) - torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME)) + for adapter_name, peft_config in self.peft_config.items(): + # save only the trainable weights + output_state_dict = get_peft_model_state_dict(self, adapter_name, kwargs.get("state_dict", None)) + output_dir = os.path.join(save_directory, adapter_name) if adapter_name != "default" else save_directory + os.makedirs(output_dir, exist_ok=True) + torch.save(output_state_dict, os.path.join(output_dir, WEIGHTS_NAME)) - # save the config and change the inference mode to `True` - if self.peft_config.base_model_name_or_path is None: - self.peft_config.base_model_name_or_path = ( - self.base_model.__dict__.get("name_or_path", None) - if isinstance(self.peft_config, PromptLearningConfig) - else self.base_model.model.__dict__.get("name_or_path", None) - ) - inference_mode = self.peft_config.inference_mode - self.peft_config.inference_mode = True - self.peft_config.save_pretrained(save_directory) - self.peft_config.inference_mode = inference_mode + # save the config and change the inference mode to `True` + if peft_config.base_model_name_or_path is None: + peft_config.base_model_name_or_path = ( + self.base_model.__dict__.get("name_or_path", None) + if isinstance(self.peft_config, PromptLearningConfig) + else self.base_model.model.__dict__.get("name_or_path", None) + ) + inference_mode = self.peft_config.inference_mode + peft_config.inference_mode = True + peft_config.save_pretrained(output_dir) + peft_config.inference_mode = inference_mode @classmethod - def from_pretrained(cls, model, model_id, **kwargs): + def from_pretrained(cls, model, model_id, adapter_name="default", **kwargs): r""" Args: - Instantiate a `LoraModel` from a pretrained Lora configuration and weights. + Instantiate a `PeftModel` from a pretrained Peft configuration and weights. model (`transformers.PreTrainedModel`): The model to be adapted. The model should be initialized with the `from_pretrained` method. from `transformers` library. @@ -132,58 +143,26 @@ class PeftModel(PushToHubMixin, torch.nn.Module): from .mapping import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PEFT_TYPE_TO_CONFIG_MAPPING # load the config - config = PEFT_TYPE_TO_CONFIG_MAPPING[PeftConfig.from_pretrained(model_id).peft_type].from_pretrained(model_id) + config = PEFT_TYPE_TO_CONFIG_MAPPING[ + PeftConfig.from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)).peft_type + ].from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)) - if getattr(model, "hf_device_map", None) is not None: + if (getattr(model, "hf_device_map", None) is not None) and len( + set(model.hf_device_map.values()).intersection({"cpu", "disk"}) + ) > 0: remove_hook_from_submodules(model) if config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys(): - model = cls(model, config) + model = cls(model, config, adapter_name) else: - model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config) - - # load weights if any - if os.path.exists(os.path.join(model_id, WEIGHTS_NAME)): - filename = os.path.join(model_id, WEIGHTS_NAME) - else: - try: - filename = hf_hub_download(model_id, WEIGHTS_NAME) - except: # noqa - raise ValueError( - f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. " - f"Please check that the file {WEIGHTS_NAME} is present at {model_id}." - ) - - adapters_weights = torch.load( - filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu") - ) - # load the weights into the model - model = set_peft_model_state_dict(model, adapters_weights) - if getattr(model, "hf_device_map", None) is not None: - device_map = kwargs.get("device_map", "auto") - max_memory = kwargs.get("max_memory", None) - no_split_module_classes = model._no_split_modules - if device_map != "sequential": - max_memory = get_balanced_memory( - model, - max_memory=max_memory, - no_split_module_classes=no_split_module_classes, - low_zero=(device_map == "balanced_low_0"), - ) - if isinstance(device_map, str): - device_map = infer_auto_device_map( - model, max_memory=max_memory, no_split_module_classes=no_split_module_classes - ) - model = dispatch_model(model, device_map=device_map) - hook = AlignDevicesHook(io_same_device=True) - if model.peft_config.peft_type == PeftType.LORA: - add_hook_to_module(model.base_model.model, hook) - else: - remove_hook_from_submodules(model.prompt_encoder) - add_hook_to_module(model.base_model, hook) + model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[config.task_type](model, config, adapter_name) + model.load_adapter(model_id, adapter_name, **kwargs) return model - def _setup_prompt_encoder(self): + def _setup_prompt_encoder(self, adapter_name): + config = self.peft_config[adapter_name] + self.prompt_encoder = torch.nn.ModuleDict({}) + self.prompt_tokens = {} transformer_backbone = None for name, module in self.base_model.named_children(): for param in module.parameters(): @@ -194,51 +173,50 @@ class PeftModel(PushToHubMixin, torch.nn.Module): transformer_backbone = module self.transformer_backbone_name = name - if self.peft_config.num_transformer_submodules is None: - self.peft_config.num_transformer_submodules = ( - 2 if self.peft_config.task_type == TaskType.SEQ_2_SEQ_LM else 1 - ) + if config.num_transformer_submodules is None: + config.num_transformer_submodules = 2 if config.task_type == TaskType.SEQ_2_SEQ_LM else 1 for named_param, value in list(transformer_backbone.named_parameters()): if value.shape[0] == self.base_model.config.vocab_size: self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", "")) break - if self.peft_config.peft_type == PeftType.PROMPT_TUNING: - prompt_encoder = PromptEmbedding(self.peft_config, self.word_embeddings) - elif self.peft_config.peft_type == PeftType.P_TUNING: - prompt_encoder = PromptEncoder(self.peft_config) - elif self.peft_config.peft_type == PeftType.PREFIX_TUNING: - prompt_encoder = PrefixEncoder(self.peft_config) + if config.peft_type == PeftType.PROMPT_TUNING: + prompt_encoder = PromptEmbedding(config, self.word_embeddings) + elif config.peft_type == PeftType.P_TUNING: + prompt_encoder = PromptEncoder(config) + elif config.peft_type == PeftType.PREFIX_TUNING: + prompt_encoder = PrefixEncoder(config) else: raise ValueError("Not supported") - self.prompt_encoder = prompt_encoder - self.prompt_tokens = torch.arange( - self.peft_config.num_virtual_tokens * self.peft_config.num_transformer_submodules + self.prompt_encoder.update(torch.nn.ModuleDict({adapter_name: prompt_encoder})) + self.prompt_tokens[adapter_name] = torch.arange( + config.num_virtual_tokens * config.num_transformer_submodules ).long() - def get_prompt_embedding_to_save(self): + def get_prompt_embedding_to_save(self, adapter_name): """ Returns the prompt embedding to save when saving the model. Only applicable when `peft_config.peft_type != PeftType.LORA`. """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: - prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens] - prompt_embeddings = self.prompt_encoder(prompt_tokens) + prompt_tokens = self.prompt_tokens[adapter_name].unsqueeze(0).expand(1, -1).to(self.device) + if self.peft_config[adapter_name].peft_type == PeftType.PREFIX_TUNING: + prompt_tokens = prompt_tokens[:, : self.peft_config[adapter_name].num_virtual_tokens] + prompt_embeddings = self.prompt_encoder[adapter_name](prompt_tokens) return prompt_embeddings[0].detach().cpu() def get_prompt(self, batch_size): """ Returns the virtual prompts to use for Peft. Only applicable when `peft_config.peft_type != PeftType.LORA`. """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.device) + prompt_encoder = self.prompt_encoder[self.active_adapter] + prompt_tokens = self.prompt_tokens[self.active_adapter].unsqueeze(0).expand(batch_size, -1).to(self.device) if self.peft_config.peft_type == PeftType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens] if self.peft_config.inference_mode: - past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + past_key_values = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: - past_key_values = self.prompt_encoder(prompt_tokens) + past_key_values = prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, self.peft_config.num_virtual_tokens, @@ -257,9 +235,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): return past_key_values else: if self.peft_config.inference_mode: - prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + prompts = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: - prompts = self.prompt_encoder(prompt_tokens) + prompts = prompt_encoder(prompt_tokens) return prompts def print_trainable_parameters(self): @@ -299,13 +277,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Disables the adapter module. """ - if isinstance(self.peft_config, PromptLearningConfig): + if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): old_forward = self.forward self.forward = self.base_model.forward else: self.base_model.disable_adapter_layers() yield - if isinstance(self.peft_config, PromptLearningConfig): + if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): self.forward = old_forward else: self.base_model.enable_adapter_layers() @@ -314,7 +292,91 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Returns the base model. """ - return self.base_model if isinstance(self.peft_config, PromptLearningConfig) else self.base_model.model + return ( + self.base_model + if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig) + else self.base_model.model + ) + + def add_adapter(self, adapter_name, peft_config): + self.peft_config[adapter_name] = peft_config + if isinstance(peft_config, PromptLearningConfig): + self._setup_prompt_encoder(adapter_name) + else: + self.base_model.add_adapter(adapter_name, peft_config) + if getattr(peft_config, "modules_to_save", None) is not None: + if self.modules_to_save is None: + self.modules_to_save = set(peft_config.modules_to_save) + else: + self.modules_to_save = self.modules_to_save.update(peft_config.modules_to_save) + _set_trainable(self, adapter_name) + + def load_adapter(self, model_id, adapter_name, **kwargs): + from .mapping import PEFT_TYPE_TO_CONFIG_MAPPING + + if adapter_name not in self.peft_config: + # load the config + peft_config = PEFT_TYPE_TO_CONFIG_MAPPING[ + PeftConfig.from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)).peft_type + ].from_pretrained(model_id, subfolder=kwargs.get("subfolder", None)) + self.add_adapter(adapter_name, peft_config) + + # load weights if any + if kwargs.get("subfolder", None) is not None: + path = os.path.join(model_id, kwargs["subfolder"]) + if os.path.exists(os.path.join(path, WEIGHTS_NAME)): + filename = os.path.join(path, WEIGHTS_NAME) + else: + try: + filename = hf_hub_download(model_id, WEIGHTS_NAME, subfolder=kwargs.get("subfolder", None)) + except: # noqa + raise ValueError( + f"Can't find weights for {model_id} in {model_id} or in the Hugging Face Hub. " + f"Please check that the file {WEIGHTS_NAME} is present at {model_id}." + ) + + adapters_weights = torch.load( + filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu") + ) + # load the weights into the model + set_peft_model_state_dict(self, adapter_name, adapters_weights) + if ( + (getattr(self, "hf_device_map", None) is not None) + and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) + and len(self.peft_config == 1) + ): + device_map = kwargs.get("device_map", "auto") + max_memory = kwargs.get("max_memory", None) + no_split_module_classes = self._no_split_modules + if device_map != "sequential": + max_memory = get_balanced_memory( + self, + max_memory=max_memory, + no_split_module_classes=no_split_module_classes, + low_zero=(device_map == "balanced_low_0"), + ) + if isinstance(device_map, str): + device_map = infer_auto_device_map( + self, max_memory=max_memory, no_split_module_classes=no_split_module_classes + ) + dispatch_model(self, device_map=device_map) + hook = AlignDevicesHook(io_same_device=True) + if not isinstance(self.peft_config[adapter_name]) == PeftType.LORA: + add_hook_to_module(self.base_model.model, hook) + else: + remove_hook_from_submodules(self.prompt_encoder) + add_hook_to_module(self.base_model, hook) + + def set_adapter(self, adapter_name): + """ + Sets the active adapter. + """ + if adapter_name not in self.peft_config: + raise ValueError(f"Adapter {adapter_name} not found.") + self.active_adapter = adapter_name + if not isinstance(self.peft_config[adapter_name], PromptLearningConfig): + self.base_model.set_adapter(adapter_name) + _set_adapter(self, adapter_name) class PeftModelForSequenceClassification(PeftModel): @@ -343,9 +405,12 @@ class PeftModelForSequenceClassification(PeftModel): params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) - self.modules_to_save = ["classifier", "score"] + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) + if self.modules_to_save is None: + self.modules_to_save = {"classifier", "score"} + else: + self.modules_to_save.update({"classifier", "score"}) for name, _ in self.base_model.named_children(): if any(module_name in name for module_name in self.modules_to_save): @@ -353,7 +418,7 @@ class PeftModelForSequenceClassification(PeftModel): break # to make sure classifier layer is trainable - _set_trainable(self) + _set_trainable(self, adapter_name) def forward( self, @@ -510,8 +575,8 @@ class PeftModelForCausalLM(PeftModel): params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation def forward( @@ -647,8 +712,8 @@ class PeftModelForSeq2SeqLM(PeftModel): params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) + def __init__(self, model, peft_config: PeftConfig, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) self.base_model_prepare_inputs_for_generation = self.base_model.prepare_inputs_for_generation self.base_model_prepare_encoder_decoder_kwargs_for_generation = ( self.base_model._prepare_encoder_decoder_kwargs_for_generation @@ -818,9 +883,12 @@ class PeftModelForTokenClassification(PeftModel): params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117 """ - def __init__(self, model, peft_config: PeftConfig): - super().__init__(model, peft_config) - self.modules_to_save = ["classifier", "score"] + def __init__(self, model, peft_config: PeftConfig = None, adapter_name="default"): + super().__init__(model, peft_config, adapter_name) + if self.modules_to_save is None: + self.modules_to_save = {"classifier", "score"} + else: + self.modules_to_save.update({"classifier", "score"}) for name, _ in self.base_model.named_children(): if any(module_name in name for module_name in self.modules_to_save): @@ -828,7 +896,7 @@ class PeftModelForTokenClassification(PeftModel): break # to make sure classifier layer is trainable - _set_trainable(self) + _set_trainable(self, adapter_name) def forward( self, diff --git a/src/peft/tuners/__init__.py b/src/peft/tuners/__init__.py index 38b7926..8f93079 100644 --- a/src/peft/tuners/__init__.py +++ b/src/peft/tuners/__init__.py @@ -17,7 +17,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .lora import LoraConfig, LoraModel from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType from .prefix_tuning import PrefixEncoder, PrefixTuningConfig from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit +from .lora import LoraConfig, LoraModel diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 0f65cbf..34cee7d 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -25,7 +25,13 @@ import torch.nn as nn import torch.nn.functional as F from transformers.pytorch_utils import Conv1D -from ..utils import PeftConfig, PeftType, transpose +from ..utils import ( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + PeftConfig, + PeftType, + _get_submodules, + transpose, +) def is_bnb_available(): @@ -48,8 +54,9 @@ class LoraConfig(PeftConfig): lora_dropout (`float`): The dropout probability for Lora layers. merge_weights (`bool`): Whether to merge the weights of the Lora layers with the base transformer model in `eval` mode. - fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out) - enable_lora ( `List[bool]`): Used with `lora.MergedLinear`. + fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out). + For example, gpt-2 uses `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`.: + enable_lora ( `List[bool]`): Used with `lora.MergedLinear`. Usually set to [True, False, True]. bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only' modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint. @@ -72,7 +79,6 @@ class LoraConfig(PeftConfig): default=False, metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, ) - enable_lora: Optional[List[bool]] = field(default=None, metadata={"help": "Used with `lora.MergedLinear`."}) bias: str = field(default="none", metadata={"help": "Bias type for Lora. Can be 'none', 'all' or 'lora_only'"}) modules_to_save: Optional[List[str]] = field( default=None, @@ -88,38 +94,26 @@ class LoraConfig(PeftConfig): class LoraModel(torch.nn.Module): - """ - Creates Low Rank Adapter (Lora) model from a pretrained transformers model. - - Args: - model ([`transformers.PreTrainedModel`]): The model to be adapted. - config ([`LoraConfig`]): The configuration of the Lora model. - - Returns: - `torch.nn.Module`: The Lora model. - - Example:: - - >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>> - config = LoraConfig( - peft_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"], - lora_dropout=0.01, ) - >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoraModel(config, model) - - **Attributes**: - - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. - - **peft_config** ([`LoraConfig`]): The configuration of the Lora model. - """ - - def __init__(self, config, model): + def __init__(self, model, config, adapter_name): super().__init__() - self.peft_config = config self.model = model - self._find_and_replace() - mark_only_lora_as_trainable(self.model, self.peft_config.bias) self.forward = self.model.forward + self.config = config + self.add_adapter(adapter_name) - def _find_and_replace(self): + def add_adapter(self, adapter_name, config=None): + if config is not None: + config = self._prepare_lora_config(config, self.model.config.to_dict()) + self.config[adapter_name] = config + self._find_and_replace(adapter_name) + if len(self.config) > 1 and self.config[adapter_name].bias != "none": + raise ValueError( + "LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters." + ) + mark_only_lora_as_trainable(self.model, self.config[adapter_name].bias) + + def _find_and_replace(self, adapter_name): + lora_config = self.config[adapter_name] loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) if loaded_in_8bit and not is_bnb_available(): raise ImportError( @@ -129,68 +123,72 @@ class LoraModel(torch.nn.Module): is_target_modules_in_base_model = False is_hf_device_map_available = hasattr(self.model, "hf_device_map") kwargs = { - "r": self.peft_config.r, - "lora_alpha": self.peft_config.lora_alpha, - "lora_dropout": self.peft_config.lora_dropout, - "fan_in_fan_out": self.peft_config.fan_in_fan_out, - "merge_weights": (self.peft_config.merge_weights or self.peft_config.inference_mode) + "r": lora_config.r, + "lora_alpha": lora_config.lora_alpha, + "lora_dropout": lora_config.lora_dropout, + "fan_in_fan_out": lora_config.fan_in_fan_out, + "merge_weights": (lora_config.merge_weights or lora_config.inference_mode) and not is_hf_device_map_available, } key_list = [key for key, _ in self.model.named_modules()] for key in key_list: - if isinstance(self.peft_config.target_modules, str): - target_module_found = re.fullmatch(self.peft_config.target_modules, key) + if isinstance(lora_config.target_modules, str): + target_module_found = re.fullmatch(lora_config.target_modules, key) else: - target_module_found = any(key.endswith(target_key) for target_key in self.peft_config.target_modules) + target_module_found = any(key.endswith(target_key) for target_key in lora_config.target_modules) if target_module_found: if not is_target_modules_in_base_model: is_target_modules_in_base_model = True - parent, target, target_name = self._get_submodules(key) + parent, target, target_name = _get_submodules(key) bias = target.bias is not None - if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): - kwargs.update( - { - "has_fp16_weights": target.state.has_fp16_weights, - "memory_efficient_backward": target.state.memory_efficient_backward, - "threshold": target.state.threshold, - "index": target.index, - } - ) - if self.peft_config.enable_lora is None: - new_module = Linear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - else: - kwargs.update({"enable_lora": self.peft_config.enable_lora}) - new_module = MergedLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - elif isinstance(target, torch.nn.Linear) and self.peft_config.enable_lora is None: - new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs) - elif self.peft_config.enable_lora is not None: - kwargs.update({"enable_lora": self.peft_config.enable_lora}) - if isinstance(target, Conv1D): - in_features, out_features = ( - target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape + if isinstance(target, LoraLayer): + target.update_layer(adapter_name, lora_config.r, lora_config.lora_alpha, lora_config.lora_dropout) + else: + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): + kwargs.update( + { + "has_fp16_weights": target.state.has_fp16_weights, + "memory_efficient_backward": target.state.memory_efficient_backward, + "threshold": target.state.threshold, + "index": target.index, + } + ) + new_module = Linear8bitLt( + adapter_name, target.in_features, target.out_features, bias=bias, **kwargs ) else: - in_features, out_features = target.in_features, target.out_features - if kwargs["fan_in_fan_out"]: - warnings.warn( - "fan_in_fan_out is set to True but the target module is not a Conv1D. " - "Setting fan_in_fan_out to False." + if isinstance(target, torch.nn.Linear): + in_features, out_features = target.in_features, target.out_features + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = False + elif isinstance(target, Conv1D): + in_features, out_features = ( + target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape ) - kwargs["fan_in_fan_out"] = self.peft_config.fan_in_fan_out = False - new_module = MergedLinear(in_features, out_features, bias=bias, **kwargs) - self._replace_module(parent, target_name, new_module, target) + if not kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to False but the target module is `Conv1D`. " + "Setting fan_in_fan_out to True." + ) + kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True + else: + raise ValueError( + f"Target module {target} is not supported. " + f"Currently, only `torch.nn.Linear` and `Conv1D` are supported." + ) + new_module = Linear(adapter_name, in_features, out_features, bias=bias, **kwargs) + + self._replace_module(parent, target_name, new_module, target) if not is_target_modules_in_base_model: raise ValueError( - f"Target modules {self.peft_config.target_modules} not found in the base model. " + f"Target modules {lora_config.target_modules} not found in the base model. " f"Please check the target modules and try again." ) - def _get_submodules(self, key): - parent = self.model.get_submodule(".".join(key.split(".")[:-1])) - target_name = key.split(".")[-1] - target = self.model.get_submodule(key) - return parent, target, target_name - def _replace_module(self, parent_module, child_name, new_module, old_module): setattr(parent_module, child_name, new_module) new_module.weight = old_module.weight @@ -217,9 +215,12 @@ class LoraModel(torch.nn.Module): return None def get_peft_config_as_dict(self, inference: bool = False): - config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(self.peft_config).items()} - if inference: - config["inference_mode"] = True + config_dict = {} + for key, value in self.config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config return config def _set_adapter_layers(self, enabled=True): @@ -233,6 +234,34 @@ class LoraModel(torch.nn.Module): def disable_adapter_layers(self): self._set_adapter_layers(enabled=False) + def set_adapter(self, adapter_name): + for module in self.model.modules(): + if isinstance(module, LoraLayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.active_adapter = adapter_name + + def merge_adapter(self): + for module in self.model.modules(): + if isinstance(module, LoraLayer): + module.merge() + + def unmerge_adapter(self): + for module in self.model.modules(): + if isinstance(module, LoraLayer): + module.unmerge() + + @staticmethod + def _prepare_lora_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] + if peft_config.inference_mode: + peft_config.merge_weights = True + return peft_config + # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py # and modified to work with PyTorch FSDP @@ -266,28 +295,53 @@ def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: class LoraLayer: def __init__( self, - r: int, - lora_alpha: int, - lora_dropout: float, merge_weights: bool, + in_features: int, + out_features: int, ): - self.r = r - self.lora_alpha = lora_alpha - # Optional dropout - if lora_dropout > 0.0: - self.lora_dropout = nn.Dropout(p=lora_dropout) - else: - self.lora_dropout = lambda x: x + self.r = {} + self.lora_alpha = {} + self.scaling = {} + self.lora_dropout = nn.ModuleDict({}) + self.lora_A = nn.ModuleDict({}) + self.lora_B = nn.ModuleDict({}) # Mark the weight as unmerged self.merged = False self.merge_weights = merge_weights self.disable_adapters = False + self.in_features = in_features + self.out_features = out_features + + def update_layer(self, adapter_name, r, lora_alpha, lora_dropout): + self.r[adapter_name] = r + self.lora_alpha[adapter_name] = lora_alpha + if lora_dropout > 0.0: + lora_dropout_layer = nn.Dropout(p=lora_dropout) + else: + + def lora_dropout_layer(x): + return x + + self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer})) + # Actual trainable parameters + if r > 0: + self.lora_A.update(nn.ModuleDict({nn.Linear(self.in_features, r, bias=False)})) + self.lora_B.update(nn.ModuleDict({nn.Linear(r, self.out_features, bias=False)})) + self.scaling[adapter_name] = lora_alpha / r + self.reset_lora_parameters(adapter_name) + + def reset_lora_parameters(self, adapter_name): + if adapter_name in self.lora_A.keys(): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A[adapter_name].weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B[adapter_name].weight) -class Linear(nn.Linear, LoraLayer): +class Linear(nn.Linear): # Lora implemented in a dense layer def __init__( self, + adapter_name: str, in_features: int, out_features: int, r: int = 0, @@ -298,185 +352,67 @@ class Linear(nn.Linear, LoraLayer): **kwargs, ): nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + LoraLayer.__init__(self, merge_weights=merge_weights, in_features=in_features, out_features=out_features) + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False self.fan_in_fan_out = fan_in_fan_out - # Actual trainable parameters - if r > 0: - self.lora_A = nn.Linear(in_features, r, bias=False) - self.lora_B = nn.Linear(r, out_features, bias=False) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - self.reset_parameters() if fan_in_fan_out: self.weight.data = self.weight.data.T - def reset_parameters(self): nn.Linear.reset_parameters(self) - if hasattr(self, "lora_A"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.active_adapter = adapter_name - def train(self, mode: bool = True): - nn.Linear.train(self, mode) - self.lora_A.train(mode) - self.lora_B.train(mode) - if not mode and self.merge_weights and not self.merged: - # Merge the weights and mark it - if self.r > 0: - self.weight.data += ( - transpose(self.lora_B.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling + def merge(self): + if not self.merge_weights: + warnings.warn("Nothing to merge. Set merge_weights to True to enable merging.") + return + if self.merged: + warnings.warn("Already merged. Nothing to do.") + return + if self.r[self.active_adapter] > 0: + self.weight.data += ( + transpose( + self.lora_B[self.active_adapter].weight @ self.lora_A[self.active_adapter].weight, + self.fan_in_fan_out, ) - self.merged = True - elif self.merge_weights and self.merged: - # Make sure that the weights are not merged - if self.r > 0: - self.weight.data -= ( - transpose(self.lora_B.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling - ) - self.merged = False - - def eval(self): - nn.Linear.eval(self) - self.lora_A.eval() - self.lora_B.eval() - - def forward(self, x: torch.Tensor): - if self.disable_adapters: - if self.r > 0 and self.merged: - self.weight.data -= ( - transpose(self.lora_B.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling - ) - self.merged = False - - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - elif self.r > 0 and not self.merged: - result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - if self.r > 0: - result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling - return result - else: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - - -class MergedLinear(nn.Linear, LoraLayer): - # Lora implemented in a dense layer - def __init__( - self, - in_features: int, - out_features: int, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0.0, - enable_lora: List[bool] = [False], - fan_in_fan_out: bool = False, - merge_weights: bool = True, - **kwargs, - ): - nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) - if out_features % len(enable_lora) != 0: - raise ValueError("The length of enable_lora must divide out_features") - self.enable_lora = enable_lora - self.fan_in_fan_out = fan_in_fan_out - # Actual trainable parameters - if r > 0 and any(enable_lora): - self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False) - self.lora_B = nn.Conv1d( - r * sum(enable_lora), - out_features // len(enable_lora) * sum(enable_lora), - kernel_size=1, - groups=2, - bias=False, + * self.scaling[self.active_adapter] ) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - # Compute the indices - self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1) - self.lora_ind[enable_lora, :] = True - self.lora_ind = self.lora_ind.view(-1) - self.reset_parameters() - if fan_in_fan_out: - self.weight.data = self.weight.data.T - - def reset_parameters(self): - nn.Linear.reset_parameters(self) - if hasattr(self, "lora_A"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) - - def zero_pad(self, x): - result = x.new_zeros((*x.shape[:-1], self.out_features)) - result = result.view(-1, self.out_features) - result[:, self.lora_ind] = x.reshape(-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)) - return result.view((*x.shape[:-1], self.out_features)) - - def train(self, mode: bool = True): - nn.Linear.train(self, mode) - self.lora_A.train(mode) - self.lora_B.train(mode) - if not mode and self.merge_weights and not self.merged: - # Merge the weights and mark it - if self.r > 0 and any(self.enable_lora): - delta_w = ( - F.conv1d( - self.lora_A.weight.data.unsqueeze(0), - self.lora_B.weight.data, - groups=sum(self.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - self.weight.data += transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) self.merged = True - elif self.merge_weights and self.merged: - # Make sure that the weights are not merged - if self.r > 0 and any(self.enable_lora): - delta_w = ( - F.conv1d( - self.lora_A.weight.data.unsqueeze(0), - self.lora_B.weight.data, - groups=sum(self.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) - self.merged = False - def eval(self): - nn.Linear.eval(self) - self.lora_A.eval() - self.lora_B.eval() + def unmerge(self): + if not self.merge_weights: + warnings.warn("Nothing to unmerge. Set merge_weights to True to enable (un)merging.") + return + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + if self.r[self.active_adapter] > 0: + self.weight.data -= ( + transpose( + self.lora_B[self.active_adapter].weight @ self.lora_A[self.active_adapter].weight, + self.fan_in_fan_out, + ) + * self.scaling[self.active_adapter] + ) + self.merged = False def forward(self, x: torch.Tensor): if self.disable_adapters: - if self.r > 0 and self.merged and any(self.enable_lora): - delta_w = ( - F.conv1d( - self.lora_A.weight.data.unsqueeze(0), - self.lora_B.weight.data, - groups=sum(self.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) - self.merged = False + if self.r[self.active_adapter] > 0 and self.merged: + self.unmerge() return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - elif self.merged: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - else: + elif self.r[self.active_adapter] > 0 and not self.merged: result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - if self.r > 0: - after_A = self.lora_A(self.lora_dropout(x)) - after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) - result += self.zero_pad(after_B) * self.scaling - return result + result += ( + self.lora_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) + ) + * self.scaling[self.active_adapter] + ) + else: + return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if is_bnb_available(): @@ -485,6 +421,7 @@ if is_bnb_available(): # Lora implemented in a dense layer def __init__( self, + adapter_name, in_features, out_features, r: int = 0, @@ -502,115 +439,37 @@ if is_bnb_available(): threshold=kwargs.get("threshold", 0.0), index=kwargs.get("index", None), ) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False) - # Actual trainable parameters - if r > 0: - self.lora_A = nn.Linear(in_features, r, bias=False) - self.lora_B = nn.Linear(r, out_features, bias=False) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - self.reset_parameters() + LoraLayer.__init__(self, merge_weights=False, in_features=in_features, out_features=out_features) - def reset_parameters(self): - if hasattr(self, "lora_A"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + + self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.active_adapter = adapter_name def forward(self, x: torch.Tensor): result = super().forward(x) if self.disable_adapters: return result - elif self.r > 0: + elif self.r[self.active_adapter] > 0: if not torch.is_autocast_enabled(): expected_dtype = result.dtype if x.dtype != torch.float32: x = x.float() - output = self.lora_B(self.lora_A(self.lora_dropout(x))).to(expected_dtype) * self.scaling - result += output + output = ( + self.lora_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) + ).to(expected_dtype) + * self.scaling[self.active_adapter] + ) else: - output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling - result += output - return result - - class MergedLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): - # Lora implemented in a dense layer - def __init__( - self, - in_features: int, - out_features: int, - r: int = 0, - lora_alpha: int = 1, - lora_dropout: float = 0.0, - enable_lora: List[bool] = [False], - **kwargs, - ): - bnb.nn.Linear8bitLt.__init__( - self, - in_features, - out_features, - bias=kwargs.get("bias", True), - has_fp16_weights=kwargs.get("has_fp16_weights", True), - memory_efficient_backward=kwargs.get("memory_efficient_backward", False), - threshold=kwargs.get("threshold", 0.0), - index=kwargs.get("index", None), - ) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False) - if out_features % len(enable_lora) != 0: - raise ValueError("The length of enable_lora must divide out_features") - self.enable_lora = enable_lora - # Actual trainable parameters - if r > 0 and any(enable_lora): - self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False) - self.lora_B = nn.Conv1d( - r * sum(enable_lora), - out_features // len(enable_lora) * sum(enable_lora), - kernel_size=1, - groups=2, - bias=False, - ) - self.scaling = self.lora_alpha / self.r - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - # Compute the indices - self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1) - self.lora_ind[enable_lora, :] = True - self.lora_ind = self.lora_ind.view(-1) - self.reset_parameters() - - def reset_parameters(self): - if hasattr(self, "lora_A"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) - - def zero_pad(self, x): - result = x.new_zeros((*x.shape[:-1], self.out_features)) - result = result.view(-1, self.out_features) - result[:, self.lora_ind] = x.reshape( - -1, self.out_features // len(self.enable_lora) * sum(self.enable_lora) - ) - return result.view((*x.shape[:-1], self.out_features)) - - def forward(self, x: torch.Tensor): - result = super().forward(x) - if self.disable_adapters: - return result - elif self.r > 0: - if not torch.is_autocast_enabled(): - expected_dtype = result.dtype - if x.dtype != torch.float32: - x = x.float() - after_A = self.lora_A(self.lora_dropout(x)) - after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) - output = self.zero_pad(after_B).to(expected_dtype) * self.scaling - result += output - else: - after_A = self.lora_A(self.lora_dropout(x)) - after_B = self.lora_B(after_A.transpose(-2, -1)).transpose(-2, -1) - output = self.zero_pad(after_B) * self.scaling - result += output + output = ( + self.lora_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) + ) + * self.scaling[self.active_adapter] + ) + result += output return result diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py index dd949c0..bfaabe8 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -17,14 +17,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .adapters_utils import CONFIG_NAME, WEIGHTS_NAME from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType from .other import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + CONFIG_NAME, + WEIGHTS_NAME, _set_trainable, bloom_model_postprocess_past_key_value, prepare_model_for_int8_training, shift_tokens_right, transpose, + _get_submodules, + _set_adapter, ) from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py deleted file mode 100644 index f2f8a95..0000000 --- a/src/peft/utils/adapters_utils.py +++ /dev/null @@ -1,18 +0,0 @@ -# coding=utf-8 -# Copyright 2023-present the HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -WEIGHTS_NAME = "adapter_model.bin" -CONFIG_NAME = "adapter_config.json" - -# TODO: add automapping and superclass here? diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index 3e2cf5b..544ec13 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -21,7 +21,7 @@ from typing import Optional, Union from huggingface_hub import hf_hub_download from transformers.utils import PushToHubMixin -from .adapters_utils import CONFIG_NAME +from .other import CONFIG_NAME class PeftType(str, enum.Enum): @@ -29,6 +29,7 @@ class PeftType(str, enum.Enum): P_TUNING = "P_TUNING" PREFIX_TUNING = "PREFIX_TUNING" LORA = "LORA" + MULTI_LORA = "MULTI_LORA" class TaskType(str, enum.Enum): @@ -82,7 +83,7 @@ class PeftConfigMixin(PushToHubMixin): writer.write(json.dumps(output_dict, indent=2, sort_keys=True)) @classmethod - def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None, **kwargs): r""" This method loads the configuration of your adapter model from a directory. @@ -92,11 +93,16 @@ class PeftConfigMixin(PushToHubMixin): **kwargs: Additional keyword arguments passed along to the child class initialization. """ - if os.path.isfile(os.path.join(pretrained_model_name_or_path, CONFIG_NAME)): - config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME) + path = ( + os.path.join(pretrained_model_name_or_path, subfolder) + if subfolder is not None + else pretrained_model_name_or_path + ) + if os.path.isfile(os.path.join(path, CONFIG_NAME)): + config_file = os.path.join(path, CONFIG_NAME) else: try: - config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME) + config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME, subfolder=subfolder) except Exception: raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'") diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 132b033..271bf89 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy + import torch @@ -86,11 +88,6 @@ def prepare_model_for_int8_training( return model -TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = { - "bloom": bloom_model_postprocess_past_key_value, -} - - # copied from transformers.models.bart.modeling_bart def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): """ @@ -113,11 +110,48 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start return shifted_input_ids -def _set_trainable(model): - if model.modules_to_save is not None: - for name, param in model.named_parameters(): - if any(module_name in name for module_name in model.modules_to_save): - param.requires_grad = True +class ModulesToSaveWrapper(torch.nn.Module): + def __init__(self, module_to_save, adapter_name): + super().__init__() + self.original_module = module_to_save + self.modules_to_save = torch.nn.ModuleDict({}) + self.update(adapter_name) + self.active_adapter = adapter_name + + def update(self, adapter_name): + self.modules_to_save.update(torch.nn.ModuleDict({adapter_name: copy.deepcopy(self.original_module)})) + + def forward(self, *args, **kwargs): + if self.active_adapter not in self.modules_to_save: + return self.original_module(*args, **kwargs) + return self.modules_to_save[self.active_adapter](*args, **kwargs) + + +def _get_submodules(model, key): + parent = model.get_submodule(".".join(key.split(".")[:-1])) + target_name = key.split(".")[-1] + target = model.get_submodule(key) + return parent, target, target_name + + +def _set_trainable(model, adapter_name): + key_list = [key for key, _ in model.named_modules()] + for key in key_list: + target_module_found = any(key.endswith(target_key) for target_key in model.modules_to_save) + if target_module_found: + parent, target, target_name = _get_submodules(key) + if isinstance(target, ModulesToSaveWrapper): + target.update(adapter_name) + else: + for param in target.parameters(): + param.requires_grad = True + setattr(parent, target_name, ModulesToSaveWrapper(target, adapter_name)) + + +def _set_adapter(model, adapter_name): + for module in model.modules(): + if isinstance(module, ModulesToSaveWrapper): + module.active_adapter = adapter_name def fsdp_auto_wrap_policy(model): @@ -157,3 +191,32 @@ def fsdp_auto_wrap_policy(model): def transpose(weight, fan_in_fan_out): return weight.T if fan_in_fan_out else weight + + +TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { + "t5": ["q", "v"], + "mt5": ["q", "v"], + "bart": ["q_proj", "v_proj"], + "gpt2": ["c_attn"], + "bloom": ["query_key_value"], + "opt": ["q_proj", "v_proj"], + "gptj": ["q_proj", "v_proj"], + "gpt_neox": ["query_key_value"], + "gpt_neo": ["q_proj", "v_proj"], + "bert": ["query", "value"], + "roberta": ["query", "value"], + "xlm-roberta": ["query", "value"], + "electra": ["query", "value"], + "deberta-v2": ["query_proj", "value_proj"], + "deberta": ["in_proj"], + "layoutlm": ["query", "value"], + "llama": ["q_proj", "v_proj"], + "chatglm": ["query_key_value"], +} + +TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = { + "bloom": bloom_model_postprocess_past_key_value, +} + +WEIGHTS_NAME = "adapter_model.bin" +CONFIG_NAME = "adapter_config.json" diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index c6596c7..43680fe 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -13,10 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .config import PeftType +from .config import PeftType, PromptLearningConfig -def get_peft_model_state_dict(model, state_dict=None): +def get_peft_model_state_dict(model, adapter_name, state_dict=None): """ Get the state dict of the Peft model. @@ -27,13 +27,14 @@ def get_peft_model_state_dict(model, state_dict=None): The state dict of the model. If not provided, the state dict of the model will be used. """ + config = model.peft_config[adapter_name] if state_dict is None: state_dict = model.state_dict() - if model.peft_config.peft_type == PeftType.LORA: + if config.peft_type == PeftType.LORA: # to_return = lora_state_dict(model, bias=model.peft_config.bias) # adapted from `https://github.com/microsoft/LoRA/blob/main/loralib/utils.py` - # to directly with the state dict which is necessary when using DeepSpeed or FSDP - bias = model.peft_config.bias + # to be used directly with the state dict which is necessary when using DeepSpeed or FSDP + bias = config.bias if bias == "none": to_return = {k: state_dict[k] for k in state_dict if "lora_" in k} elif bias == "all": @@ -48,21 +49,26 @@ def get_peft_model_state_dict(model, state_dict=None): to_return[bias_name] = state_dict[bias_name] else: raise NotImplementedError - else: + to_return = {k: v for k, v in to_return.items() if (("lora_" in k and adapter_name in k) or ("bias" in k))} + elif isinstance(config, PromptLearningConfig): to_return = {} - if model.peft_config.inference_mode: + if config.inference_mode: prompt_embeddings = model.prompt_encoder.embedding.weight else: - prompt_embeddings = model.get_prompt_embedding_to_save() + prompt_embeddings = model.get_prompt_embedding_to_save(adapter_name) to_return["prompt_embeddings"] = prompt_embeddings + else: + raise NotImplementedError if model.modules_to_save is not None: for key, value in state_dict.items(): - if any(module_name in key for module_name in model.modules_to_save): - to_return[key] = value + if any(f"{module_name}.modules_to_save.{adapter_name}" in key for module_name in model.modules_to_save): + to_return[key.replace("modules_to_save.", "")] = value + + to_return = {k.replace(f"{adapter_name}.", ""): v for k, v in to_return.items()} return to_return -def set_peft_model_state_dict(model, peft_model_state_dict): +def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): """ Set the state dict of the Peft model. @@ -70,10 +76,33 @@ def set_peft_model_state_dict(model, peft_model_state_dict): model ([`PeftModel`]): The Peft model. peft_model_state_dict (`dict`): The state dict of the Peft model. """ + config = model.peft_config[adapter_name] + state_dict = {} + if model.modules_to_save is not None: + for key, value in peft_model_state_dict.items(): + if any(module_name in key for module_name in model.modules_to_save): + for module_name in model.modules_to_save: + if module_name in key: + key = key.replace(module_name, f"{module_name}.modules_to_save.{adapter_name}") + break + state_dict[key] = value + + if config.peft_type == PeftType.LORA: + peft_model_state_dict = {} + for k, v in state_dict.items(): + if "lora_" in k: + suffix_to_replace = ".".join(k.split("lora_")[1].split(".")[1:]) + k = k.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + peft_model_state_dict[k] = v + else: + peft_model_state_dict[k] = v + elif isinstance(config, PromptLearningConfig): + peft_model_state_dict = state_dict + else: + raise NotImplementedError model.load_state_dict(peft_model_state_dict, strict=False) - if model.peft_config.peft_type != PeftType.LORA: - model.prompt_encoder.embedding.load_state_dict( + if isinstance(config, PromptLearningConfig): + model.prompt_encoder[adapter_name].embedding.load_state_dict( {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True ) - return model From af252b709bd8f69dc47c10efb0c9798ec7f639de Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:29:24 +0530 Subject: [PATCH 02/35] Update peft_model.py --- src/peft/peft_model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index cbc8a89..b54a148 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -86,9 +86,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.active_adapter = adapter_name if not isinstance(peft_config, PromptLearningConfig): self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( - self.base_model, peft_config, adapter_name + self.base_model, self.peft_config, adapter_name ) - self.add_adapter(adapter_name, peft_config) + else: + self.add_adapter(adapter_name, peft_config) def save_pretrained(self, save_directory, **kwargs): r""" From 7d7c598647a68418705231286d1e6f90eb44646f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:32:21 +0530 Subject: [PATCH 03/35] Update peft_model.py --- src/peft/peft_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index b54a148..71b9ef2 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -85,6 +85,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.peft_config = {} self.active_adapter = adapter_name if not isinstance(peft_config, PromptLearningConfig): + self.peft_config[adapter_name] = peft_config self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( self.base_model, self.peft_config, adapter_name ) From 64cae2aab2174214987758f0486c3d6ae8a96563 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:34:04 +0530 Subject: [PATCH 04/35] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 34cee7d..5459381 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -139,7 +139,7 @@ class LoraModel(torch.nn.Module): if target_module_found: if not is_target_modules_in_base_model: is_target_modules_in_base_model = True - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) bias = target.bias is not None if isinstance(target, LoraLayer): target.update_layer(adapter_name, lora_config.r, lora_config.lora_alpha, lora_config.lora_dropout) From e9d45da4c5868c4a5b0624101a39bdebc2d4d138 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:35:49 +0530 Subject: [PATCH 05/35] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 5459381..215a315 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -361,7 +361,7 @@ class Linear(nn.Linear): self.weight.data = self.weight.data.T nn.Linear.reset_parameters(self) - self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.update_layer(adapter_name, r, lora_alpha, lora_dropout) self.active_adapter = adapter_name def merge(self): From 8ec7cb84350239d6ca51aa1223c91d40714f6a1b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:36:41 +0530 Subject: [PATCH 06/35] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 215a315..058a3f3 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -444,7 +444,7 @@ if is_bnb_available(): # Freezing the pre-trained weight matrix self.weight.requires_grad = False - self.update_layer(self, adapter_name, r, lora_alpha, lora_dropout) + self.update_layer(adapter_name, r, lora_alpha, lora_dropout) self.active_adapter = adapter_name def forward(self, x: torch.Tensor): From 090d0743992c5ac43ae931b8e02d109b1b27d900 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:38:22 +0530 Subject: [PATCH 07/35] Update lora.py --- src/peft/tuners/lora.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 058a3f3..45d88ec 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -325,8 +325,8 @@ class LoraLayer: self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer})) # Actual trainable parameters if r > 0: - self.lora_A.update(nn.ModuleDict({nn.Linear(self.in_features, r, bias=False)})) - self.lora_B.update(nn.ModuleDict({nn.Linear(r, self.out_features, bias=False)})) + self.lora_A.update(nn.ModuleDict({adapter_name: nn.Linear(self.in_features, r, bias=False)})) + self.lora_B.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)})) self.scaling[adapter_name] = lora_alpha / r self.reset_lora_parameters(adapter_name) From 7c8ee5814a1a4ced386fc39cc13b3caf83a30930 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 19:40:06 +0530 Subject: [PATCH 08/35] Update peft_model.py --- src/peft/peft_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 71b9ef2..f2274fe 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -324,8 +324,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.add_adapter(adapter_name, peft_config) # load weights if any - if kwargs.get("subfolder", None) is not None: - path = os.path.join(model_id, kwargs["subfolder"]) + path = os.path.join(model_id, kwargs["subfolder"]) if kwargs.get("subfolder", None) is not None else model_id + if os.path.exists(os.path.join(path, WEIGHTS_NAME)): filename = os.path.join(path, WEIGHTS_NAME) else: From 002da1b450a1c3d4e206f9bd1556a45b7e90e8e6 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 28 Mar 2023 20:19:06 +0530 Subject: [PATCH 09/35] fix bugs --- src/peft/tuners/lora.py | 1 + src/peft/utils/save_and_load.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 45d88ec..4739dce 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -329,6 +329,7 @@ class LoraLayer: self.lora_B.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)})) self.scaling[adapter_name] = lora_alpha / r self.reset_lora_parameters(adapter_name) + self.to(self.weight.device) def reset_lora_parameters(self, adapter_name): if adapter_name in self.lora_A.keys(): diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 43680fe..fb4b252 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -86,6 +86,8 @@ def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): key = key.replace(module_name, f"{module_name}.modules_to_save.{adapter_name}") break state_dict[key] = value + else: + state_dict = peft_model_state_dict if config.peft_type == PeftType.LORA: peft_model_state_dict = {} @@ -100,7 +102,6 @@ def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): peft_model_state_dict = state_dict else: raise NotImplementedError - model.load_state_dict(peft_model_state_dict, strict=False) if isinstance(config, PromptLearningConfig): model.prompt_encoder[adapter_name].embedding.load_state_dict( From bd80d61b2a77f439c2e3428151c164840a5c020d Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 17:44:24 +0530 Subject: [PATCH 10/35] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 126 ++++++++++++++++++++++------------------ src/peft/tuners/lora.py | 12 +--- 2 files changed, 70 insertions(+), 68 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 4df5cd9..b98fc9f 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -84,6 +84,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.peft_config = {} self.active_adapter = adapter_name + self.peft_type = peft_config.peft_type if not isinstance(peft_config, PromptLearningConfig): self.peft_config[adapter_name] = peft_config self.base_model = PEFT_TYPE_TO_MODEL_MAPPING[peft_config.peft_type]( @@ -120,10 +121,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if peft_config.base_model_name_or_path is None: peft_config.base_model_name_or_path = ( self.base_model.__dict__.get("name_or_path", None) - if isinstance(self.peft_config, PromptLearningConfig) + if isinstance(peft_config, PromptLearningConfig) else self.base_model.model.__dict__.get("name_or_path", None) ) - inference_mode = self.peft_config.inference_mode + inference_mode = peft_config.inference_mode peft_config.inference_mode = True peft_config.save_pretrained(output_dir) peft_config.inference_mode = inference_mode @@ -213,32 +214,33 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Returns the virtual prompts to use for Peft. Only applicable when `peft_config.peft_type != PeftType.LORA`. """ + peft_config = self.active_peft_config prompt_encoder = self.prompt_encoder[self.active_adapter] prompt_tokens = self.prompt_tokens[self.active_adapter].unsqueeze(0).expand(batch_size, -1).to(self.device) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: - prompt_tokens = prompt_tokens[:, : self.peft_config.num_virtual_tokens] - if self.peft_config.inference_mode: + if peft_config.peft_type == PeftType.PREFIX_TUNING: + prompt_tokens = prompt_tokens[:, : peft_config.num_virtual_tokens] + if peft_config.inference_mode: past_key_values = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: past_key_values = prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, - self.peft_config.num_virtual_tokens, - self.peft_config.num_layers * 2, - self.peft_config.num_attention_heads, - self.peft_config.token_dim // self.peft_config.num_attention_heads, + peft_config.num_virtual_tokens, + peft_config.num_layers * 2, + peft_config.num_attention_heads, + peft_config.token_dim // peft_config.num_attention_heads, ) - if self.peft_config.num_transformer_submodules == 2: + if peft_config.num_transformer_submodules == 2: past_key_values = torch.cat([past_key_values, past_key_values], dim=2) past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( - self.peft_config.num_transformer_submodules * 2 + peft_config.num_transformer_submodules * 2 ) if TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING.get(self.config.model_type, None) is not None: post_process_fn = TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING[self.config.model_type] past_key_values = post_process_fn(past_key_values) return past_key_values else: - if self.peft_config.inference_mode: + if peft_config.inference_mode: prompts = prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: prompts = prompt_encoder(prompt_tokens) @@ -281,13 +283,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Disables the adapter module. """ - if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): + if isinstance(self.active_peft_config, PromptLearningConfig): old_forward = self.forward self.forward = self.base_model.forward else: self.base_model.disable_adapter_layers() yield - if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig): + if isinstance(self.active_peft_config, PromptLearningConfig): self.forward = old_forward else: self.base_model.enable_adapter_layers() @@ -296,13 +298,14 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Returns the base model. """ - return ( - self.base_model - if isinstance(self.peft_config[self.active_adapter], PromptLearningConfig) - else self.base_model.model - ) + return self.base_model if isinstance(self.active_peft_config, PromptLearningConfig) else self.base_model.model def add_adapter(self, adapter_name, peft_config): + if peft_config.peft_type != self.peft_type: + raise ValueError( + f"Cannot combine adapters with different peft types. " + f"Found {self.peft_type} and {peft_config.peft_type}." + ) self.peft_config[adapter_name] = peft_config if isinstance(peft_config, PromptLearningConfig): self._setup_prompt_encoder(adapter_name) @@ -380,11 +383,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): **dispatch_model_kwargs, ) hook = AlignDevicesHook(io_same_device=True) - if not isinstance(self.peft_config[adapter_name]) == PeftType.LORA: - add_hook_to_module(self.base_model.model, hook) - else: + if isinstance(self.peft_config[adapter_name], PromptLearningConfig): remove_hook_from_submodules(self.prompt_encoder) - add_hook_to_module(self.base_model, hook) + add_hook_to_module(self.get_base_model(), hook) def set_adapter(self, adapter_name): """ @@ -397,6 +398,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): self.base_model.set_adapter(adapter_name) _set_adapter(self, adapter_name) + @property + def active_peft_config(self): + return self.peft_config[self.active_adapter] + class PeftModelForSequenceClassification(PeftModel): """ @@ -465,8 +470,8 @@ class PeftModelForSequenceClassification(PeftModel): **kwargs, ): return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if not isinstance(self.peft_config, PromptLearningConfig): + peft_config = self.active_peft_config + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -481,7 +486,7 @@ class PeftModelForSequenceClassification(PeftModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -496,13 +501,13 @@ class PeftModelForSequenceClassification(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: return self._prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.peft_config.num_virtual_tokens).to(self.device), + torch.zeros(batch_size, peft_config.num_virtual_tokens).to(self.device), kwargs["token_type_ids"], ), dim=1, @@ -638,7 +643,8 @@ class PeftModelForCausalLM(PeftModel): return_dict=None, **kwargs, ): - if not isinstance(self.peft_config, PromptLearningConfig): + peft_config = self.active_peft_config + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -653,7 +659,7 @@ class PeftModelForCausalLM(PeftModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -672,7 +678,7 @@ class PeftModelForCausalLM(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.base_model(input_ids=input_ids, past_key_values=past_key_values, **kwargs) else: @@ -680,7 +686,7 @@ class PeftModelForCausalLM(PeftModel): inputs_embeds = self.word_embeddings(input_ids) # concat prompt labels if labels is not None: - prefix_labels = torch.full((batch_size, self.peft_config.num_virtual_tokens), -100).to(self.device) + prefix_labels = torch.full((batch_size, peft_config.num_virtual_tokens), -100).to(self.device) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) prompts = prompts.to(inputs_embeds.dtype) @@ -688,9 +694,10 @@ class PeftModelForCausalLM(PeftModel): return self.base_model(inputs_embeds=inputs_embeds, **kwargs) def generate(self, **kwargs): + peft_config = self.active_peft_config self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation try: - if not isinstance(self.peft_config, PromptLearningConfig): + if not isinstance(peft_config, PromptLearningConfig): outputs = self.base_model.generate(**kwargs) else: if "input_ids" not in kwargs: @@ -698,13 +705,13 @@ class PeftModelForCausalLM(PeftModel): # For gpt2 models, we construct postion_ids on the fly by using attention mask, and position ids need to match input_shape. # for prefix tuning, input shape is determined using `input_ids`. Thus we should not expand 'attention_mask' here # for prompt tuning input_ids is not passed but a concatenated input_embeds is passed. Thus attention_mask needs to be of same size of num_virtual_tokens + input_ids - if kwargs.get("attention_mask", None) is not None and self.peft_config.peft_type in [ + if kwargs.get("attention_mask", None) is not None and peft_config.peft_type in [ PeftType.PROMPT_TUNING, PeftType.P_TUNING, ]: # concat prompt attention mask prefix_attention_mask = torch.ones( - kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens + kwargs["input_ids"].shape[0], peft_config.num_virtual_tokens ).to(kwargs["input_ids"].device) kwargs["attention_mask"] = torch.cat((prefix_attention_mask, kwargs["attention_mask"]), dim=1) @@ -728,17 +735,18 @@ class PeftModelForCausalLM(PeftModel): return outputs def prepare_inputs_for_generation(self, *args, **kwargs): + peft_config = self.active_peft_config model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs) - if isinstance(self.peft_config, PromptLearningConfig): - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if isinstance(peft_config, PromptLearningConfig): + if peft_config.peft_type == PeftType.PREFIX_TUNING: prefix_attention_mask = torch.ones( - model_kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens + model_kwargs["input_ids"].shape[0], peft_config.num_virtual_tokens ).to(model_kwargs["input_ids"].device) model_kwargs["attention_mask"] = torch.cat( (prefix_attention_mask, model_kwargs["attention_mask"]), dim=1 ) - if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0]) model_kwargs["past_key_values"] = past_key_values else: @@ -810,7 +818,8 @@ class PeftModelForSeq2SeqLM(PeftModel): return_dict=None, **kwargs, ): - if not isinstance(self.peft_config, PromptLearningConfig): + peft_config = self.active_peft_config + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -828,7 +837,7 @@ class PeftModelForSeq2SeqLM(PeftModel): batch_size = input_ids.shape[0] if decoder_attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -848,7 +857,7 @@ class PeftModelForSeq2SeqLM(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.base_model( input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs @@ -864,35 +873,36 @@ class PeftModelForSeq2SeqLM(PeftModel): if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1) # concat prompt labels if labels is not None: - if self.peft_config.num_transformer_submodules == 1: + if peft_config.num_transformer_submodules == 1: kwargs["labels"] = labels - elif self.peft_config.num_transformer_submodules == 2: - prefix_labels = torch.full((batch_size, self.peft_config.num_virtual_tokens), -100).to(self.device) + elif peft_config.num_transformer_submodules == 2: + prefix_labels = torch.full((batch_size, peft_config.num_virtual_tokens), -100).to(self.device) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) prompts = prompts.to(inputs_embeds.dtype) - inputs_embeds = torch.cat((prompts[:, : self.peft_config.num_virtual_tokens], inputs_embeds), dim=1) - if self.peft_config.num_transformer_submodules == 1: + inputs_embeds = torch.cat((prompts[:, : peft_config.num_virtual_tokens], inputs_embeds), dim=1) + if peft_config.num_transformer_submodules == 1: return self.base_model(inputs_embeds=inputs_embeds, **kwargs) - elif self.peft_config.num_transformer_submodules == 2: + elif peft_config.num_transformer_submodules == 2: decoder_inputs_embeds = torch.cat( - (prompts[:, self.peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 + (prompts[:, peft_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 ) return self.base_model( inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs ) def generate(self, **kwargs): + peft_config = self.active_peft_config self.base_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation self.base_model._prepare_encoder_decoder_kwargs_for_generation = ( self._prepare_encoder_decoder_kwargs_for_generation ) try: - if not isinstance(self.peft_config, PromptLearningConfig): + if not isinstance(peft_config, PromptLearningConfig): outputs = self.base_model.generate(**kwargs) else: if "input_ids" not in kwargs: @@ -908,7 +918,7 @@ class PeftModelForSeq2SeqLM(PeftModel): ) kwargs["token_type_ids"] = None - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: outputs = self.base_model.generate(**kwargs) else: raise NotImplementedError @@ -926,8 +936,9 @@ class PeftModelForSeq2SeqLM(PeftModel): return outputs def prepare_inputs_for_generation(self, *args, **kwargs): + peft_config = self.active_peft_config model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs) - if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING: batch_size = model_kwargs["decoder_input_ids"].shape[0] past_key_values = self.get_prompt(batch_size) model_kwargs["past_key_values"] = past_key_values @@ -1000,9 +1011,10 @@ class PeftModelForTokenClassification(PeftModel): return_dict=None, **kwargs, ): + peft_config = self.active_peft_config return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if not isinstance(self.peft_config, PromptLearningConfig): + if not isinstance(peft_config, PromptLearningConfig): return self.base_model( input_ids=input_ids, attention_mask=attention_mask, @@ -1017,7 +1029,7 @@ class PeftModelForTokenClassification(PeftModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(self.device) + prefix_attention_mask = torch.ones(batch_size, peft_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -1032,13 +1044,13 @@ class PeftModelForTokenClassification(PeftModel): } ) - if self.peft_config.peft_type == PeftType.PREFIX_TUNING: + if peft_config.peft_type == PeftType.PREFIX_TUNING: return self._prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.peft_config.num_virtual_tokens).to(self.device), + torch.zeros(batch_size, peft_config.num_virtual_tokens).to(self.device), kwargs["token_type_ids"], ), dim=1, diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 977a3b7..0b96a77 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -323,17 +323,7 @@ class LoraModel(torch.nn.Module): if isinstance(target, LoraLayer): bias = target.bias is not None new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias) - - # manually merge if not merged - if not target.merged: - # merge weights per: https://arxiv.org/pdf/2106.09685.pdf / page 4 - if target.r > 0: - target.weight.data += ( - transpose(target.lora_B.weight @ target.lora_A.weight, target.fan_in_fan_out) - * target.scaling - ).to(target.weight.dtype) - target.merged = True - + target.merge() self._replace_module(parent, target_name, new_module, target) return self.model From 96ca100e34e1d0207a517d34ccc9a73fb4abf69e Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 18:03:32 +0530 Subject: [PATCH 11/35] Update lora.py --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 0b96a77..aab5cdf 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -166,7 +166,7 @@ class LoraModel(torch.nn.Module): "fan_in_fan_out": lora_config.fan_in_fan_out, "merge_weights": (lora_config.merge_weights or lora_config.inference_mode) and not is_hf_device_map_available, - "init_lora_weights": self.peft_config.init_lora_weights, + "init_lora_weights": lora_config.init_lora_weights, } key_list = [key for key, _ in self.model.named_modules()] for key in key_list: From d4b64c82801b9a939bd909d94d147307c07c7926 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 18:27:23 +0530 Subject: [PATCH 12/35] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index aab5cdf..417bb7a 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -434,6 +434,8 @@ class Linear(nn.Linear): self.active_adapter = adapter_name def merge(self): + if self.active_adapter not in self.lora_A.keys(): + return if not self.merge_weights: warnings.warn("Nothing to merge. Set merge_weights to True to enable merging.") return @@ -451,6 +453,8 @@ class Linear(nn.Linear): self.merged = True def unmerge(self): + if self.active_adapter not in self.lora_A.keys(): + return if not self.merge_weights: warnings.warn("Nothing to unmerge. Set merge_weights to True to enable (un)merging.") return @@ -468,6 +472,8 @@ class Linear(nn.Linear): self.merged = False def forward(self, x: torch.Tensor): + if self.active_adapter not in self.lora_A.keys(): + return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if self.disable_adapters: if self.r[self.active_adapter] > 0 and self.merged: self.unmerge() @@ -520,7 +526,7 @@ if is_bnb_available(): def forward(self, x: torch.Tensor): result = super().forward(x) - if self.disable_adapters: + if self.disable_adapters or self.active_adapter not in self.lora_A.keys(): return result elif self.r[self.active_adapter] > 0: if not torch.is_autocast_enabled(): From 18ccde8e86a1dcebc1d2ddf85350d09a341342d2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 19:44:58 +0530 Subject: [PATCH 13/35] =?UTF-8?q?fixing=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 6 ++++-- src/peft/tuners/lora.py | 2 +- src/peft/utils/save_and_load.py | 4 ++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index b98fc9f..1c3e7e1 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -112,7 +112,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module): for adapter_name, peft_config in self.peft_config.items(): # save only the trainable weights - output_state_dict = get_peft_model_state_dict(self, adapter_name, kwargs.get("state_dict", None)) + output_state_dict = get_peft_model_state_dict( + self, state_dict=kwargs.get("state_dict", None), adapter_name=adapter_name + ) output_dir = os.path.join(save_directory, adapter_name) if adapter_name != "default" else save_directory os.makedirs(output_dir, exist_ok=True) torch.save(output_state_dict, os.path.join(output_dir, WEIGHTS_NAME)) @@ -346,7 +348,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu") ) # load the weights into the model - set_peft_model_state_dict(self, adapter_name, adapters_weights) + set_peft_model_state_dict(self, adapters_weights, adapter_name=adapter_name) if ( (getattr(self, "hf_device_map", None) is not None) and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 417bb7a..0023717 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -136,7 +136,7 @@ class LoraModel(torch.nn.Module): self.model = model self.forward = self.model.forward self.config = config - self.add_adapter(adapter_name) + self.add_adapter(adapter_name, self.config[adapter_name]) def add_adapter(self, adapter_name, config=None): if config is not None: diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index fb4b252..7ebdfaf 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -16,7 +16,7 @@ from .config import PeftType, PromptLearningConfig -def get_peft_model_state_dict(model, adapter_name, state_dict=None): +def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): """ Get the state dict of the Peft model. @@ -68,7 +68,7 @@ def get_peft_model_state_dict(model, adapter_name, state_dict=None): return to_return -def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict): +def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="default"): """ Set the state dict of the Peft model. From 122f708ae8551f5b88cfbe3b9579fcdb3cb1f7ec Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 4 Apr 2023 20:05:59 +0530 Subject: [PATCH 14/35] =?UTF-8?q?=F0=9F=98=85.=20Fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 2 +- src/peft/utils/save_and_load.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 0023717..63201cc 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -404,7 +404,7 @@ class LoraLayer: nn.init.zeros_(self.lora_B[adapter_name].weight) -class Linear(nn.Linear): +class Linear(nn.Linear, LoraLayer): # Lora implemented in a dense layer def __init__( self, diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 7ebdfaf..a258c32 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -53,7 +53,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): elif isinstance(config, PromptLearningConfig): to_return = {} if config.inference_mode: - prompt_embeddings = model.prompt_encoder.embedding.weight + prompt_embeddings = model.prompt_encoder[adapter_name].embedding.weight else: prompt_embeddings = model.get_prompt_embedding_to_save(adapter_name) to_return["prompt_embeddings"] = prompt_embeddings From d4c2bc60e4bcbe47d0ba25cae23783b1503c5b49 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:04:42 +0530 Subject: [PATCH 15/35] =?UTF-8?q?fix=20more=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 63201cc..13ae8f5 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -135,22 +135,22 @@ class LoraModel(torch.nn.Module): super().__init__() self.model = model self.forward = self.model.forward - self.config = config - self.add_adapter(adapter_name, self.config[adapter_name]) + self.peft_config = config + self.add_adapter(adapter_name, self.peft_config[adapter_name]) def add_adapter(self, adapter_name, config=None): if config is not None: config = self._prepare_lora_config(config, self.model.config.to_dict()) - self.config[adapter_name] = config + self.peft_config[adapter_name] = config self._find_and_replace(adapter_name) - if len(self.config) > 1 and self.config[adapter_name].bias != "none": + if len(self.peft_config) > 1 and self.peft_config[adapter_name].bias != "none": raise ValueError( "LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters." ) - mark_only_lora_as_trainable(self.model, self.config[adapter_name].bias) + mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias) def _find_and_replace(self, adapter_name): - lora_config = self.config[adapter_name] + lora_config = self.peft_config[adapter_name] loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) if loaded_in_8bit and not is_bnb_available(): raise ImportError( @@ -260,7 +260,7 @@ class LoraModel(torch.nn.Module): def get_peft_config_as_dict(self, inference: bool = False): config_dict = {} - for key, value in self.config.items(): + for key, value in self.peft_config.items(): config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} if inference: config["inference_mode"] = True From 41b2fd770f50aa1c632068ec865a34ab2379dc34 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:13:41 +0530 Subject: [PATCH 16/35] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 13ae8f5..35f355e 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -477,7 +477,7 @@ class Linear(nn.Linear, LoraLayer): if self.disable_adapters: if self.r[self.active_adapter] > 0 and self.merged: self.unmerge() - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) elif self.r[self.active_adapter] > 0 and not self.merged: result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) result += ( @@ -487,7 +487,8 @@ class Linear(nn.Linear, LoraLayer): * self.scaling[self.active_adapter] ) else: - return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + return result if is_bnb_available(): From dbdb8f375708862d9e286a2962613025eb24fbae Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:19:19 +0530 Subject: [PATCH 17/35] =?UTF-8?q?fix=20more=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 35f355e..53cbba7 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -319,7 +319,7 @@ class LoraModel(torch.nn.Module): key_list = [key for key, _ in self.model.named_modules() if "lora" not in key] for key in key_list: - parent, target, target_name = self._get_submodules(key) + parent, target, target_name = _get_submodules(key) if isinstance(target, LoraLayer): bias = target.bias is not None new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias) From 6f1f26f426861e402ddcf67ff23a26b66f62e877 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:30:07 +0530 Subject: [PATCH 18/35] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/lora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 53cbba7..74ed008 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -319,7 +319,7 @@ class LoraModel(torch.nn.Module): key_list = [key for key, _ in self.model.named_modules() if "lora" not in key] for key in key_list: - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) if isinstance(target, LoraLayer): bias = target.bias is not None new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias) From b9433a82082f7bb56504d5e0d04558f69d5e929f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 01:39:51 +0530 Subject: [PATCH 19/35] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/utils/other.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 7cd285f..53e2ee8 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -139,7 +139,7 @@ def _set_trainable(model, adapter_name): for key in key_list: target_module_found = any(key.endswith(target_key) for target_key in model.modules_to_save) if target_module_found: - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(model, key) if isinstance(target, ModulesToSaveWrapper): target.update(adapter_name) else: From 75131959d1d5b9c00acd7a59aefcc196b54d5ef9 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 02:03:31 +0530 Subject: [PATCH 20/35] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 1c3e7e1..fbe98e4 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -352,7 +352,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if ( (getattr(self, "hf_device_map", None) is not None) and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) - and len(self.peft_config == 1) + and len(self.peft_config) == 1 ): device_map = kwargs.get("device_map", "auto") max_memory = kwargs.get("max_memory", None) From 44f3e86b6285901f4adea14c88642f0acbc1c70c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 03:35:08 +0530 Subject: [PATCH 21/35] Update config.py --- src/peft/utils/config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index 927daf0..34e98e9 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -29,7 +29,6 @@ class PeftType(str, enum.Enum): P_TUNING = "P_TUNING" PREFIX_TUNING = "PREFIX_TUNING" LORA = "LORA" - MULTI_LORA = "MULTI_LORA" class TaskType(str, enum.Enum): From 405f68f54abae0ffb281c16b1abbdb86bd79001f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 12:40:02 +0530 Subject: [PATCH 22/35] fix doc failure --- docs/source/package_reference/tuners.mdx | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/source/package_reference/tuners.mdx b/docs/source/package_reference/tuners.mdx index 2ec0824..9404266 100644 --- a/docs/source/package_reference/tuners.mdx +++ b/docs/source/package_reference/tuners.mdx @@ -14,8 +14,6 @@ For finetuning a model with LoRA. [[autodoc]] tuners.lora.Linear -[[autodoc]] tuners.lora.MergedLinear - ## P-tuning [[autodoc]] tuners.p_tuning.PromptEncoderConfig From d936aa9349a25e7acc2e954070652ac6297babdf Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 5 Apr 2023 19:54:02 +0530 Subject: [PATCH 23/35] fix tests --- tests/testing_common.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/testing_common.py b/tests/testing_common.py index cfe6cf2..38aee14 100644 --- a/tests/testing_common.py +++ b/tests/testing_common.py @@ -268,8 +268,12 @@ class PeftCommonTester: logits_transformers = transformers_model(**dummy_input)[0] - self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4)) - self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10)) + if config.merge_weights: + self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4)) + self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10)) + else: + self.assertFalse(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4)) + self.assertTrue(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10)) with tempfile.TemporaryDirectory() as tmp_dirname: model.save_pretrained(tmp_dirname) From 37e1f9ba340b2025ff2f415827abdc613e1acbf2 Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Wed, 5 Apr 2023 17:17:01 +0000 Subject: [PATCH 24/35] fix test --- tests/test_encoder_decoder_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py index 6cb5a36..c7c73e6 100644 --- a/tests/test_encoder_decoder_models.py +++ b/tests/test_encoder_decoder_models.py @@ -22,7 +22,7 @@ from .testing_common import PeftCommonTester, PeftTestConfigManager PEFT_ENCODER_DECODER_MODELS_TO_TEST = [ - "hf-internal-testing/tiny-random-T5ForConditionalGeneration", + "ybelkada/tiny-random-T5ForConditionalGeneration-calibrated", "hf-internal-testing/tiny-random-BartForConditionalGeneration", ] @@ -74,7 +74,7 @@ class PeftEncoderDecoderModelTester(unittest.TestCase, PeftCommonTester): PeftTestConfigManager.get_grid_parameters( { "model_ids": PEFT_ENCODER_DECODER_MODELS_TO_TEST, - "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [True, False]}, + "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]}, "task_type": "SEQ_2_SEQ_LM", }, ) From 739716043504f8f8980c51d6202be0e1df7a39ec Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:06:33 +0530 Subject: [PATCH 25/35] making adalora compatible with multiple adapters --- src/peft/tuners/adalora.py | 433 ++++++++++++++++++--------- src/peft/tuners/lora.py | 27 +- src/peft/utils/__init__.py | 2 + src/peft/utils/other.py | 26 ++ tests/test_decoder_models.py | 2 +- tests/test_encoder_decoder_models.py | 2 +- tests/testing_common.py | 8 +- 7 files changed, 328 insertions(+), 172 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index f98b3c6..938b83e 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -1,14 +1,27 @@ import importlib import re +import warnings from dataclasses import dataclass, field from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F +from transformers.pytorch_utils import Conv1D -from ..utils import PeftType, transpose -from .lora import LoraConfig, LoraLayer, LoraModel, mark_only_lora_as_trainable +from ..utils import ( + TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING, + PeftType, + _freeze_adapter, + _get_submodules, + transpose, +) +from .lora import ( + LoraConfig, + LoraLayer, + LoraModel, + mark_only_lora_as_trainable, +) def is_bnb_available(): @@ -78,17 +91,42 @@ class AdaLoraModel(LoraModel): - **peft_config** ([`AdaLoraConfig`]): The configuration of the AdaLora model. """ - def __init__(self, config, model): + def __init__(self, model, config, adapter_name): nn.Module.__init__(self) - self.peft_config = config self.model = model - self._find_and_replace() - mark_only_lora_as_trainable(self.model, self.peft_config.bias) + self.peft_config = config self.rankallocator = RankAllocator(config, self.model) - if config.enable_lora is not None: - raise NotImplementedError("MergedLinear has not been implemented for AdaLoRA.") + self.add_adapter(adapter_name, self.peft_config[adapter_name]) - def _find_and_replace(self): + def add_adapter(self, adapter_name, config=None): + if config is not None: + config = self._prepare_adalora_config(config, self.model.config.to_dict()) + self.peft_config[adapter_name] = config + self._find_and_replace(adapter_name) + if len(self.peft_config) > 1 and self.peft_config[adapter_name].bias != "none": + raise ValueError( + "AdaLoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters." + ) + traininable_mode_counter = 0 + for config in self.peft_config.values(): + if not config.inference_mode: + traininable_mode_counter += 1 + + if traininable_mode_counter > 1: + raise ValueError( + "AdaLoraModel supports only 1 trainable adapter. " + "When using multiple adapters, set inference_mode to True for all adapters except the one you want to train." + ) + + if self.peft_config[adapter_name].inference_mode: + _freeze_adapter(self.model, adapter_name) + else: + self.trainable_adapter_name = adapter_name + mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias) + self.rankallocator = RankAllocator(self.model, self.peft_config[adapter_name], self.trainable_adapter_name) + + def _find_and_replace(self, adapter_name): + lora_config = self.peft_config[adapter_name] loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) if loaded_in_8bit and not is_bnb_available(): raise ImportError( @@ -97,39 +135,74 @@ class AdaLoraModel(LoraModel): ) is_target_modules_in_base_model = False kwargs = { - "r": self.peft_config.init_r, - "lora_alpha": self.peft_config.lora_alpha, - "lora_dropout": self.peft_config.lora_dropout, - "fan_in_fan_out": self.peft_config.fan_in_fan_out, - "merge_weights": self.peft_config.merge_weights or self.peft_config.inference_mode, + "r": lora_config.r, + "lora_alpha": lora_config.lora_alpha, + "lora_dropout": lora_config.lora_dropout, + "fan_in_fan_out": lora_config.fan_in_fan_out, + "init_lora_weights": lora_config.init_lora_weights, } key_list = [key for key, _ in self.model.named_modules()] for key in key_list: - if isinstance(self.peft_config.target_modules, str): - target_module_found = re.fullmatch(self.peft_config.target_modules, key) + if isinstance(lora_config.target_modules, str): + target_module_found = re.fullmatch(lora_config.target_modules, key) else: - target_module_found = any(key.endswith(target_key) for target_key in self.peft_config.target_modules) + target_module_found = any(key.endswith(target_key) for target_key in lora_config.target_modules) if target_module_found: if not is_target_modules_in_base_model: is_target_modules_in_base_model = True - parent, target, target_name = self._get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) bias = target.bias is not None - if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): - kwargs.update( - { - "has_fp16_weights": target.state.has_fp16_weights, - "memory_efficient_backward": target.state.memory_efficient_backward, - "threshold": target.state.threshold, - "index": target.index, - } + if isinstance(target, LoraLayer): + target.update_layer( + adapter_name, + lora_config.r, + lora_config.lora_alpha, + lora_config.lora_dropout, + lora_config.init_lora_weights, ) - new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) - elif isinstance(target, torch.nn.Linear): - new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs) - self._replace_module(parent, target_name, new_module, target) + else: + if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): + kwargs.update( + { + "has_fp16_weights": target.state.has_fp16_weights, + "memory_efficient_backward": target.state.memory_efficient_backward, + "threshold": target.state.threshold, + "index": target.index, + } + ) + new_module = SVDLinear8bitLt( + adapter_name, target.in_features, target.out_features, bias=bias, **kwargs + ) + else: + if isinstance(target, torch.nn.Linear): + in_features, out_features = target.in_features, target.out_features + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = False + elif isinstance(target, Conv1D): + in_features, out_features = ( + target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape + ) + if not kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to False but the target module is `Conv1D`. " + "Setting fan_in_fan_out to True." + ) + kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True + else: + raise ValueError( + f"Target module {target} is not supported. " + f"Currently, only `torch.nn.Linear` and `Conv1D` are supported." + ) + new_module = SVDLinear(adapter_name, in_features, out_features, bias=bias, **kwargs) + + self._replace_module(parent, target_name, new_module, target) if not is_target_modules_in_base_model: raise ValueError( - f"Target modules {self.peft_config.target_modules} not found in the base model. " + f"Target modules {lora_config.target_modules} not found in the base model. " f"Please check the target modules and try again." ) @@ -144,14 +217,14 @@ class AdaLoraModel(LoraModel): outputs = self.model.forward(*args, **kwargs) # Calculate the orthogonal regularization - orth_reg_weight = self.peft_config.orth_reg_weight + orth_reg_weight = self.peft_config[self.trainable_adapter_name].orth_reg_weight assert orth_reg_weight > 0 if hasattr(outputs, "loss"): regu_loss = 0 num_param = 0 for n, p in self.model.named_parameters(): - if "lora_A" in n or "lora_B" in n: + if ("lora_A" in n or "lora_B" in n) and self.trainable_adapter_name in n: 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)) I.requires_grad = False @@ -161,7 +234,7 @@ class AdaLoraModel(LoraModel): outputs.loss += orth_reg_weight * regu_loss return outputs - def _prepare_new_module(self, target, rank_idx): + def _prepare_new_module(self, target, rank_idx, adapter_name): if isinstance(rank_idx, list): rank = sum(rank_idx) elif isinstance(rank_idx, torch.Tensor): @@ -169,12 +242,14 @@ class AdaLoraModel(LoraModel): rank = rank_idx.sum().item() else: raise ValueError("Unexcepted type of rank_idx") + + lora_config = self.peft_config[adapter_name] kwargs = { "r": rank, - "lora_alpha": self.peft_config.lora_alpha, - "lora_dropout": self.peft_config.lora_dropout, - "fan_in_fan_out": self.peft_config.fan_in_fan_out, - "merge_weights": self.peft_config.merge_weights or self.peft_config.inference_mode, + "lora_alpha": lora_config.lora_alpha, + "lora_dropout": lora_config.lora_dropout, + "fan_in_fan_out": lora_config.fan_in_fan_out, + "init_lora_weights": lora_config.init_lora_weights, } bias = target.bias is not None loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) @@ -187,9 +262,9 @@ class AdaLoraModel(LoraModel): "index": target.index, } ) - new_module = SVDLinear8bitLt(target.in_features, target.out_features, bias=bias, **kwargs) + new_module = SVDLinear8bitLt(adapter_name, target.in_features, target.out_features, bias=bias, **kwargs) elif isinstance(target, torch.nn.Linear): - new_module = SVDLinear(target.in_features, target.out_features, bias=bias, **kwargs) + new_module = SVDLinear(adapter_name, target.in_features, target.out_features, bias=bias, **kwargs) new_module = new_module.to(target.weight.device) with torch.no_grad(): @@ -197,120 +272,195 @@ class AdaLoraModel(LoraModel): if bias: new_module.bias.copy_(target.bias) if rank > 0: - new_module.lora_E.copy_(target.lora_E[rank_idx]) - new_module.lora_A.copy_(target.lora_A[rank_idx]) - new_module.lora_B.copy_(target.lora_B[:, rank_idx]) + new_module.lora_E[adapter_name].copy_(target.lora_E[rank_idx]) + new_module.lora_A[adapter_name].copy_(target.lora_A[rank_idx]) + new_module.lora_B[adapter_name].copy_(target.lora_B[:, rank_idx]) # The scaling is exactly as the previous - new_module.ranknum.copy_(target.ranknum) + new_module.ranknum[adapter_name].copy_(target.ranknum) return new_module - def resize_modules_by_rank_pattern(self, rank_pattern): + def resize_modules_by_rank_pattern(self, rank_pattern, adapter_name): for name, rank_idx in rank_pattern.items(): key = ".".join(name.split(".")[0:-1]) - parent, target, target_name = self._get_submodules(key) - new_module = self._prepare_new_module(target, rank_idx) + key = f"{key}.{adapter_name}" if adapter_name not in key else key + parent, target, target_name = _get_submodules(key) + new_module = self._prepare_new_module(target, rank_idx, adapter_name) self._replace_module(parent, target_name, new_module, target) def update_and_allocate(self, global_step): + lora_config = self.peft_config[self.trainable_adapter_name] # Update the importance score and allocate the budget - if global_step < self.peft_config.total_step - self.peft_config.tfinal: - budget, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step) + if global_step < lora_config.total_step - lora_config.tfinal: + _, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step) if rank_pattern: - self.peft_config.rank_pattern = rank_pattern + lora_config.rank_pattern = rank_pattern # Finalize the budget allocation - elif global_step == self.peft_config.total_step - self.peft_config.tfinal: - budget, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True) - self.resize_modules_by_rank_pattern(rank_pattern) - self.peft_config.rank_pattern = rank_pattern + elif global_step == lora_config.total_step - lora_config.tfinal: + _, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True) + self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name) + lora_config.rank_pattern = rank_pattern self.rankallocator.reset_ipt() # Pass the function and do forward propagation else: return None + @staticmethod + def _prepare_adalora_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING[ + model_config["model_type"] + ] + if peft_config.inference_mode: + peft_config.merge_weights = True + return peft_config -class SVDLinear(nn.Linear, LoraLayer): + +class AdaLoraLayer(LoraLayer): + def __init__( + self, + in_features: int, + out_features: int, + ): + super().__init__(in_features, out_features) + self.lora_E = nn.ParameterDict({}) + self.lora_A = nn.ParameterDict({}) + self.lora_B = nn.ParameterDict({}) + self.ranknum = nn.ParameterDict({}) + + def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights): + self.r[adapter_name] = r + self.lora_alpha[adapter_name] = lora_alpha + if lora_dropout > 0.0: + lora_dropout_layer = nn.Dropout(p=lora_dropout) + else: + + def lora_dropout_layer(x): + return x + + self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer})) + # Actual trainable parameters + if r > 0: + # Right singular vectors + self.lora_A.update( + nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))}) + ) + # Singular values + self.lora_E.update(nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros(r, 1))})) + # Left singular vectors + self.lora_B.update( + nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))}) + ) + # The current rank + self.ranknum.update( + nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(1), requires_grad=False)}) + ) + self.ranknum[adapter_name].data.fill_(float(self.r)) + self.ranknum[adapter_name].requires_grad = False + self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) + if init_lora_weights: + self.reset_lora_parameters(adapter_name) + self.to(self.weight.device) + + def reset_lora_parameters(self, adapter_name): + if adapter_name in self.lora_A.keys(): + nn.init.zeros_(self.lora_E[adapter_name]) + nn.init.normal_(self.lora_A[adapter_name], mean=0.0, std=0.02) + nn.init.normal_(self.lora_B[adapter_name], mean=0.0, std=0.02) + + +class SVDLinear(nn.Linear, AdaLoraLayer): # SVD-based adaptation by a dense layer def __init__( self, + adapter_name: str, in_features: int, out_features: int, r: int = 0, lora_alpha: int = 1, lora_dropout: float = 0.0, fan_in_fan_out: bool = False, - merge_weights: bool = True, **kwargs, ): + init_lora_weights = kwargs.pop("init_lora_weights", True) nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + AdaLoraLayer.__init__(self, in_features=in_features, out_features=out_features) + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False self.fan_in_fan_out = fan_in_fan_out - # Actual trainable parameters - if r > 0: - # Right singular vectors - self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) - # Singular values - self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) - # Left singular vectors - self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) - # The current rank - self.ranknum = nn.Parameter(self.weight.new_zeros(1), requires_grad=False) - self.ranknum.data.fill_(float(self.r)) - self.scaling = self.lora_alpha if self.lora_alpha > 0 else float(self.r) - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - self.ranknum.requires_grad = False - self.reset_parameters() if fan_in_fan_out: self.weight.data = self.weight.data.T - def reset_parameters(self): nn.Linear.reset_parameters(self) - if hasattr(self, "lora_A"): - nn.init.zeros_(self.lora_E) - nn.init.normal_(self.lora_A, mean=0.0, std=0.02) - nn.init.normal_(self.lora_B, mean=0.0, std=0.02) + self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights) + self.active_adapter = adapter_name - def train(self, mode: bool = True): - nn.Linear.train(self, mode) - if self.merge_weights and self.merged: - # Make sure that the weights are not merged - if self.r > 0: - self.weight.data -= ( - transpose(self.lora_B @ (self.lora_A * self.lora_E)) * self.scaling / (self.ranknum + 1e-5) - ) - self.merged = False - - def eval(self): - nn.Linear.eval(self) - if self.merge_weights and not self.merged: - # Merge the weights and mark it - if self.r > 0: - self.weight.data += ( - transpose(self.lora_B @ (self.lora_A * self.lora_E)) * self.scaling / (self.ranknum + 1e-5) + def merge(self): + if self.active_adapter not in self.lora_A.keys(): + return + if self.merged: + warnings.warn("Already merged. Nothing to do.") + return + if self.r[self.active_adapter] > 0: + self.weight.data += ( + transpose( + self.lora_B[self.active_adapter] + @ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]) ) + * self.scaling[self.active_adapter] + / (self.ranknum[self.active_adapter] + 1e-5) + ) self.merged = True - def forward(self, x: torch.Tensor): - if self.r > 0 and not self.merged: - result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) - if self.r > 0: - result += ( - (self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T) - * self.scaling - / (self.ranknum + 1e-5) + def unmerge(self): + if self.active_adapter not in self.lora_A.keys(): + return + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + if self.r[self.active_adapter] > 0: + self.weight.data -= ( + transpose( + self.lora_B[self.active_adapter] + @ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]) ) - return result - else: + * self.scaling[self.active_adapter] + / (self.ranknum[self.active_adapter] + 1e-5) + ) + self.merged = False + + def forward(self, x: torch.Tensor): + if self.active_adapter not in self.lora_A.keys(): return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + if self.disable_adapters: + if self.r[self.active_adapter] > 0 and self.merged: + self.unmerge() + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + elif self.r[self.active_adapter] > 0 and not self.merged: + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + result += ( + ( + self.lora_dropout[self.active_adapter](x) + @ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]).T + @ self.lora_B[self.active_adapter].T + ) + * self.scaling[self.active_adapter] + / (self.ranknum[self.active_adapter] + 1e-5) + ) + else: + result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) + return result if is_bnb_available(): - class SVDLinear8bitLt(bnb.nn.Linear8bitLt, LoraLayer): + class SVDLinear8bitLt(bnb.nn.Linear8bitLt, AdaLoraLayer): # Low-rank matrix for SVD-based adaptation def __init__( self, + adapter_name, in_features, out_features, r: int = 0, @@ -328,51 +478,45 @@ if is_bnb_available(): threshold=kwargs.get("threshold", 0.0), index=kwargs.get("index", None), ) - LoraLayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=False) - # Actual trainable parameters - if r > 0: - # Right singular vectors - self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) - # Singular values - self.lora_E = nn.Parameter(self.weight.new_zeros(r, 1)) - # Left singular vectors - self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) - # The current rank - self.ranknum = nn.Parameter(self.weight.new_zeros(1), requires_grad=False) - self.ranknum.data.fill_(float(self.r)) - self.scaling = self.lora_alpha if self.lora_alpha > 0 else float(self.r) - # Freezing the pre-trained weight matrix - self.weight.requires_grad = False - self.ranknum.requires_grad = False - self.reset_parameters() + AdaLoraLayer.__init__(self, in_features=in_features, out_features=out_features) + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False - def reset_parameters(self): - if hasattr(self, "lora_A"): - # initialize A the same way as the default for nn.Linear and B to zero - nn.init.zeros_(self.lora_E) - nn.init.normal_(self.lora_A, mean=0.0, std=0.02) - nn.init.normal_(self.lora_B, mean=0.0, std=0.02) + init_lora_weights = kwargs.pop("init_lora_weights", True) + self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights) + self.active_adapter = adapter_name def forward(self, x: torch.Tensor): result = super().forward(x) - if self.disable_adapters: + if self.disable_adapters or self.active_adapter not in self.lora_A.keys(): return result - elif self.r > 0: + elif self.r[self.active_adapter] > 0: if not torch.is_autocast_enabled(): expected_dtype = result.dtype if x.dtype != torch.float32: x = x.float() output = ( - self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5) - ).to(expected_dtype) * self.scaling - result += output + ( + self.lora_dropout[self.active_adapter](x) + @ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]).T + @ self.lora_B[self.active_adapter].T + ).to(expected_dtype) + * self.scaling[self.active_adapter] + / (self.ranknum[self.active_adapter] + 1e-5) + ) else: output = ( - self.lora_dropout(x) @ (self.lora_A * self.lora_E).T @ self.lora_B.T / (self.ranknum + 1e-5) - ) * self.scaling - result += output + ( + self.lora_dropout[self.active_adapter](x) + @ (self.lora_A[self.active_adapter] * self.lora_E[self.active_adapter]).T + @ self.lora_B[self.active_adapter].T + ) + * self.scaling[self.active_adapter] + / (self.ranknum[self.active_adapter] + 1e-5) + ) + result += output return result @@ -386,8 +530,9 @@ class RankAllocator(object): """ - def __init__(self, peft_config, model): + def __init__(self, model, peft_config, adapter_name): self.peft_config = peft_config + self.adapter_name = adapter_name self.beta1 = peft_config.beta1 self.beta2 = peft_config.beta2 assert self.beta1 > 0 and self.beta1 < 1 @@ -408,7 +553,7 @@ class RankAllocator(object): self.init_bgt = 0 self.name_set = set() for n, p in model.named_parameters(): - if "lora_A" in n: + if f"lora_A.{self.adapter_name}" in n: self.init_bgt += p.size(0) self.name_set.add(n.replace("lora_A", "%s")) self.name_set = sorted(self.name_set) @@ -437,7 +582,7 @@ class RankAllocator(object): def update_ipt(self, model): # Update the sensitivity and uncertainty for every weight for n, p in model.named_parameters(): - if "lora_" in n: + if "lora_" in n and self.adapter_name in n: if n not in self.ipt: self.ipt[n] = torch.zeros_like(p) self.exp_avg_ipt[n] = torch.zeros_like(p) @@ -465,7 +610,7 @@ class RankAllocator(object): triplet_ipt = {} # Get the importance score for A, E, B for n, p in model.named_parameters(): - if "lora_A" in n: + 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") @@ -473,7 +618,7 @@ class RankAllocator(object): vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) - if "lora_B" in n: + 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") @@ -481,7 +626,7 @@ class RankAllocator(object): vector_ipt[name_m] = [comb_ipt] else: vector_ipt[name_m].append(comb_ipt) - if "lora_E" in n: + 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 @@ -506,7 +651,7 @@ class RankAllocator(object): # Mask the unimportant triplets with torch.no_grad(): for n, p in model.named_parameters(): - if "lora_E" in n: + 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 diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 74ed008..90e4a23 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -29,6 +29,7 @@ from ..utils import ( TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, PeftConfig, PeftType, + _freeze_adapter, _get_submodules, transpose, ) @@ -52,8 +53,6 @@ class LoraConfig(PeftConfig): target_modules (`Union[List[str],str]`): The names of the modules to apply Lora to. lora_alpha (`float`): The alpha parameter for Lora scaling. lora_dropout (`float`): The dropout probability for Lora layers. - merge_weights (`bool`): - Whether to merge the weights of the Lora layers with the base transformer model in `eval` mode. fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out). For example, gpt-2 uses `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`.: bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only' @@ -71,9 +70,6 @@ class LoraConfig(PeftConfig): ) lora_alpha: int = field(default=None, metadata={"help": "Lora alpha"}) lora_dropout: float = field(default=None, metadata={"help": "Lora dropout"}) - merge_weights: bool = field( - default=False, metadata={"help": "Merge weights of the original model and the Lora model"} - ) fan_in_fan_out: bool = field( default=False, metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, @@ -147,7 +143,10 @@ class LoraModel(torch.nn.Module): raise ValueError( "LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters." ) - mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias) + if self.peft_config[adapter_name].inference_mode: + _freeze_adapter(self.model, adapter_name) + else: + mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias) def _find_and_replace(self, adapter_name): lora_config = self.peft_config[adapter_name] @@ -158,14 +157,11 @@ class LoraModel(torch.nn.Module): "You can install it with `pip install bitsandbytes`." ) is_target_modules_in_base_model = False - is_hf_device_map_available = hasattr(self.model, "hf_device_map") kwargs = { "r": lora_config.r, "lora_alpha": lora_config.lora_alpha, "lora_dropout": lora_config.lora_dropout, "fan_in_fan_out": lora_config.fan_in_fan_out, - "merge_weights": (lora_config.merge_weights or lora_config.inference_mode) - and not is_hf_device_map_available, "init_lora_weights": lora_config.init_lora_weights, } key_list = [key for key, _ in self.model.named_modules()] @@ -360,7 +356,6 @@ def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: class LoraLayer: def __init__( self, - merge_weights: bool, in_features: int, out_features: int, ): @@ -372,7 +367,6 @@ class LoraLayer: self.lora_B = nn.ModuleDict({}) # Mark the weight as unmerged self.merged = False - self.merge_weights = merge_weights self.disable_adapters = False self.in_features = in_features self.out_features = out_features @@ -415,13 +409,12 @@ class Linear(nn.Linear, LoraLayer): lora_alpha: int = 1, lora_dropout: float = 0.0, fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) - merge_weights: bool = True, **kwargs, ): init_lora_weights = kwargs.pop("init_lora_weights", True) nn.Linear.__init__(self, in_features, out_features, **kwargs) - LoraLayer.__init__(self, merge_weights=merge_weights, in_features=in_features, out_features=out_features) + LoraLayer.__init__(self, in_features=in_features, out_features=out_features) # Freezing the pre-trained weight matrix self.weight.requires_grad = False @@ -436,9 +429,6 @@ class Linear(nn.Linear, LoraLayer): def merge(self): if self.active_adapter not in self.lora_A.keys(): return - if not self.merge_weights: - warnings.warn("Nothing to merge. Set merge_weights to True to enable merging.") - return if self.merged: warnings.warn("Already merged. Nothing to do.") return @@ -455,9 +445,6 @@ class Linear(nn.Linear, LoraLayer): def unmerge(self): if self.active_adapter not in self.lora_A.keys(): return - if not self.merge_weights: - warnings.warn("Nothing to unmerge. Set merge_weights to True to enable (un)merging.") - return if not self.merged: warnings.warn("Already unmerged. Nothing to do.") return @@ -515,7 +502,7 @@ if is_bnb_available(): threshold=kwargs.get("threshold", 0.0), index=kwargs.get("index", None), ) - LoraLayer.__init__(self, merge_weights=False, in_features=in_features, out_features=out_features) + LoraLayer.__init__(self, in_features=in_features, out_features=out_features) # Freezing the pre-trained weight matrix self.weight.requires_grad = False diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py index bfaabe8..346b667 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -21,6 +21,7 @@ from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType from .other import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING, CONFIG_NAME, WEIGHTS_NAME, _set_trainable, @@ -30,5 +31,6 @@ from .other import ( transpose, _get_submodules, _set_adapter, + _freeze_adapter, ) from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 53e2ee8..1bfbbfb 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -134,6 +134,12 @@ def _get_submodules(model, key): return parent, target, target_name +def _freeze_adapter(model, adapter_name): + for n, p in model.named_parameters(): + if adapter_name in n: + p.requires_grad = False + + def _set_trainable(model, adapter_name): key_list = [key for key, _ in model.named_modules()] for key in key_list: @@ -199,6 +205,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "bart": ["q_proj", "v_proj"], "gpt2": ["c_attn"], "bloom": ["query_key_value"], + "blip-2": ["q", "v", "q_proj", "v_proj"], "opt": ["q_proj", "v_proj"], "gptj": ["q_proj", "v_proj"], "gpt_neox": ["query_key_value"], @@ -214,6 +221,25 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "chatglm": ["query_key_value"], } +TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING = { + "t5": ["q", "k", "v", "o", "wi", "wo"], + "mt5": ["q", "k", "v", "o", "wi_0", "wi_1", "wo"], + "bart": ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"], + # "gpt2": ["c_attn"], + # "bloom": ["query_key_value"], + "opt": ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"], + # "gptj": ["q_proj", "v_proj"], + # "gpt_neox": ["query_key_value"], + # "gpt_neo": ["q_proj", "v_proj"], + # "bert": ["query", "value"], + "roberta": ["query", "key", "value", "dense"], + # "xlm-roberta": ["query", "value"], + # "electra": ["query", "value"], + "deberta-v2": ["query_proj", "key_proj", "value_proj", "dense"], + # "deberta": ["in_proj"], + # "layoutlm": ["query", "value"], +} + TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = { "bloom": bloom_model_postprocess_past_key_value, } diff --git a/tests/test_decoder_models.py b/tests/test_decoder_models.py index f0a575e..209b4df 100644 --- a/tests/test_decoder_models.py +++ b/tests/test_decoder_models.py @@ -72,7 +72,7 @@ class PeftDecoderModelTester(unittest.TestCase, PeftCommonTester): PeftTestConfigManager.get_grid_parameters( { "model_ids": PEFT_DECODER_MODELS_TO_TEST, - "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]}, + "lora_kwargs": {"init_lora_weights": [False]}, "task_type": "CAUSAL_LM", }, ) diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py index c7c73e6..cdf9571 100644 --- a/tests/test_encoder_decoder_models.py +++ b/tests/test_encoder_decoder_models.py @@ -74,7 +74,7 @@ class PeftEncoderDecoderModelTester(unittest.TestCase, PeftCommonTester): PeftTestConfigManager.get_grid_parameters( { "model_ids": PEFT_ENCODER_DECODER_MODELS_TO_TEST, - "lora_kwargs": {"init_lora_weights": [False], "merge_weights": [False, True]}, + "lora_kwargs": {"init_lora_weights": [False]}, "task_type": "SEQ_2_SEQ_LM", }, ) diff --git a/tests/testing_common.py b/tests/testing_common.py index 38aee14..cfe6cf2 100644 --- a/tests/testing_common.py +++ b/tests/testing_common.py @@ -268,12 +268,8 @@ class PeftCommonTester: logits_transformers = transformers_model(**dummy_input)[0] - if config.merge_weights: - self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4)) - self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10)) - else: - self.assertFalse(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4)) - self.assertTrue(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10)) + self.assertTrue(torch.allclose(logits_lora, logits_merged, atol=1e-4, rtol=1e-4)) + self.assertFalse(torch.allclose(logits_merged, logits_transformers, atol=1e-10, rtol=1e-10)) with tempfile.TemporaryDirectory() as tmp_dirname: model.save_pretrained(tmp_dirname) From 74e2a3da50e37c700cade042b240d82d4e45a1e2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:16:37 +0530 Subject: [PATCH 26/35] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 938b83e..048832f 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -95,7 +95,6 @@ class AdaLoraModel(LoraModel): nn.Module.__init__(self) self.model = model self.peft_config = config - self.rankallocator = RankAllocator(config, self.model) self.add_adapter(adapter_name, self.peft_config[adapter_name]) def add_adapter(self, adapter_name, config=None): From b728f5f559c65e9051eb48fbea86874b068d3a5b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:20:13 +0530 Subject: [PATCH 27/35] =?UTF-8?q?=F0=9F=90=9B=20fixing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 048832f..6872f64 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -343,13 +343,13 @@ class AdaLoraLayer(LoraLayer): if r > 0: # Right singular vectors self.lora_A.update( - nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))}) + nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))}) ) # Singular values - self.lora_E.update(nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros(r, 1))})) + self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(r, 1))})) # Left singular vectors self.lora_B.update( - nn.ModuleDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))}) + nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))}) ) # The current rank self.ranknum.update( From dee2a96fea700743a723e608cfa9756e517b102c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:22:56 +0530 Subject: [PATCH 28/35] Update adalora.py --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 6872f64..53abce3 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -355,7 +355,7 @@ class AdaLoraLayer(LoraLayer): self.ranknum.update( nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(1), requires_grad=False)}) ) - self.ranknum[adapter_name].data.fill_(float(self.r)) + self.ranknum[adapter_name].data.fill_(float(r)) self.ranknum[adapter_name].requires_grad = False self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) if init_lora_weights: From b6c751455e9286290aa9da9d7f9dfe0c0148ee8f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:31:21 +0530 Subject: [PATCH 29/35] =?UTF-8?q?=F0=9F=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/peft_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 2b7b286..9bcaf80 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -331,7 +331,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if isinstance(peft_config, PromptLearningConfig) and is_trainable: raise ValueError("Cannot set a prompt learning adapter to trainable when loading pretrained adapter.") else: - peft_config[adapter_name].inference_mode = not is_trainable + peft_config.inference_mode = not is_trainable self.add_adapter(adapter_name, peft_config) # load weights if any From 07a4b8aacc840ca32fa3040a317644672da03e97 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 19:56:55 +0530 Subject: [PATCH 30/35] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 53abce3..8c84a55 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -134,7 +134,7 @@ class AdaLoraModel(LoraModel): ) is_target_modules_in_base_model = False kwargs = { - "r": lora_config.r, + "r": lora_config.init_r, "lora_alpha": lora_config.lora_alpha, "lora_dropout": lora_config.lora_dropout, "fan_in_fan_out": lora_config.fan_in_fan_out, @@ -154,7 +154,7 @@ class AdaLoraModel(LoraModel): if isinstance(target, LoraLayer): target.update_layer( adapter_name, - lora_config.r, + lora_config.init_r, lora_config.lora_alpha, lora_config.lora_dropout, lora_config.init_lora_weights, From 3aaf482704a5348e51d36819deac06d759d957b4 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 20:02:31 +0530 Subject: [PATCH 31/35] fix --- src/peft/utils/save_and_load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index dc0a391..22792e7 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -52,7 +52,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): if config.peft_type == PeftType.ADALORA: rank_pattern = config.rank_pattern if rank_pattern is not None: - rank_pattern = {k.replace(f"{adapter_name}.", ""): v for k, v in rank_pattern.items()} + rank_pattern = {k.replace(f".{adapter_name}", ""): v for k, v in rank_pattern.items()} config.rank_pattern = rank_pattern to_return = {k: v for k, v in to_return.items() if (("lora_" in k and adapter_name in k) or ("bias" in k))} From a591b4b905a419295a74f816c33c697d92a6a5eb Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 20:04:04 +0530 Subject: [PATCH 32/35] final fix I guess --- src/peft/tuners/adalora.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 8c84a55..06b18de 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -282,7 +282,7 @@ class AdaLoraModel(LoraModel): for name, rank_idx in rank_pattern.items(): key = ".".join(name.split(".")[0:-1]) key = f"{key}.{adapter_name}" if adapter_name not in key else key - parent, target, target_name = _get_submodules(key) + parent, target, target_name = _get_submodules(self.model, key) new_module = self._prepare_new_module(target, rank_idx, adapter_name) self._replace_module(parent, target_name, new_module, target) From 3258b709a3b08d6c6fb16f5fa02ca4665ed525af Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 20:41:36 +0530 Subject: [PATCH 33/35] =?UTF-8?q?fix=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 76 ++++++++++++++------------------------ 1 file changed, 27 insertions(+), 49 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 06b18de..a8cd3b7 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -233,58 +233,36 @@ class AdaLoraModel(LoraModel): outputs.loss += orth_reg_weight * regu_loss return outputs - def _prepare_new_module(self, target, rank_idx, adapter_name): - if isinstance(rank_idx, list): - rank = sum(rank_idx) - elif isinstance(rank_idx, torch.Tensor): - rank_idx = rank_idx.view(-1) - rank = rank_idx.sum().item() - else: - raise ValueError("Unexcepted type of rank_idx") - - lora_config = self.peft_config[adapter_name] - kwargs = { - "r": rank, - "lora_alpha": lora_config.lora_alpha, - "lora_dropout": lora_config.lora_dropout, - "fan_in_fan_out": lora_config.fan_in_fan_out, - "init_lora_weights": lora_config.init_lora_weights, - } - bias = target.bias is not None - loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) - if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt): - kwargs.update( - { - "has_fp16_weights": target.state.has_fp16_weights, - "memory_efficient_backward": target.state.memory_efficient_backward, - "threshold": target.state.threshold, - "index": target.index, - } - ) - new_module = SVDLinear8bitLt(adapter_name, target.in_features, target.out_features, bias=bias, **kwargs) - elif isinstance(target, torch.nn.Linear): - new_module = SVDLinear(adapter_name, target.in_features, target.out_features, bias=bias, **kwargs) - new_module = new_module.to(target.weight.device) - - with torch.no_grad(): - new_module.weight.copy_(target.weight) - if bias: - new_module.bias.copy_(target.bias) - if rank > 0: - new_module.lora_E[adapter_name].copy_(target.lora_E[rank_idx]) - new_module.lora_A[adapter_name].copy_(target.lora_A[rank_idx]) - new_module.lora_B[adapter_name].copy_(target.lora_B[:, rank_idx]) - # The scaling is exactly as the previous - new_module.ranknum[adapter_name].copy_(target.ranknum) - return new_module - def resize_modules_by_rank_pattern(self, rank_pattern, adapter_name): + lora_config = self.peft_config[adapter_name] for name, rank_idx in rank_pattern.items(): + if isinstance(rank_idx, list): + rank = sum(rank_idx) + elif isinstance(rank_idx, torch.Tensor): + rank_idx = rank_idx.view(-1) + rank = rank_idx.sum().item() + else: + raise ValueError("Unexcepted type of rank_idx") key = ".".join(name.split(".")[0:-1]) - key = f"{key}.{adapter_name}" if adapter_name not in key else key - parent, target, target_name = _get_submodules(self.model, key) - new_module = self._prepare_new_module(target, rank_idx, adapter_name) - self._replace_module(parent, target_name, new_module, target) + _, target, _ = _get_submodules(self.model, key) + lora_E_weights = target.lora_E[adapter_name][rank_idx] + lora_A_weights = target.lora_A[adapter_name][rank_idx] + lora_B_weights = target.lora_B[adapter_name][:, rank_idx] + ranknum = target.ranknum[adapter_name] + target.update_layer( + adapter_name, + rank, + lora_config.lora_alpha, + lora_config.lora_dropout, + lora_config.init_lora_weights, + ) + with torch.no_grad(): + if rank > 0: + target.lora_E[adapter_name].copy_(lora_E_weights) + target.lora_A[adapter_name].copy_(lora_A_weights) + target.lora_B[adapter_name].copy_(lora_B_weights) + # The scaling is exactly as the previous + target.ranknum[adapter_name].copy_(ranknum) def update_and_allocate(self, global_step): lora_config = self.peft_config[self.trainable_adapter_name] From d5feb8b787624bd9b886a4eb447eabf6d01b8bb2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 6 Apr 2023 21:17:54 +0530 Subject: [PATCH 34/35] =?UTF-8?q?fixing=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/peft/tuners/adalora.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index a8cd3b7..1bbf7a2 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -243,7 +243,7 @@ class AdaLoraModel(LoraModel): rank = rank_idx.sum().item() else: raise ValueError("Unexcepted type of rank_idx") - key = ".".join(name.split(".")[0:-1]) + key = ".".join(name.split(".")[0:-2]) _, target, _ = _get_submodules(self.model, key) lora_E_weights = target.lora_E[adapter_name][rank_idx] lora_A_weights = target.lora_A[adapter_name][rank_idx] @@ -320,19 +320,13 @@ class AdaLoraLayer(LoraLayer): # Actual trainable parameters if r > 0: # Right singular vectors - self.lora_A.update( - nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((r, self.in_features)))}) - ) + self.lora_A.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, self.in_features))})) # Singular values - self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(r, 1))})) + self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, 1))})) # Left singular vectors - self.lora_B.update( - nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros((self.out_features, r)))}) - ) + self.lora_B.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(self.out_features, r))})) # The current rank - self.ranknum.update( - nn.ParameterDict({adapter_name: nn.Parameter(self.weight.new_zeros(1), requires_grad=False)}) - ) + self.ranknum.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(1), requires_grad=False)})) self.ranknum[adapter_name].data.fill_(float(r)) self.ranknum[adapter_name].requires_grad = False self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) From e8b0085d2b09735566a38e420747222834b80262 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 7 Apr 2023 04:08:10 +0530 Subject: [PATCH 35/35] fixing adalora saving and loading --- src/peft/peft_model.py | 4 +-- src/peft/tuners/adalora.py | 61 +++++++++++++++++++++++++-------- src/peft/utils/save_and_load.py | 14 +++++--- 3 files changed, 58 insertions(+), 21 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 9bcaf80..dd68deb 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -393,8 +393,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): remove_hook_from_submodules(self.prompt_encoder) add_hook_to_module(self.get_base_model(), hook) - # Set model in evaluation mode to deactivate Dropout modules by default - self.eval() + # Set model in evaluation mode to deactivate Dropout modules by default + self.eval() def set_adapter(self, adapter_name): """ diff --git a/src/peft/tuners/adalora.py b/src/peft/tuners/adalora.py index 1bbf7a2..fc6261f 100644 --- a/src/peft/tuners/adalora.py +++ b/src/peft/tuners/adalora.py @@ -243,7 +243,7 @@ class AdaLoraModel(LoraModel): rank = rank_idx.sum().item() else: raise ValueError("Unexcepted type of rank_idx") - key = ".".join(name.split(".")[0:-2]) + key = ".".join(name.split(".")[0:-2]) if adapter_name in name else ".".join(name.split(".")[0:-1]) _, target, _ = _get_submodules(self.model, key) lora_E_weights = target.lora_E[adapter_name][rank_idx] lora_A_weights = target.lora_A[adapter_name][rank_idx] @@ -264,6 +264,22 @@ class AdaLoraModel(LoraModel): # The scaling is exactly as the previous target.ranknum[adapter_name].copy_(ranknum) + def resize_state_dict_by_rank_pattern(self, rank_pattern, state_dict, adapter_name): + for name, rank_idx in rank_pattern.items(): + rank = sum(rank_idx) + prefix = ".".join(name.split(".")[0:-2]) if adapter_name in name else ".".join(name.split(".")[0:-1]) + for layer in ["lora_E", "lora_A", "lora_B"]: + key = f"base_model.model.{prefix}.{layer}.{adapter_name}" + if layer != "lora_B": + state_dict[key] = ( + state_dict[key][rank_idx] if rank != state_dict[key].shape[0] else state_dict[key] + ) + else: + state_dict[key] = ( + state_dict[key][:, rank_idx] if rank != state_dict[key].shape[1] else state_dict[key] + ) + return state_dict + def update_and_allocate(self, global_step): lora_config = self.peft_config[self.trainable_adapter_name] # Update the importance score and allocate the budget @@ -274,9 +290,14 @@ class AdaLoraModel(LoraModel): # Finalize the budget allocation elif global_step == lora_config.total_step - lora_config.tfinal: _, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True) - self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name) + # for some reason, this freezes the trainable parameters and nothing gets updates + # self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name) lora_config.rank_pattern = rank_pattern self.rankallocator.reset_ipt() + # Currently using inefficient way to mask the unimportant weights using the rank pattern + # due to problem mentioned above + elif global_step > lora_config.total_step - lora_config.tfinal: + self.rankallocator.mask_using_rank_pattern(self.model, lora_config.rank_pattern) # Pass the function and do forward propagation else: return None @@ -318,18 +339,17 @@ class AdaLoraLayer(LoraLayer): self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer})) # Actual trainable parameters - if r > 0: - # Right singular vectors - self.lora_A.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, self.in_features))})) - # Singular values - self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, 1))})) - # Left singular vectors - self.lora_B.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(self.out_features, r))})) - # The current rank - self.ranknum.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(1), requires_grad=False)})) - self.ranknum[adapter_name].data.fill_(float(r)) - self.ranknum[adapter_name].requires_grad = False - self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) + # Right singular vectors + self.lora_A.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, self.in_features))})) + # Singular values + self.lora_E.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(r, 1))})) + # Left singular vectors + self.lora_B.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(self.out_features, r))})) + # The current rank + self.ranknum.update(nn.ParameterDict({adapter_name: nn.Parameter(torch.zeros(1), requires_grad=False)})) + self.ranknum[adapter_name].data.fill_(float(r)) + self.ranknum[adapter_name].requires_grad = False + self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) if init_lora_weights: self.reset_lora_parameters(adapter_name) self.to(self.weight.device) @@ -638,3 +658,16 @@ class RankAllocator(object): else: rank_pattern = None return budget, rank_pattern + + def mask_using_rank_pattern(self, model, rank_pattern): + # Mask the unimportant triplets + is_adapter_name_truncated = False + if self.adapter_name not in next(iter(rank_pattern.keys())): + is_adapter_name_truncated = True + + with torch.no_grad(): + for n, p in model.named_parameters(): + if f"lora_E.{self.adapter_name}" in n: + key = n if not is_adapter_name_truncated else n.replace(f".{self.adapter_name}", "") + mask = torch.Tensor(rank_pattern[key]).unsqueeze(-1).to(p.device) + p.masked_fill_(~mask.bool(), 0.0) diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py index 22792e7..2876bbe 100644 --- a/src/peft/utils/save_and_load.py +++ b/src/peft/utils/save_and_load.py @@ -49,13 +49,13 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): to_return[bias_name] = state_dict[bias_name] else: raise NotImplementedError + to_return = {k: v for k, v in to_return.items() if (("lora_" in k and adapter_name in k) or ("bias" in k))} if config.peft_type == PeftType.ADALORA: rank_pattern = config.rank_pattern if rank_pattern is not None: rank_pattern = {k.replace(f".{adapter_name}", ""): v for k, v in rank_pattern.items()} config.rank_pattern = rank_pattern - - to_return = {k: v for k, v in to_return.items() if (("lora_" in k and adapter_name in k) or ("bias" in k))} + to_return = model.resize_state_dict_by_rank_pattern(rank_pattern, to_return, adapter_name) elif isinstance(config, PromptLearningConfig): to_return = {} if config.inference_mode: @@ -70,7 +70,7 @@ def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"): if any(f"{module_name}.modules_to_save.{adapter_name}" in key for module_name in model.modules_to_save): to_return[key.replace("modules_to_save.", "")] = value - to_return = {k.replace(f"{adapter_name}.", ""): v for k, v in to_return.items()} + to_return = {k.replace(f".{adapter_name}", ""): v for k, v in to_return.items()} return to_return @@ -99,8 +99,12 @@ def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="defaul peft_model_state_dict = {} for k, v in state_dict.items(): if "lora_" in k: - suffix_to_replace = ".".join(k.split("lora_")[1].split(".")[1:]) - k = k.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + suffix = k.split("lora_")[1] + if "." in suffix: + suffix_to_replace = ".".join(suffix.split(".")[1:]) + k = k.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}") + else: + k = f"{k}.{adapter_name}" peft_model_state_dict[k] = v else: peft_model_state_dict[k] = v