From 6c9534e660de73fa44abaefbf1f4f5b0df818fbb Mon Sep 17 00:00:00 2001 From: younesbelkada Date: Sun, 29 Jan 2023 11:18:31 +0000 Subject: [PATCH] adapt for other models --- src/peft/mapping.py | 2 + src/peft/peft_model.py | 84 ++++++++++++++++++++++++++++- src/peft/tuners/lora.py | 73 +------------------------ tests/test_lora.py | 114 ++++++++++++++++++++++++---------------- 4 files changed, 154 insertions(+), 119 deletions(-) diff --git a/src/peft/mapping.py b/src/peft/mapping.py index 66c07ff..18e292b 100644 --- a/src/peft/mapping.py +++ b/src/peft/mapping.py @@ -138,4 +138,6 @@ def get_peft_model(model, peft_config): else: peft_config = _prepare_lora_config(peft_config, model_config) + peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None) + 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 d43baf8..f3e01c1 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -14,18 +14,31 @@ # limitations under the License. import inspect +import os import warnings import torch from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput, TokenClassifierOutput +from transformers.utils import PushToHubMixin + +from huggingface_hub import hf_hub_download from .tuners import LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder -from .utils import PeftConfig, PeftType, TaskType, _set_trainable, shift_tokens_right +from .utils import ( + WEIGHTS_NAME, + PeftConfig, + PeftType, + TaskType, + _set_trainable, + get_peft_model_state_dict, + set_peft_model_state_dict, + shift_tokens_right, +) -class PeftModel(torch.nn.Module): +class PeftModel(PushToHubMixin, torch.nn.Module): """ Parameter-Efficient Fine-Tuning Model. Base model encompassing various Peft methods. @@ -61,6 +74,73 @@ class PeftModel(torch.nn.Module): self.base_model = LoraModel(peft_config, model) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def save_pretrained(self, save_directory, **kwargs): + r""" + Args: + This function saves the adapter model and the adapter configuration files to a directory, so that it can be + re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub` + method. + save_directory (`str`): + Directory where the adapter model and configuration files will be saved (will be created if it does not + exist). + **kwargs: + Additional keyword arguments passed along to the `push_to_hub` method. + """ + if os.path.isfile(save_directory): + raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file") + os.makedirs(save_directory, exist_ok=True) + + # save the config + 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) + self.peft_config.inference_mode = True + self.peft_config.save_pretrained(save_directory) + + for param in self.parameters(): + param.requires_grad = False # freeze the model + + # 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)) + + @classmethod + def from_pretrained(cls, model, model_id, **kwargs): + r""" + Args: + Instantiate a `LoraModel` from a pretrained Lora configuration and weights. + model (`transformers.PreTrainedModel`): + The model to be adapted. The model should be initialized with the `from_pretrained` method. from + `transformers` library. + model_id (`str`): + The name of the Lora configuration to use. Can be either: + - A string, the `model id` of a Lora configuration hosted inside a model repo on + huggingface Hub + - A path to a directory containing a Lora configuration file saved using the + `save_pretrained` method, e.g., ``./my_lora_config_directory/``. + """ + 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) + + 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) + # load the weights into the model + return set_peft_model_state_dict(model, adapters_weights) + def _setup_prompt_encoder(self): num_transformer_submodules = 0 transformer_backbone = None diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 071bc90..a1eb3b1 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -14,7 +14,6 @@ # limitations under the License. import importlib import math -import os import warnings from dataclasses import asdict, dataclass, field from enum import Enum @@ -24,12 +23,10 @@ import torch import torch.nn as nn import torch.nn.functional as F from transformers.pytorch_utils import Conv1D -from transformers.utils import PushToHubMixin import bitsandbytes as bnb -from huggingface_hub import hf_hub_download -from ..utils import WEIGHTS_NAME, PeftConfig, PeftType, get_peft_model_state_dict, transpose +from ..utils import PeftConfig, PeftType, transpose def is_loralib_available(): @@ -76,7 +73,7 @@ class LoraConfig(PeftConfig): self.peft_type = PeftType.LORA -class LoraModel(PushToHubMixin, torch.nn.Module): +class LoraModel(torch.nn.Module): """ Creates Low Rank Adapter (Lora) model from a pretrained transformers model. @@ -165,72 +162,6 @@ class LoraModel(PushToHubMixin, torch.nn.Module): new_module.state = old_module.state new_module.to(old_module.weight.device) - def save_pretrained(self, save_directory, **kwargs): - r""" - This function saves the adapter model and the adapter configuration files to a directory, so that it can be - re-loaded using the `LoraModel.from_pretrained` class method, and also used by the `LoraModel.push_to_hub` - method. - - Args: - save_directory (`str`): - Directory where the adapter model and configuration files will be saved (will be created if it does not - exist). - **kwargs: - Additional keyword arguments passed along to the `push_to_hub` method. - """ - if os.path.isfile(save_directory): - raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file") - os.makedirs(save_directory, exist_ok=True) - - # save the config - self.peft_config.save_pretrained(save_directory) - - for param in self.parameters(): - param.requires_grad = False # freeze the model - - # save only the trainable weights - output_state_dict = get_peft_model_state_dict(self) - torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME)) - - @classmethod - def from_pretrained(cls, model, lora_id, **kwargs): - r""" - Instantiate a `LoraModel` from a pretrained Lora configuration and weights. - - Args: - model (`transformers.PreTrainedModel`): - The model to be adapted. The model should be initialized with the `from_pretrained` method. from - `transformers` library. - lora_id (`str`): - The name of the Lora configuration to use. Can be either: - - A string, the `model id` of a Lora configuration hosted inside a model repo on - huggingface Hub - - A path to a directory containing a Lora configuration file saved using the - `save_pretrained` method, e.g., ``./my_lora_config_directory/``. - """ - # load the config - config = LoraConfig.from_pretrained(lora_id) - - model = cls(config, model) - - # load weights if any - if os.path.exists(os.path.join(lora_id, WEIGHTS_NAME)): - filename = os.path.join(lora_id, WEIGHTS_NAME) - else: - try: - filename = hf_hub_download(lora_id, WEIGHTS_NAME) - except: # noqa - raise ValueError( - f"Can't find weights for {lora_id} in {lora_id} or in the Hugging Face Hub. " - f"Please check that the file {WEIGHTS_NAME} is present at {lora_id}." - ) - - adapters_weights = torch.load(filename) - # load the weights into the model - model.load_state_dict(adapters_weights, strict=False) - - return model - def __getattr__(self, name: str): """Forward missing attributes to the wrapped module.""" try: diff --git a/tests/test_lora.py b/tests/test_lora.py index c69a32a..17b710a 100644 --- a/tests/test_lora.py +++ b/tests/test_lora.py @@ -17,36 +17,62 @@ import torch import tempfile import unittest -from peft import LoraModel, LoraConfig, get_peft_model_state_dict +from peft import PeftConfig, PeftModel, LoraConfig, get_peft_model_state_dict, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from transformers import AutoModelForCausalLM class LoraTestMixin: checkpoints_to_test = [ - "trl-internal-testing/tiny-random-OPTForCausalLM", + "hf-internal-testing/tiny-random-OPTForCausalLM", ] + config_classes = ( + LoraConfig, + # PrefixTuningConfig, + # PromptEncoderConfig, + # PromptTuningConfig, + ) + config_kwargs = ( + dict( + r = 8, + lora_alpha=32, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", + ), + # dict( + # encoder_hidden_size=32, + # task_type="CAUSAL_LM", + # ), + # dict( + # encoder_hidden_size=32, + # task_type="CAUSAL_LM", + # ), + # dict( + # task_type="CAUSAL_LM", + # ) -class LoraTester(unittest.TestCase, LoraTestMixin): + ) + +class PeftModelTester(unittest.TestCase, LoraTestMixin): r""" - Test if the LoraModel behaves as expected. This includes: + Test if the PeftModel behaves as expected. This includes: - test if the model has the expected methods """ def test_attributes_lora_model(self): for model_id in self.checkpoints_to_test: model = AutoModelForCausalLM.from_pretrained(model_id) - config = LoraConfig( - r = 8, - lora_alpha=32, - target_modules=["q_proj", "v_proj"], - lora_dropout=0.05, - bias="none", - ) - model = LoraModel(config, model) + for i, config_cls in enumerate(self.config_classes): + config = config_cls( + base_model_name_or_path=model_id, + **self.config_kwargs[i], + ) + model = PeftModel(model, config) - self.assertTrue(hasattr(model, 'save_pretrained')) - self.assertTrue(hasattr(model, 'from_pretrained')) - self.assertTrue(hasattr(model, 'push_to_hub')) + self.assertTrue(hasattr(model, 'save_pretrained')) + self.assertTrue(hasattr(model, 'from_pretrained')) + self.assertTrue(hasattr(model, 'push_to_hub')) def test_save_pretrained(self): r""" @@ -62,42 +88,38 @@ class LoraTester(unittest.TestCase, LoraTestMixin): for model_id in self.checkpoints_to_test: model = AutoModelForCausalLM.from_pretrained(model_id) - config = LoraConfig( - r = 8, - lora_alpha=32, - target_modules=["q_proj", "v_proj"], - lora_dropout=0.05, - bias="none", - ) - model = LoraModel(config, model) + for i, config_cls in enumerate(self.config_classes): + config = config_cls( + base_model_name_or_path=model_id, + **self.config_kwargs[i], + ) + model = PeftModel(model, config) - with tempfile.TemporaryDirectory() as tmp_dirname: - model.save_pretrained(tmp_dirname) + with tempfile.TemporaryDirectory() as tmp_dirname: + model.save_pretrained(tmp_dirname) - model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id) - model_from_pretrained = LoraModel.from_pretrained(model_from_pretrained, tmp_dirname) - - # check if the state dicts are equal - state_dict = get_peft_model_state_dict(model) - state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained) + model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id) + model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname) + + # check if the state dicts are equal + state_dict = get_peft_model_state_dict(model) + state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained) - # check if same keys - self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys()) + # check if same keys + self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys()) - # check if tensors equal - for key in state_dict.keys(): - self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key])) + # check if tensors equal + for key in state_dict.keys(): + self.assertTrue(torch.allclose(state_dict[key], state_dict_from_pretrained[key])) - # check if `adapter_model.bin` is present - self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin"))) + # check if `adapter_model.bin` is present + self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin"))) - # check if `adapter_config.json` is present - self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json"))) - - # check if `pytorch_model.bin` is not present - self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin"))) - - # check if `config.json` is not present - self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json"))) + # check if `adapter_config.json` is present + self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json"))) + # check if `pytorch_model.bin` is not present + self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin"))) + # check if `config.json` is not present + self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json"))) \ No newline at end of file