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 diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 1878854..9c03978 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -39,47 +39,6 @@ PEFT_TYPE_TO_CONFIG_MAPPING = { "ADALORA": AdaLoraConfig, } -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"], - "blip-2": ["q", "v", "q_proj", "v_proj"], - "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_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"], -} - def get_peft_config(config_dict): """ @@ -134,29 +93,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 _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 - - def get_peft_model(model, peft_config): """ Returns a Peft model object from a model and a config. @@ -167,14 +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 isinstance(peft_config, AdaLoraConfig): - peft_config = _prepare_adalora_config(peft_config, model_config) - elif 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 c8dde62..dd68deb 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -28,7 +28,7 @@ from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput from transformers.utils import PushToHubMixin -from .tuners import AdaLoraConfig, AdaLoraModel, LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder +from .tuners import AdaLoraModel, LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import ( TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, @@ -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,15 @@ from .utils import ( ) +PEFT_TYPE_TO_MODEL_MAPPING = { + PeftType.LORA: LoraModel, + PeftType.PROMPT_TUNING: PromptEmbedding, + PeftType.P_TUNING: PromptEncoder, + PeftType.PREFIX_TUNING: PrefixEncoder, + PeftType.ADALORA: AdaLoraModel, +} + + class PeftModel(PushToHubMixin, torch.nn.Module): """ Base model encompassing various Peft methods. @@ -67,22 +77,21 @@ class PeftModel(PushToHubMixin, torch.nn.Module): in the base model if using [`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() - elif isinstance(self.peft_config, AdaLoraConfig): - self.base_model = AdaLoraModel(peft_config, model) + 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]( + self.base_model, self.peft_config, adapter_name + ) 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.add_adapter(adapter_name, peft_config) def save_pretrained(self, save_directory, **kwargs): r""" @@ -101,24 +110,29 @@ 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)) - - # 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) + for adapter_name, peft_config in self.peft_config.items(): + # save only the trainable weights + output_state_dict = get_peft_model_state_dict( + self, state_dict=kwargs.get("state_dict", None), adapter_name=adapter_name ) - 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 + 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 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(peft_config, PromptLearningConfig) + else self.base_model.model.__dict__.get("name_or_path", None) + ) + inference_mode = 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""" Instantiate a [`LoraModel`] from a pretrained Lora configuration and weights. @@ -136,75 +150,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) - offload_dir = kwargs.get("offload_folder", None) - offload_index = kwargs.get("offload_index", None) - - dispatch_model_kwargs = {} - # Safety checker for previous `accelerate` versions - # `offload_index` was introduced in https://github.com/huggingface/accelerate/pull/873/ - if "offload_index" in inspect.signature(dispatch_model).parameters: - dispatch_model_kwargs["offload_index"] = offload_index - - 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, - offload_dir=offload_dir, - **dispatch_model_kwargs, - ) - hook = AlignDevicesHook(io_same_device=True) - if model.peft_config.peft_type == PeftType.LORA: - add_hook_to_module(model.base_model.model, hook) - elif model.peft_config.peft_type == PeftType.ADALORA: - 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(): @@ -215,72 +180,72 @@ 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) - 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) + 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 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 = 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, - 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: - prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + if peft_config.inference_mode: + 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): @@ -320,13 +285,13 @@ class PeftModel(PushToHubMixin, torch.nn.Module): """ Disables the adapter module. """ - if isinstance(self.peft_config, 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, PromptLearningConfig): + if isinstance(self.active_peft_config, PromptLearningConfig): self.forward = old_forward else: self.base_model.enable_adapter_layers() @@ -335,7 +300,116 @@ 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.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) + 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, is_trainable=False, **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)) + 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.inference_mode = not is_trainable + self.add_adapter(adapter_name, peft_config) + + # load weights if any + 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: + 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, 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) + and len(self.peft_config) == 1 + ): + device_map = kwargs.get("device_map", "auto") + max_memory = kwargs.get("max_memory", None) + offload_dir = kwargs.get("offload_folder", None) + offload_index = kwargs.get("offload_index", None) + + dispatch_model_kwargs = {} + # Safety checker for previous `accelerate` versions + # `offload_index` was introduced in https://github.com/huggingface/accelerate/pull/873/ + if "offload_index" in inspect.signature(dispatch_model).parameters: + dispatch_model_kwargs["offload_index"] = offload_index + + 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, + offload_dir=offload_dir, + **dispatch_model_kwargs, + ) + hook = AlignDevicesHook(io_same_device=True) + if isinstance(self.peft_config[adapter_name], PromptLearningConfig): + 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() + + 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) + + @property + def active_peft_config(self): + return self.peft_config[self.active_adapter] class PeftModelForSequenceClassification(PeftModel): @@ -378,9 +452,12 @@ class PeftModelForSequenceClassification(PeftModel): ``` """ - 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): @@ -388,7 +465,7 @@ class PeftModelForSequenceClassification(PeftModel): break # to make sure classifier layer is trainable - _set_trainable(self) + _set_trainable(self, adapter_name) def forward( self, @@ -402,8 +479,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, @@ -418,7 +495,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.") @@ -433,13 +510,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, @@ -560,8 +637,8 @@ class PeftModelForCausalLM(PeftModel): ``` """ - 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( @@ -575,7 +652,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, @@ -590,7 +668,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: @@ -609,7 +687,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: @@ -617,7 +695,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) @@ -625,9 +703,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: @@ -635,13 +714,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) @@ -665,17 +744,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: @@ -726,8 +806,8 @@ class PeftModelForSeq2SeqLM(PeftModel): ``` """ - 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 @@ -747,7 +827,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, @@ -765,7 +846,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: @@ -785,7 +866,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 @@ -801,35 +882,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: @@ -845,7 +927,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 @@ -863,8 +945,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 @@ -911,9 +994,12 @@ class PeftModelForTokenClassification(PeftModel): ``` """ - 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): @@ -921,7 +1007,7 @@ class PeftModelForTokenClassification(PeftModel): break # to make sure classifier layer is trainable - _set_trainable(self) + _set_trainable(self, adapter_name) def forward( self, @@ -934,9 +1020,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, @@ -951,7 +1038,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.") @@ -966,13 +1053,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/adalora.py b/src/peft/tuners/adalora.py index f98b3c6..fc6261f 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,41 @@ 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.rankallocator = RankAllocator(config, self.model) - if config.enable_lora is not None: - raise NotImplementedError("MergedLinear has not been implemented for AdaLoRA.") + self.peft_config = config + 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 +134,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.init_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.init_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 +216,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,156 +233,225 @@ class AdaLoraModel(LoraModel): outputs.loss += orth_reg_weight * regu_loss return outputs - def _prepare_new_module(self, target, rank_idx): - 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") - 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, - } - 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(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 = 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.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]) - # The scaling is exactly as the previous - new_module.ranknum.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): + lora_config = self.peft_config[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) - self._replace_module(parent, target_name, new_module, target) + 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:-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] + 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 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 - 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) + # 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 + @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 + # 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) + + 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 +469,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 +521,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 +544,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 +573,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 +601,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 +609,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 +617,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 +642,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 @@ -522,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/tuners/lora.py b/src/peft/tuners/lora.py index ab5259d..90e4a23 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -25,7 +25,14 @@ 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, + _freeze_adapter, + _get_submodules, + transpose, +) def is_bnb_available(): @@ -46,12 +53,10 @@ 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`). - enable_lora ( `List[bool]`): Used with [`lora.MergedLinear`]. - 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 + 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' + modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint. """ @@ -65,14 +70,10 @@ 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)"}, ) - 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, @@ -126,15 +127,29 @@ class LoraModel(torch.nn.Module): - **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.peft_config = config + 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_lora_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( + "LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters." + ) + 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] loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False) if loaded_in_8bit and not is_bnb_available(): raise ImportError( @@ -142,71 +157,78 @@ 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": 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) - and not is_hf_device_map_available, - "init_lora_weights": self.peft_config.init_lora_weights, + "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, ) - 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 + 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 @@ -233,9 +255,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.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 + config_dict[key] = config return config def _set_adapter_layers(self, enabled=True): @@ -249,6 +274,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 + def merge_and_unload(self): r""" This method merges the LoRa layers into the base model. This is needed if someone wants to use the base model @@ -262,36 +315,11 @@ 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(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) - - # manually merge if not merged - if not target.merged: - # merge weights per: https://arxiv.org/pdf/2106.09685.pdf / page 4 - if isinstance(target, Linear): - 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) - else: - if target.r > 0: - delta_w = ( - F.conv1d( - target.lora_A.weight.data.unsqueeze(0), - target.lora_B.weight.data, - groups=sum(target.enable_lora), - ) - .squeeze(0) - .transpose(-2, -1) - ) - target.weight.data += transpose( - target.zero_pad(delta_w * target.scaling), not target.fan_in_fan_out - ).to(target.weight.dtype) - target.merged = True - + target.merge() self._replace_module(parent, target_name, new_module, target) return self.model @@ -328,224 +356,126 @@ 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, 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: + 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 + 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(): + # 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): # Lora implemented in 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, # 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, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + LoraLayer.__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: - 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 - if init_lora_weights: - 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(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) - 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 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].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, - ): - 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) - 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) - - if init_lora_weights: - 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) + + 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].weight @ self.lora_A[self.active_adapter].weight, + self.fan_in_fan_out, ) - self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) + * self.scaling[self.active_adapter] + ) 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.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 > 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) + 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_B[self.active_adapter]( + self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x)) ) - self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out) - self.merged = False - 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) + * self.scaling[self.active_adapter] + ) else: 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 + return result if is_bnb_available(): @@ -554,6 +484,7 @@ if is_bnb_available(): # Lora implemented in a dense layer def __init__( self, + adapter_name, in_features, out_features, r: int = 0, @@ -571,115 +502,38 @@ 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, 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 + + 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_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..346b667 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -17,14 +17,20 @@ # 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, + TRANSFORMERS_MODELS_TO_ADALORA_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, + _freeze_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 ceeb57e..31c4669 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): @@ -83,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. @@ -93,11 +93,16 @@ class PeftConfigMixin(PushToHubMixin): kwargs (additional keyword arguments, *optional*): 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_NAME}' at '{pretrained_model_name_or_path}'") diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 585da64..1bfbbfb 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,54 @@ 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 _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: + 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(model, 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 +197,52 @@ 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"], + "blip-2": ["q", "v", "q_proj", "v_proj"], + "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_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, +} + +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 0b10717..2876bbe 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, state_dict=None, adapter_name="default"): """ 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 in (PeftType.LORA, PeftType.ADALORA): + if config.peft_type in (PeftType.LORA, PeftType.ADALORA): # 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,32 @@ 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))} + 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 = model.resize_state_dict_by_rank_pattern(rank_pattern, to_return, adapter_name) + elif isinstance(config, PromptLearningConfig): to_return = {} - if model.peft_config.inference_mode: - prompt_embeddings = model.prompt_encoder.embedding.weight + if config.inference_mode: + prompt_embeddings = model.prompt_encoder[adapter_name].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, peft_model_state_dict, adapter_name="default"): """ Set the state dict of the Peft model. @@ -70,14 +82,43 @@ 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. """ - if model.peft_config.peft_type == PeftType.ADALORA: - rank_pattern = model.peft_config.rank_pattern - if rank_pattern: - model.base_model.resize_modules_by_rank_pattern(rank_pattern) - model.load_state_dict(peft_model_state_dict, strict=False) + 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 + else: + state_dict = peft_model_state_dict - if model.peft_config.peft_type not in (PeftType.LORA, PeftType.ADALORA): - model.prompt_encoder.embedding.load_state_dict( + if config.peft_type in (PeftType.LORA, PeftType.ADALORA): + peft_model_state_dict = {} + for k, v in state_dict.items(): + if "lora_" in k: + 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 + if config.peft_type == PeftType.ADALORA: + rank_pattern = config.rank_pattern + if rank_pattern is not None: + model.resize_modules_by_rank_pattern(rank_pattern, adapter_name) + elif isinstance(config, PromptLearningConfig): + 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( {"weight": peft_model_state_dict["prompt_embeddings"]}, strict=True ) - return model 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 6cb5a36..cdf9571 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]}, "task_type": "SEQ_2_SEQ_LM", }, )