diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py index cfd3fd8..5040dc9 100644 --- a/src/peft/tuners/lora.py +++ b/src/peft/tuners/lora.py @@ -12,7 +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. - +import os import importlib import math import warnings @@ -25,7 +25,9 @@ import torch.nn as nn import torch.nn.functional as F from transformers.pytorch_utils import Conv1D -from ..utils import PeftConfig, PeftType, transpose +from huggingface_hub import hf_hub_download + +from ..utils import PeftConfig, PeftType, transpose, WEIGHTS_NAME, CONFIG_NAME, get_peft_model_state_dict def is_loralib_available(): @@ -85,8 +87,9 @@ class LoraModel(torch.nn.Module): Example:: - >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig >>> from peft import LoraModel, LoraConfig >>> - config = LoraConfig( + >>> from transformers import AutoModelForSeq2SeqLM, LoraConfig + >>> from peft import LoraModel, LoraConfig + >>> config = LoraConfig( peft_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"], lora_dropout=0.01, ) >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoraModel(config, model) @@ -167,6 +170,57 @@ class LoraModel(torch.nn.Module): config["inference_mode"] = True return config + def save_pretrained(self, save_directory): + 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""" + + Args: + model (`transformers.PreTrainedModel`): + The model to be adapted. + lora_id (`str`): + The name of the Lora configuration to use. + """ + # 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 + + + # Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py # and modified to work with PyTorch FSDP diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py index 45b0cd5..5a35b1e 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -20,3 +20,4 @@ from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType from .other import _set_trainable, bloom_model_postprocess_past_key_value, shift_tokens_right, transpose from .save_and_load import get_peft_model_state_dict, peft_model_load_and_dispatch, set_peft_model_state_dict +from .adapters_utils import WEIGHTS_NAME, CONFIG_NAME \ No newline at end of file diff --git a/src/peft/utils/adapters_utils.py b/src/peft/utils/adapters_utils.py new file mode 100644 index 0000000..9f3761c --- /dev/null +++ b/src/peft/utils/adapters_utils.py @@ -0,0 +1,17 @@ +# 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" \ No newline at end of file diff --git a/src/peft/utils/config.py b/src/peft/utils/config.py index 17690a0..41bef20 100644 --- a/src/peft/utils/config.py +++ b/src/peft/utils/config.py @@ -12,11 +12,16 @@ # 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 json import enum -from dataclasses import dataclass, field +from dataclasses import dataclass, field, asdict from typing import Optional, Union +from huggingface_hub import hf_hub_download + +from .adapters_utils import CONFIG_NAME + class PeftType(str, enum.Enum): PROMPT_TUNING = "PROMPT_TUNING" @@ -31,9 +36,53 @@ class TaskType(str, enum.Enum): CAUSAL_LM = "CAUSAL_LM" TOKEN_CLS = "TOKEN_CLS" - @dataclass -class PeftConfig: +class PeftConfigMixin(object): + @property + def __dict__(self): + return asdict(self) + + def save_pretrained(self, save_directory): + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + + output_dict = self.__dict__ + output_path = os.path.join(save_directory, CONFIG_NAME) + + # save it + with open(output_path, "w") as writer: + writer.write(json.dumps(output_dict)) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + 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: + try: + config_file = hf_hub_download(pretrained_model_name_or_path, CONFIG_NAME) + except: + 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(): + if hasattr(config, key): + setattr(config, key, value) + + return config + + @classmethod + def from_json_file(cls, json_file, **kwargs): + with open(json_file, 'r') as file: + json_object = json.load(file) + + return json_object + + +class PeftConfig(PeftConfigMixin): """ This is the base configuration class to store the configuration of a :class:`~peft.PeftModel`. @@ -42,13 +91,11 @@ class PeftConfig: task_type (Union[[`~peft.utils.config.TaskType`], `str`]): The type of task to perform. inference_mode (`bool`, defaults to `False`): Whether to use the Peft model in inference mode. """ - peft_type: Union[str, PeftType] = field(default=None, metadata={"help": "Peft type"}) task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"}) inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"}) -@dataclass class PromptLearningConfig(PeftConfig): """ This is the base configuration class to store the configuration of a Union[[`~peft.PrefixTuning`],