add lora support

README.md

@@ -1,7 +1,8 @@
-# 🤗 pets
-Parameter-Efficient Tuning at Scale with 🤗 Accelerate
+# 🤗 PET
+Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. 
 
 Supported methods:
+
 1. Prefix Tuning
 2. P-Tuning
 3. Prompt Tuning
@@ -38,4 +39,7 @@ Supported methods:
 | BART      | ✅  | ✅  | ✅  |   |
 
 
+## Caveats:
+1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog.
+
 

setup.py

@@ -22,7 +22,7 @@ extras["dev"] = extras["quality"]
 setup(
     name="pets",
     version="0.1.0.dev0",
-    description="Parameter-Efficient Tuning at Scale (PETS)",
+    description="Parameter-Efficient Tuning (PET)",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
     keywords="deep learning",

src/pet/__init__.py

@@ -15,5 +15,7 @@ from .tuners import (
     PromptEncoderReparameterizationType,
     PromptTuningConfig,
     PromptTuningInit,
+    LoRAModel,
+    LoRAConfig,
 )
 from .utils import PETConfig, PETType, PromptLearningConfig, TaskType

src/pet/tuners/__init__.py

@@ -5,3 +5,4 @@
 from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType
 from .prefix_tuning import PrefixEncoder, PrefixTuningConfig
 from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit
+from .lora import LoRAModel, LoRAConfig

src/pet/tuners/lora.py

@@ -1,8 +1,29 @@
 # todo
+from typing import Callable, Optional
 import torch
-from transformers import Conv1D
+from transformers.pytorch_utils import Conv1D
+from dataclasses import dataclass, asdict, field
 
 import loralib as lora
+from loralib import mark_only_lora_as_trainable, lora_state_dict  # flake8: noqa
+
+from ..utils import PETConfig
+
+
+@dataclass
+class LoRAConfig(PETConfig):
+    r: int = field(default=None, metadata={"help": "LoRA attention dimension"})
+    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)"},
+    )
+    target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
+    bias: str = field(default="none", metadata={"help": "Bias type for LoRA. Can be 'none', 'all' or 'lora_only'"})
 
 
 class LoRAModel(torch.nn.Module):
@@ -10,27 +31,26 @@ class LoRAModel(torch.nn.Module):
         super().__init__()
         self.config = config
         self.model = model
+        self.find_and_replace()
+        mark_only_lora_as_trainable(self.model, self.config.bias)
 
     def find_and_replace(self):
         key_list = [key for key, _ in self.model.named_modules()]
         for key in key_list:
-            if any(key.endswith(target_key) for target_key in self.config["target_module_keys"]):
-                parent, target_name, target = self.get_submodules(key)
+            if any(key.endswith(target_key) for target_key in self.config.target_module_keys):
+                parent, target, target_name = self.get_submodules(key)
                 if isinstance(target, torch.nn.Linear):
-                    new_module = lora.Linear(
-                        target.in_features, target.out_features, **self.config["prompt_encoder_config"]
-                    )
-                elif isinstance(target, torch.nn.Conv1d, Conv1D):
-                    new_module = lora.LoRAConv1d(
-                        target.in_channels, target.out_channels, target.kernel_size, bias=target.bias is not None
-                    )
+                    new_module = lora.Linear(target.in_features, target.out_features, **asdict(self.config))
+                elif isinstance(target, Conv1D):
+                    in_features, out_features = target.weight.shape
+                    new_module = lora.MergedLinear(in_features, out_features, **asdict(self.config))
                 self.replace_module(parent, target_name, new_module)
 
     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_name, target
+        return parent, target, target_name
 
     def replace_module(self, parent_module, child_name, new_module, old_module):
         setattr(parent_module, child_name, new_module)