Refactor LoRA (#8)

- refactor: lora (e0ea168f256f54e7728033ff8afbf9bb71c617cc)
- refactor: remove pooling layer stuff (c6a5a4d6aa2d39e1b0691e4f48869aa8d7e34b09)
- refactor: restructure the class (a2b7c8644033cc4318c6caa6730836023776faa9)
- refactor: disable lora by default (5418705c2b50051908c38d6df1055df4f21274a2)
- refactor: set task in lora class rather than xlm roberta (851aaca7b1e7f9dbffaacd3b070231e0f94401cb)
- refactor: stuff (370394630d973381185d21d153a5e46d3b9fc6da)

configuration_xlm_roberta.py

@@ -22,7 +22,11 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             position_embedding_type="absolute",
             use_cache=True,
             classifier_dropout=None,
-            num_loras=1,
+            lora_adaptations=None,
+            lora_rank=4,
+            lora_dropout_p=0.0,
+            lora_alpha=1,
+            lora_main_params_trainable=False,
             load_trained_adapters=False,
             use_flash_attn=True,
             torch_dtype=None,
@@ -47,8 +51,12 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.position_embedding_type = position_embedding_type
         self.use_cache = use_cache
         self.classifier_dropout = classifier_dropout
-        self.num_loras = num_loras
         self.load_trained_adapters = load_trained_adapters
+        self.lora_adaptations = lora_adaptations
+        self.lora_rank = lora_rank
+        self.lora_dropout_p = lora_dropout_p
+        self.lora_alpha = lora_alpha
+        self.lora_main_params_trainable = lora_main_params_trainable
         self.use_flash_attn = use_flash_attn
         self.emb_pooler = emb_pooler
         if torch_dtype and hasattr(torch, torch_dtype) and type(getattr(torch, torch_dtype)) is torch.dtype:

modeling_lora.py

@@ -1,22 +1,27 @@
 import math
 import os
+import warnings
 from functools import partial
-from typing import Iterator, Optional, Tuple, Union
+from typing import Iterator, List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 import torch.nn.utils.parametrize as parametrize
 from torch import nn
 from torch.nn import Parameter
 from transformers import PretrainedConfig
 
-from .modeling_xlm_roberta import XLMRobertaModel, XLMRobertaPreTrainedModel, XLMRobertaFlashConfig
+from .modeling_xlm_roberta import XLMRobertaFlashConfig, XLMRobertaModel
+
+
+LORA_NO_UPDATE = '__lora_no_update__'
 
 
 def initialized_weights(
-    shape: Tuple[int], num_adaptions: int, init: str = "kaiming"
+    shape: Tuple[int], num_adaptations: int, init: str = "kaiming"
 ) -> torch.Tensor:
     weight_data = []
-    for _ in range(num_adaptions):
+    for _ in range(num_adaptations):
         new_adaption = torch.zeros(shape)
         if init == "kaiming":
             nn.init.kaiming_uniform_(new_adaption, a=math.sqrt(5))
@@ -45,15 +50,16 @@ class LoRAParametrization(nn.Module):
     WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
     SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
     """
+
     def __init__(
         self,
         fan_in: int,
         fan_out: int,
         layer_type: str = "linear",
-        num_adaptions: int = 1,
+        num_adaptations: int = 1,
         rank: int = 4,
-        lora_dropout_p: float = 0.0,
-        lora_alpha: float = 1,
+        dropout_p: float = 0.0,
+        alpha: float = 1,
     ):
         super().__init__()
         # if weight is stored as (fan_out, fan_in), the memory layout of A & B follows (W + BA)x
@@ -63,25 +69,23 @@ class LoRAParametrization(nn.Module):
 
         if layer_type == "linear":
             self.lora_A = nn.Parameter(
-                initialized_weights((rank, fan_in), num_adaptions, init="kaiming")
+                initialized_weights((rank, fan_in), num_adaptations, init="kaiming")
             )
-            self.lora_B = nn.Parameter(torch.zeros((num_adaptions, fan_out, rank)))
+            self.lora_B = nn.Parameter(torch.zeros((num_adaptations, fan_out, rank)))
         elif layer_type == "embedding":
-            self.lora_A = nn.Parameter(torch.zeros((num_adaptions, fan_in, rank)))
+            self.lora_A = nn.Parameter(torch.zeros((num_adaptations, fan_in, rank)))
             self.lora_B = nn.Parameter(
                 initialized_weights(
-                    (rank, fan_out), num_adaptions=num_adaptions, init="normal"
+                    (rank, fan_out), num_adaptations=num_adaptations, init="normal"
                 )
             )
         else:
             raise NotImplementedError
 
-        self.lora_alpha, self.rank = lora_alpha, rank
-        self.scaling = lora_alpha / rank
-        self.lora_dropout = (
-            nn.Dropout(p=lora_dropout_p) if lora_dropout_p > 0 else lambda x: x
-        )
-        self.dropout_fn = self._dropout if lora_dropout_p > 0 else lambda x: x
+        self.lora_alpha, self.rank = alpha, rank
+        self.scaling = alpha / rank
+        self.lora_dropout = nn.Dropout(p=dropout_p) if dropout_p > 0 else lambda x: x
+        self.dropout_fn = self._dropout if dropout_p > 0 else lambda x: x
         self.register_buffer(
             "lora_dropout_mask",
             torch.ones(self.swap((1, fan_in)), dtype=self.lora_A.dtype),
@@ -128,42 +132,52 @@ class LoRAParametrization(nn.Module):
     def from_linear(
         cls,
         layer: nn.Module,
-        num_adaptions: int = 1,
-        rank: int = 4,
-        lora_dropout_p: float = 0.0,
-        lora_alpha: int = 1,
+        num_adaptations: int,
+        rank: int,
+        dropout_p: float,
+        alpha: float,
     ):
         assert isinstance(layer, nn.Linear)
         fan_out, fan_in = layer.weight.shape
         return cls(
             fan_in,
             fan_out,
-            num_adaptions=num_adaptions,
+            num_adaptations=num_adaptations,
             layer_type="linear",
             rank=rank,
-            lora_dropout_p=lora_dropout_p,
-            lora_alpha=lora_alpha,
+            dropout_p=dropout_p,
+            alpha=alpha,
         )
 
     @classmethod
     def from_embedding(
-        cls, layer, num_adaptions=1, rank=4, lora_dropout_p=0.0, lora_alpha=1
+        cls,
+        layer: nn.Module,
+        num_adaptations: int,
+        rank: int,
+        dropout_p: float,
+        alpha: float,
     ):
         assert isinstance(layer, nn.Embedding)
         fan_in, fan_out = layer.weight.shape
         return cls(
             fan_in,
             fan_out,
-            num_adaptions=num_adaptions,
+            num_adaptations=num_adaptations,
             layer_type="embedding",
             rank=rank,
-            lora_dropout_p=lora_dropout_p,
-            lora_alpha=lora_alpha,
+            dropout_p=dropout_p,
+            alpha=alpha,
         )
 
     @classmethod
     def add_to_layer(
-        cls, layer, num_adaptions=1, rank=4, lora_dropout_p=0.0, lora_alpha=1
+        cls,
+        layer: nn.Module,
+        num_adaptations: int,
+        rank: int,
+        dropout_p: float,
+        alpha: float,
     ):
         if isinstance(layer, nn.Linear):
             parametrize.register_parametrization(
@@ -171,10 +185,10 @@ class LoRAParametrization(nn.Module):
                 "weight",
                 cls.from_linear(
                     layer,
-                    num_adaptions=num_adaptions,
+                    num_adaptations=num_adaptations,
                     rank=rank,
-                    lora_dropout_p=lora_dropout_p,
-                    lora_alpha=lora_alpha,
+                    dropout_p=dropout_p,
+                    alpha=alpha,
                 ),
             )
         elif isinstance(layer, nn.Embedding):
@@ -183,10 +197,10 @@ class LoRAParametrization(nn.Module):
                 "weight",
                 cls.from_embedding(
                     layer,
-                    num_adaptions=num_adaptions,
+                    num_adaptations=num_adaptations,
                     rank=rank,
-                    lora_dropout_p=lora_dropout_p,
-                    lora_alpha=lora_alpha,
+                    dropout_p=dropout_p,
+                    alpha=alpha,
                 ),
             )
 
@@ -195,30 +209,39 @@ class LoRAParametrization(nn.Module):
         if isinstance(layer, LoRAParametrization):
             layer.current_task = task_idx
 
-    @staticmethod
-    def merge_lora_into_layer(layer: nn.Module):
-        if hasattr(layer, "parametrizations"):
-            for attr_name in layer.parametrizations.keys():
-                parametrize.remove_parametrizations(layer, attr_name, leave_parametrized=True)
-
 
-class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
-    def __init__(self, config: XLMRobertaFlashConfig, roberta: Optional[XLMRobertaModel] = None, add_pooling_layer=True):
+class XLMRobertaLoRA(XLMRobertaModel):
+    def __init__(
+        self,
+        config: XLMRobertaFlashConfig,
+    ):
         super().__init__(config)
 
-        if roberta is None:
-            self.roberta = XLMRobertaModel(config, add_pooling_layer=add_pooling_layer)
-        else:
-            self.roberta = roberta
-
-        self._is_merged = False
-        self._num_adaptions = config.num_loras
-        self._register_lora(self._num_adaptions)
-
-        self.main_params_trainable = False
+        self._lora_adaptations = config.lora_adaptations
+        if (
+            not isinstance(self._lora_adaptations, list)
+            or len(self._lora_adaptations) < 1
+        ):
+            raise ValueError(
+                f'`lora_adaptations` must be a list and contain at least one element'
+            )
+        self._adaptation_map = {
+            name: idx for idx, name in enumerate(self._lora_adaptations)
+        }
+        self._rank = config.lora_rank
+        self._dropout_p = config.lora_dropout_p
+        self._alpha = config.lora_alpha
+
+        self._register_lora(
+            num_adaptations=self._num_adaptations,
+            rank=self._rank,
+            dropout_p=self._dropout_p,
+            alpha=self._alpha,
+        )
+        self.main_params_trainable = config.lora_main_params_trainable
         self._task_idx = None
-        # By default, we select the first LoRA
-        self.current_task = 0
+        # By default, disable LoRA until it's specified which adapter/task to use
+        self.current_task = None
 
     @property
     def main_params_trainable(self):
@@ -237,13 +260,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             if "lora" not in name:
                 param.requires_grad_(val)
 
-    def merge_lora(self):
-        """Merges currently selected LoRA into main weights."""
-        if self._is_merged:
-            raise Exception('LoRA has already been merged, cannot merge again')
-        self._is_merged = True
-        self.apply(LoRAParametrization.merge_lora_into_layer)
-
     @classmethod
     def from_pretrained(
         cls,
@@ -259,46 +275,52 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         use_safetensors: bool = None,
         **kwargs,
     ):
-        config = XLMRobertaFlashConfig.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        config = XLMRobertaFlashConfig.from_pretrained(
+            pretrained_model_name_or_path, *model_args, **kwargs
+        )
+
         if config.load_trained_adapters:
             return super().from_pretrained(
-                pretrained_model_name_or_path,
-                *model_args,
-                **kwargs
+                pretrained_model_name_or_path, *model_args, **kwargs
             )
         else:
-            roberta = XLMRobertaModel.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-            return cls(config, roberta=roberta)
+            dtype = config.torch_dtype if config.torch_dtype else torch.bfloat16
+            torch.set_default_dtype(dtype)
+            return cls(config)
 
-    def _register_lora(self, num_adaptions=1, rank=4, lora_dropout_p=0.0, lora_alpha=1):
+    def _register_lora(self, num_adaptations, rank, dropout_p, alpha):
         self.apply(
             partial(
                 LoRAParametrization.add_to_layer,
-                num_adaptions=num_adaptions,
+                num_adaptations=num_adaptations,
                 rank=rank,
-                lora_dropout_p=lora_dropout_p,
-                lora_alpha=lora_alpha,
+                dropout_p=dropout_p,
+                alpha=alpha,
             )
         )
 
     @property
     def current_task(self):
-        """ Which LoRA is currently selected
+        """Which LoRA is currently selected
         :return: Integer or None (when LoRA is disabled)
         """
         return self._task_idx
 
     @current_task.setter
-    def current_task(self, task_idx: Union[None, int]):
+    def current_task(self, task_name: Union[None, str]):
         """Set the LoRA that is to be used.
         The LoRA is specified by `task_idx`, which may be an integer >= 0,
         indexing the available LoRAs. If it is None, no LoRA is used.
-        :param task_idx: Which LoRA to use
+        :param task_name: Which LoRA to use
         :return:
         """
-        if self._is_merged:
-            raise Exception('LoRA has been merged, cannot select new task')
-        assert task_idx is None or 0 <= task_idx < self._num_adaptions
+        if task_name and task_name not in self._lora_adaptations:
+            raise ValueError(
+                f"Unsupported task '{task_name}'. "
+                f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
+                f"Alternatively, set `task` to `None` if you want to disable LoRA."
+            )
+        task_idx = self._adaptation_map[task_name] if task_name else None
         if self._task_idx != task_idx:
             # In this case, we need to update the LoRAs everywhere
             self._task_idx = task_idx
@@ -306,10 +328,10 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 partial(LoRAParametrization.select_task_for_layer, task_idx=task_idx)
             )
 
-    def forward(self, *args, current_task: Union[None, int] = -1, **kwargs):
-        if current_task is None or current_task >= 0:
-            self.current_task = current_task
-        return self.roberta(*args, **kwargs)
+    def forward(self, *args, task: Union[str, None] = LORA_NO_UPDATE, **kwargs):
+        if task != LORA_NO_UPDATE:
+            self.current_task = task
+        return super().forward(*args, **kwargs)
 
     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
         for _, param in self.named_parameters(recurse=recurse):
@@ -323,3 +345,32 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         ):
             if "lora" in name or self.main_params_trainable:
                 yield name, param
+
+    @torch.inference_mode()
+    def encode(
+        self,
+        *args,
+        task: Union[str, None] = LORA_NO_UPDATE,
+        **kwargs,
+    ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
+        """
+        Computes sentence embeddings
+
+        task(`str`, *optional*, defaults to `LORA_NO_UPDATE`):
+            Specifies the task for which the encoding is intended. This parameter controls the
+            use of specialized LoRA adapters that are tuned for specific tasks. If `task` is set
+            to `LORA_NO_UPDATE`, there will be no update to the current task, retaining the
+            existing adapter configuration. If `task` is explicitly set to `None`, all LoRA
+            adapters are disabled, and the model reverts to its original, general-purpose weights.
+            If `task` is set to a specific LoRA adaptation, that adaptation is activated.
+        """
+        if task != LORA_NO_UPDATE:
+            if not task:
+                warnings.warn(
+                    f"Task-specific embeddings are disabled. To enable, specify the `task` "
+                    f"argument with one of the supported tasks: {', '.join(self.config.lora_adaptations)}",
+                    category=UserWarning,
+                )
+            self.current_task = task
+
+        return super().encode(*args, **kwargs)

modeling_xlm_roberta.py

@@ -1253,4 +1253,4 @@ class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):
             logits=logits,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-        )
+        )