refactor: lora

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

configuration_xlm_roberta.py

@@ -22,7 +22,10 @@ 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,
             load_trained_adapters=False,
             use_flash_attn=True,
             torch_dtype=None,
@@ -47,8 +50,11 @@ 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.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

@@ -9,14 +9,18 @@ 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,
+    XLMRobertaPreTrainedModel,
+)
 
 
 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 +49,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 +68,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 +131,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 +184,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 +196,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,15 +208,14 @@ 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):
+    def __init__(
+        self,
+        config: XLMRobertaFlashConfig,
+        roberta: Optional[XLMRobertaModel] = None,
+        add_pooling_layer=True,
+    ):
         super().__init__(config)
 
         if roberta is None:
@@ -211,10 +223,17 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         else:
             self.roberta = roberta
 
-        self._is_merged = False
-        self._num_adaptions = config.num_loras
-        self._register_lora(self._num_adaptions)
+        self._num_adaptations = len(config.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 = False
         self._task_idx = None
         # By default, we select the first LoRA
@@ -237,13 +256,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,31 +271,33 @@ 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)
+            roberta = XLMRobertaModel.from_pretrained(
+                pretrained_model_name_or_path, *model_args, **kwargs
+            )
             return cls(config, roberta=roberta)
 
-    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
@@ -296,9 +310,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         :param task_idx: 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
+        assert task_idx is None or 0 <= task_idx < self._num_adaptations
         if self._task_idx != task_idx:
             # In this case, we need to update the LoRAs everywhere
             self._task_idx = task_idx
@@ -306,9 +318,9 @@ 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
+    def forward(self, *args, lora_adaptation: Union[None, int] = -1, **kwargs):
+        if lora_adaptation is None or lora_adaptation >= 0:
+            self.current_task = lora_adaptation
         return self.roberta(*args, **kwargs)
 
     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:

modeling_xlm_roberta.py

@@ -452,6 +452,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         convert_to_tensor: bool = False,
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
+        task: Optional[str] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -481,6 +482,12 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                 If set to true, returned vectors will have length 1. In that case, the
                 faster dot-product (util.dot_score) instead of cosine similarity can
                 be used.
+            task(`str`, *optional*, defaults to None):
+                Specifies the task for which the encoding is intended. This
+                controls the use of specialized LoRA adapters that are tuned for specific tasks.
+                If provided, the corresponding LoRA adapter is enabled, enhancing the model's
+                performance for that task. If `None` or not provided, LoRA is disabled, and the
+                model uses its original, general-purpose weights.
             tokenizer_kwargs(`Dict[str, Any]`, *optional*, defaults to {}):
                 Keyword arguments for the tokenizer
         Returns:
@@ -518,6 +525,22 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         if device is not None:
             self.to(device)
 
+        lora_adapter_num = None
+        if self.config.lora_adaptations:
+            if task:
+                if task in self.config.lora_adaptations:
+                    lora_adapter_num = self.config.lora_adaptations.index(task)
+                else:
+                    raise ValueError(
+                        f"Unsupported task '{task}'. "
+                        f"Supported tasks are: {', '.join(self.config.lora_adaptations)}.")
+            else:
+                logger.warning(
+                    f"Task-specific embeddings are disabled. To enable, specify the `task` "
+                    f"argument with one of the supported tasks: {', '.join(self.config.lora_adaptations)}"
+                )
+
+
         permutation = np.argsort([-len(i) for i in sentences])
         inverse_permutation = np.argsort(permutation)
         sentences = [sentences[idx] for idx in permutation]
@@ -547,7 +570,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                 return_tensors='pt',
                 **tokenizer_kwargs,
             ).to(self.device)
-            token_embs = self.forward(**encoded_input)[0]
+            token_embs = self.forward(**encoded_input, lora_adaptation=lora_adapter_num)[0]
 
             # Accumulate in fp32 to avoid overflow
             token_embs = token_embs.float()
@@ -1253,4 +1276,4 @@ class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):
             logits=logits,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-        )
+        )