diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index 5040dc9..17f75c3 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -27,6 +27,7 @@ from transformers.pytorch_utils import Conv1D from huggingface_hub import hf_hub_download +from transformers.utils import PushToHubMixin from ..utils import PeftConfig, PeftType, transpose, WEIGHTS_NAME, CONFIG_NAME, get_peft_model_state_dict @@ -38,6 +39,17 @@ if is_loralib_available(): import loralib as lora # noqa: F401 from loralib import mark_only_lora_as_trainable +MODEL_CARD_TEMPLATE = """--- +license: apache-2.0 +base_model: {base_model} +tags: +- peft +- lora +--- +# Lora adapters for {model_name} + +""" + @dataclass class LoraConfig(PeftConfig): @@ -74,7 +86,7 @@ class LoraConfig(PeftConfig): self.peft_type = PeftType.LORA -class LoraModel(torch.nn.Module): +class LoraModel(PushToHubMixin, torch.nn.Module): """ Creates Low Rank Adapter (Lora) model from a pretrained transformers model. @@ -150,27 +162,19 @@ class LoraModel(torch.nn.Module): if old_module.bias is not None: new_module.bias = old_module.bias - def forward(self, *args, **kwargs): - return self.model(*args, **kwargs) + 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. - def __getattr__(self, name: str): - """Forward missing attributes to the wrapped module.""" - try: - return super().__getattr__(name) # defer to nn.Module's logic - except AttributeError: - return getattr(self.model, name) - - @property - def modules_to_save(self): - 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 - return config - - def save_pretrained(self, save_directory): + 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) @@ -184,16 +188,31 @@ class LoraModel(torch.nn.Module): # 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)) + + # save model card + if 'name_or_path' in self.model.__dict__: + model_name = self.model.__dict__['name_or_path'] + else: + model_name = None + model_card_content = MODEL_CARD_TEMPLATE.format(model_name=model_name, base_model=model_name) + with open(os.path.join(save_directory, "README.md"), "w", encoding="utf-8") as f: + f.write(model_card_content) @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 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. + 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) @@ -220,6 +239,27 @@ class LoraModel(torch.nn.Module): return model + def forward(self, *args, **kwargs): + return self.model(*args, **kwargs) + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + return getattr(self.model, name) + + @property + def modules_to_save(self): + 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 + return config + + # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py index 9f3761c..fe906c4 100644 --- a/src/peft/utils/adapters_utils.py +++ b/src/peft/utils/adapters_utils.py @@ -12,6 +12,7 @@ # 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" \ No newline at end of file +CONFIG_NAME = "adapter_config.json" + +# TODO: add automapping and superclass here? \ No newline at end of file diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index ef985e3..ad30b32 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -18,11 +18,11 @@ import enum from dataclasses import dataclass, field, asdict from typing import Optional, Union +from transformers.utils import PushToHubMixin from huggingface_hub import hf_hub_download from .adapters_utils import CONFIG_NAME - class PeftType(str, enum.Enum): PROMPT_TUNING = "PROMPT_TUNING" P_TUNING = "P_TUNING" @@ -36,8 +36,19 @@ class TaskType(str, enum.Enum): CAUSAL_LM = "CAUSAL_LM" TOKEN_CLS = "TOKEN_CLS" + @dataclass -class PeftConfigMixin(object): +class PeftConfigMixin(PushToHubMixin): + r""" + This is the base configuration class for PEFT adapter models. It contains all the methods that + are common to all PEFT adapter models. + This class inherits from `transformers.utils.PushToHubMixin` which contains the methods to push your model to the Hub. + The method `save_pretrained` will save the configuration of your adapter model in a directory. + The method `from_pretrained` will load the configuration of your adapter model from a directory. + + Args: + peft_type (Union[[`~peft.utils.config.PeftType`], `str`]): The type of Peft method to use. + """ peft_type: Optional[PeftType] = field(default=None, metadata={"help": "The type of PEFT model."}) @property @@ -47,7 +58,16 @@ class PeftConfigMixin(object): def to_dict(self): return self.__dict__ - def save_pretrained(self, save_directory): + def save_pretrained(self, save_directory, **kwargs): + r""" + This method saves the configuration of your adapter model in a directory. + + Args: + save_directory (`str`): + The directory where the configuration will be saved. + **kwargs: + Additional keyword arguments passed along to the `transformers.utils.PushToHubMixin.push_to_hub` method. + """ if os.path.isfile(save_directory): raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") @@ -58,10 +78,19 @@ class PeftConfigMixin(object): # save it with open(output_path, "w") as writer: - writer.write(json.dumps(output_dict)) + writer.write(json.dumps(output_dict, indent=2, sort_keys=True)) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + r""" + This method loads the configuration of your adapter model from a directory. + + Args: + pretrained_model_name_or_path (`str`): + The directory or the hub-id where the configuration is saved. + **kwargs: + Additional keyword arguments passed along to the child class initialization. + """ if os.path.isfile(os.path.join(pretrained_model_name_or_path, CONFIG_NAME)): config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME) else: @@ -71,6 +100,7 @@ class PeftConfigMixin(object): raise ValueError(f"Can't find config.json at '{pretrained_model_name_or_path}'") loaded_attributes = cls.from_json_file(config_file) + config = cls(**kwargs) for key, value in loaded_attributes.items(): @@ -80,8 +110,15 @@ class PeftConfigMixin(object): return config @classmethod - def from_json_file(cls, json_file, **kwargs): - with open(json_file, 'r') as file: + def from_json_file(cls, path_json_file, **kwargs): + r""" + Loads a configuration file from a json file. + + Args: + path_json_file (`str`): + The path to the json file. + """ + with open(path_json_file, 'r') as file: json_object = json.load(file) return json_object diff --git a/tests/test_config.py b/tests/test_config.py index 7d3df71..98881e9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,10 +1,24 @@ +# 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. import unittest import tempfile import os from peft import LoraConfig, PromptEncoderConfig, PrefixTuningConfig, PromptTuningConfig -class PeftConfigMixin: +class PeftConfigTestMixin: all_config_classes = ( LoraConfig, PromptEncoderConfig, @@ -13,7 +27,7 @@ class PeftConfigMixin: ) -class PeftConfigTester(unittest.TestCase, PeftConfigMixin): +class PeftConfigTester(unittest.TestCase, PeftConfigTestMixin): def test_methods(self): r""" Test if all configs have the expected methods. Here we test diff --git a/tests/test_lora.py b/tests/test_lora.py new file mode 100644 index 0000000..30a8268 --- /dev/null +++ b/tests/test_lora.py @@ -0,0 +1,110 @@ +# 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. +import os +import torch +import tempfile +import unittest + +from peft import LoraModel, LoraConfig, get_peft_model_state_dict + +from transformers import AutoModelForCausalLM + +class LoraTestMixin: + checkpoints_to_test = [ + "trl-internal-testing/tiny-random-OPTForCausalLM", + ] + +class LoraTester(unittest.TestCase, LoraTestMixin): + r""" + Test if the LoraModel 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) + + 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""" + A test to check if `save_pretrained` behaves as expected. This function + should only save the state dict of the adapter model and not the state + dict of the base model. Hence inside each saved directory you should have: + + - README.md (that contains an entry `base_model`) + - adapter_config.json + - adapter_model.bin + + """ + 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) + + 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) + + # 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 `README.md` is present + self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "README.md"))) + # check if `base_model` attribute is in `README.md` + with open(os.path.join(tmp_dirname, "README.md"), "r") as f: + readme = f.read() + self.assertTrue("base_model" in readme) + + # 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"))) + +