make-lora-stateless (#25)

- draft (6cc0f517fef666dc2c5c0f0fe7c92935f7b7d7bb)
- poc (eefe43c08d4952b1e2a3c17d6e20090a72a52474)
- refactor: finalize impl (509511d7e5d7b46de5b143e5fc908784eada3f07)
- refactor: modify encode (3eb20d0a7451e244a9c2ee80e7bb762d17a9f9d1)
- fix: 0 is not none (ae40cb93420908fdd3f370e6af3cb0cbe9cbf90d)
- fix: remove prints (acffa62b444442f4933a3023f19d09148de9e2fb)
- refactor: kwargs comprehension (4e13c908bafe0f54becc8c278b3752a37c207c50)
- fix: residual is kept in kwargs (d9d83063f0fcb5bfa22cd4946f23baa1421b9930)
- feat: merge with recent changes (493416fa0cabc0de4deba3c317c58043f9998d49)

block.py

@@ -233,7 +233,7 @@ class Block(nn.Module):
                     is_rms_norm=isinstance(self.norm1, RMSNorm),
                 )
             if not isinstance(self.mlp, nn.Identity):
-                mlp_out = self.mlp(hidden_states)
+                mlp_out = self.mlp(hidden_states, task_type=mixer_kwargs.get('task_type'))
                 if self.return_residual:  # mlp out is actually a pair here
                     mlp_out, hidden_states = mlp_out
                 if not self.fused_dropout_add_ln:

embedding.py

@@ -40,14 +40,15 @@ class XLMRobertaEmbeddings(nn.Module):
         if self.type_vocab_size > 0:
             self.token_type_embeddings = nn.Embedding(type_vocab_size, embed_dim, **factory_kwargs)
 
-    def forward(self, input_ids, position_ids=None, token_type_ids=None):
+    def forward(self, input_ids, position_ids=None, token_type_ids=None, task_type=None):
         """
         input_ids: (batch, seqlen)
         position_ids: (batch, seqlen)
         token_type_ids: (batch, seqlen)
         """
         batch_size, seqlen = input_ids.shape
-        embeddings = self.word_embeddings(input_ids)
+        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
+        embeddings = self.word_embeddings(input_ids, **lora_kwargs)
         if self.max_position_embeddings > 0:
             if position_ids is None:
                 position_ids = create_position_ids_from_input_ids(input_ids, padding_idx=self.word_embeddings.padding_idx).to(input_ids.device)
@@ -57,6 +58,6 @@ class XLMRobertaEmbeddings(nn.Module):
         if self.type_vocab_size > 0:
             if token_type_ids is None:
                 token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=input_ids.device)
-            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+            token_type_embeddings = self.token_type_embeddings(token_type_ids, **lora_kwargs)
             embeddings = embeddings + token_type_embeddings
         return embeddings

mha.py

@@ -450,6 +450,7 @@ class MHA(nn.Module):
 
         if fused_bias_fc and FusedDense is None:
             raise ImportError("fused_dense is not installed")
+
         linear_cls = nn.Linear if not fused_bias_fc else FusedDense
         linear_resid_cls = (
             LinearResidual if not fused_bias_fc else partial(FusedDense, return_residual=True)
@@ -589,6 +590,7 @@ class MHA(nn.Module):
         max_seqlen=None,
         mixer_subset=None,
         inference_params=None,
+        task_type=None,
         **kwargs,
     ):
         """
@@ -643,10 +645,14 @@ class MHA(nn.Module):
         batch, seqlen = x.shape[:2]
         if not self.cross_attn and self.num_heads_kv == self.num_heads:
             assert x_kv is None and mixer_subset is None
+            lora_kwargs = {'task_type': task_type} if task_type is not None else {}
             if not self.return_residual:
-                qkv = self.Wqkv(x)
+                qkv = self.Wqkv(x, **lora_kwargs)
             else:
-                qkv, x = self.Wqkv(x)
+                if lora_kwargs:
+                    lora_kwargs['residual'] = True
+                qkv, x = self.Wqkv(x, **lora_kwargs)
+
             if self.dwconv:
                 qkv = rearrange(
                     self.dwconv_qkv(rearrange(qkv, "b s d -> b d s"))[..., :-2], "b d s -> b s d"
@@ -731,5 +737,7 @@ class MHA(nn.Module):
                     context = self._update_kvcache_attention(q, kv, inference_params)
             else:
                 context = self._apply_rotary_update_kvcache_attention(q, kv, inference_params)
-        out = self.out_proj(rearrange(context, "... h d -> ... (h d)"))
+
+        lora_kwargs.pop('residual', None)
+        out = self.out_proj(rearrange(context, "... h d -> ... (h d)"), **lora_kwargs)
         return out if not self.return_residual else (out, x)

mlp.py

@@ -47,10 +47,11 @@ class Mlp(nn.Module):
         self.activation = activation
         self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs)
 
-    def forward(self, x):
-        y = self.fc1(x)
+    def forward(self, x, task_type=None):
+        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
+        y = self.fc1(x, **lora_kwargs)
         y = self.activation(y)
-        y = self.fc2(y)
+        y = self.fc2(y, **lora_kwargs)
         return y if not self.return_residual else (y, x)
 
 

modeling_lora.py

@@ -9,6 +9,7 @@ import torch
 import torch.nn.utils.parametrize as parametrize
 from torch import nn
 from torch.nn import Parameter
+from torch.nn import functional as F
 from transformers import PretrainedConfig
 
 from .modeling_xlm_roberta import XLMRobertaFlashConfig, XLMRobertaModel, XLMRobertaPreTrainedModel
@@ -88,22 +89,19 @@ class LoRAParametrization(nn.Module):
             torch.ones(self.swap((1, fan_in)), dtype=self.lora_A.dtype),
             persistent=False,
         )
-        self.forward_fn = lambda x: x
-        self.current_task = None
 
     def _dropout(self, A):
         # to mimic the original implementation: A @ dropout(x), we do (A * dropout(ones)) @ x
         return A * self.lora_dropout(self.lora_dropout_mask)
 
-    def lora_forward(self, X):
-        assert self.current_task is not None
+    def lora_forward(self, X, current_task):
         return (
             X
             + torch.matmul(
                 *self.swap(
                     (
-                        self.lora_B[self.current_task],
-                        self.dropout_fn(self.lora_A[self.current_task]),
+                        self.lora_B[current_task],
+                        self.dropout_fn(self.lora_A[current_task]),
                     )
                 )
             ).view(X.shape)
@@ -111,19 +109,7 @@ class LoRAParametrization(nn.Module):
         )
 
     def forward(self, X):
-        return self.forward_fn(X)
-
-    @property
-    def current_task(self):
-        return self._current_task
-
-    @current_task.setter
-    def current_task(self, task: Union[None, int]):
-        self._current_task = task
-        if task is None:
-            self.forward_fn = lambda x: x
-        else:
-            self.forward_fn = self.lora_forward
+        return X
 
     @classmethod
     def from_linear(
@@ -175,6 +161,7 @@ class LoRAParametrization(nn.Module):
         rank: int,
         dropout_p: float,
         alpha: float,
+        adaptation_map: dict,
     ):
         if isinstance(layer, nn.Linear):
             parametrize.register_parametrization(
@@ -188,6 +175,22 @@ class LoRAParametrization(nn.Module):
                     alpha=alpha,
                 ),
             )
+
+            def new_forward(self, input, task_type, residual=False):
+                task_idx = adaptation_map[task_type] if task_type else None
+                if task_idx is not None:
+                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
+                else:
+                    weights = self.weight
+
+                out = F.linear(input, weights, self.bias)
+
+                if residual:
+                    return out, input
+                return out
+
+            layer.forward = new_forward.__get__(layer, layer.__class__)
+
         elif isinstance(layer, nn.Embedding):
             parametrize.register_parametrization(
                 layer,
@@ -201,10 +204,20 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-    @staticmethod
-    def select_task_for_layer(layer: nn.Module, task_idx: Optional[int] = None):
-        if isinstance(layer, LoRAParametrization):
-            layer.current_task = task_idx
+            def new_forward(self, input, task_type):
+                task_idx = adaptation_map[task_type] if task_type else None
+                if task_idx is not None:
+                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
+                else:
+                    weights = self.weight
+
+                out = F.embedding(
+                    input, weights, self.padding_idx, self.max_norm,
+                    self.norm_type, self.scale_grad_by_freq, self.sparse)
+
+                return out
+
+            layer.forward = new_forward.__get__(layer, layer.__class__)
 
 
 class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
@@ -251,9 +264,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             alpha=self._alpha,
         )
         self.main_params_trainable = config.lora_main_params_trainable
-        self._task_idx = None
-        # By default, disable LoRA until it's specified which adapter/task to use
-        self.current_task = None
+
 
     @property
     def main_params_trainable(self):
@@ -307,51 +318,11 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 rank=rank,
                 dropout_p=dropout_p,
                 alpha=alpha,
+                adaptation_map=self._adaptation_map,
             )
         )
 
-    @property
-    def current_task(self):
-        """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_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_name: Which LoRA to use
-        :return:
-        """
-        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
-            self.apply(
-                partial(LoRAParametrization.select_task_for_layer, task_idx=task_idx)
-            )
-
-    def forward(self, *args, task_type: Union[str, None] = None, **kwargs):
-        if task_type:
-            self.current_task = task_type
-        else:
-            input_ids = kwargs["input_ids"]
-            input_text = self.roberta.tokenizer.decode(input_ids[0], skip_special_tokens=True)
-            for task_name, prompt in self._lora_prompts.items():
-                if input_text.startswith(prompt):
-                    self.current_task = task_name
-                    break
-            else:
-                self.current_task = None  # No task-specific adapter is found, just use the general-purpose weights
-
+    def forward(self, *args, **kwargs):
         return self.roberta(*args, **kwargs)
 
     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
@@ -371,33 +342,22 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     def encode(
         self,
         *args,
-        task_type: Union[str, None] = None,
+        task_type: Optional[str] = None,
         **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 `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.
+        task_type(`str`, *optional*, defaults to `None`):
+            Specifies the task for which the encoding is intended. If `task_type` is not provide,
+            all LoRA adapters are disabled, and the model reverts to its original,
+            general-purpose weights.
         """
-        if task_type:
-            self.current_task = task_type
-        else:  # infer the task from the input text
-            input_text = args[0][0] if isinstance(args[0], list) else args[0]  # take only the first sentence
-            for task_name, prompt in self._lora_prompts.items():
-                if input_text.startswith(prompt):
-                    self.current_task = task_name
-                    break
-            else:
-                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 = None  # No task-specific adapter is found, just use the general-purpose weights
+        if task_type and task_type not in self._lora_adaptations:
+            raise ValueError(
+                f"Unsupported task '{task_type}'. "
+                f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
+                f"Alternatively, don't pass the `task_type` argument to disable LoRA."
+            )
 
-        return self.roberta.encode(*args, **kwargs)
+        return self.roberta.encode(*args, task_type=task_type, **kwargs)

modeling_xlm_roberta.py

@@ -204,7 +204,7 @@ class XLMRobertaEncoder(nn.Module):
     def gradient_checkpointing(self, value):
         self._grad_checkpointing = value
 
-    def forward(self, hidden_states, key_padding_mask=None, subset_mask=None):
+    def forward(self, hidden_states, key_padding_mask=None, subset_mask=None, task_type=None):
         """If subset_mask is not None, we only want output for the subset of the sequence.
         This means that we only compute the last layer output for these tokens.
         subset_mask: (batch, seqlen), dtype=torch.bool
@@ -215,6 +215,7 @@ class XLMRobertaEncoder(nn.Module):
                 if key_padding_mask is not None
                 else None
             )
+            mixer_kwargs['task_type'] = task_type
             for layer in self.layers:
                 if self._grad_checkpointing:
                     hidden_states = torch.utils.checkpoint.checkpoint(
@@ -232,7 +233,7 @@ class XLMRobertaEncoder(nn.Module):
             hidden_states, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
                 hidden_states, key_padding_mask
             )
-            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch}
+            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch, "task_type": task_type}
             if subset_mask is None:
                 for layer in self.layers:
                     if self._grad_checkpointing:
@@ -309,11 +310,13 @@ class XLMRobertaPooler(nn.Module):
         self.dense = linear_cls(config.hidden_size, config.hidden_size)
         self.activation = nn.Tanh()
 
-    def forward(self, hidden_states, pool=True):
+    def forward(self, hidden_states, pool=True, task_type=None):
         # We "pool" the model by simply taking the hidden state corresponding
         # to the first token.
+        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
+
         first_token_tensor = hidden_states[:, 0] if pool else hidden_states
-        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.dense(first_token_tensor, **lora_kwargs)
         pooled_output = self.activation(pooled_output)
         return pooled_output
 
@@ -454,6 +457,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
         truncate_dim: Optional[int] = None,
+        task_type: Optional[str] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -538,14 +542,14 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             )
         else:
             range_iter = range(0, len(sentences), batch_size)
-
+        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
         for i in range_iter:
             encoded_input = self.tokenizer(
                 sentences[i : i + batch_size],
                 return_tensors='pt',
                 **tokenizer_kwargs,
             ).to(self.device)
-            token_embs = self.forward(**encoded_input)[0]
+            token_embs = self.forward(**encoded_input, **lora_kwargs)[0]
 
             # Accumulate in fp32 to avoid overflow
             token_embs = token_embs.float()
@@ -633,7 +637,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         layer output for these tokens.
         masked_tokens_mask: (batch, seqlen), dtype=torch.bool
         """
-
+        task_type = kwargs.pop('task_type', None)
         if kwargs:
             for key, value in kwargs.items():
                 if value is not None:
@@ -647,7 +651,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         )
 
         hidden_states = self.embeddings(
-            input_ids, position_ids=position_ids, token_type_ids=token_type_ids
+            input_ids, position_ids=position_ids, token_type_ids=token_type_ids, task_type=task_type
         )
         # TD [2022-12:18]: Don't need to force residual in fp32
         # BERT puts embedding LayerNorm before embedding dropout.
@@ -671,12 +675,12 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             subset_mask = None
 
         sequence_output = self.encoder(
-            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask
+            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask, task_type=task_type
         )
 
         if masked_tokens_mask is None:
             pooled_output = (
-                self.pooler(sequence_output) if self.pooler is not None else None
+                self.pooler(sequence_output, task_type=task_type) if self.pooler is not None else None
             )
         else:
             # TD [2022-03-01]: the indexing here is very tricky.
@@ -690,7 +694,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                 pool_input = sequence_output[first_col_mask[subset_mask]]
                 sequence_output = sequence_output[masked_tokens_mask[subset_mask]]
             pooled_output = (
-                self.pooler(pool_input, pool=False) if self.pooler is not None else None
+                self.pooler(pool_input, pool=False, task_type=task_type) if self.pooler is not None else None
             )
 
         if not return_dict: