feat: merge with recent changes

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

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, task=mixer_kwargs.get('task'))
+                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:

configuration_xlm_roberta.py

@@ -23,6 +23,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             use_cache=True,
             classifier_dropout=None,
             lora_adaptations=None,
+            lora_prompts=None,
             lora_rank=4,
             lora_dropout_p=0.0,
             lora_alpha=1,
@@ -55,6 +56,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.classifier_dropout = classifier_dropout
         self.load_trained_adapters = load_trained_adapters
         self.lora_adaptations = lora_adaptations
+        self.lora_prompts = lora_prompts
         self.lora_rank = lora_rank
         self.lora_dropout_p = lora_dropout_p
         self.lora_alpha = lora_alpha

embedding.py

@@ -40,14 +40,14 @@ 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, task=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
-        lora_kwargs = {'task': task} if task is not None else {}
+        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:

mha.py

@@ -590,7 +590,7 @@ class MHA(nn.Module):
         max_seqlen=None,
         mixer_subset=None,
         inference_params=None,
-        task=None,
+        task_type=None,
         **kwargs,
     ):
         """
@@ -645,7 +645,7 @@ 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': task} if task is not None else {}
+            lora_kwargs = {'task_type': task_type} if task_type is not None else {}
             if not self.return_residual:
                 qkv = self.Wqkv(x, **lora_kwargs)
             else:

mlp.py

@@ -47,8 +47,8 @@ class Mlp(nn.Module):
         self.activation = activation
         self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs)
 
-    def forward(self, x, task):
-        lora_kwargs = {'task': task} if task is not None else {}
+    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, **lora_kwargs)

modeling_lora.py

@@ -15,9 +15,6 @@ from transformers import PretrainedConfig
 from .modeling_xlm_roberta import XLMRobertaFlashConfig, XLMRobertaModel, XLMRobertaPreTrainedModel
 
 
-LORA_NO_UPDATE = '__lora_no_update__'
-
-
 def initialized_weights(
     shape: Tuple[int], num_adaptations: int, init: str = "kaiming"
 ) -> torch.Tensor:
@@ -179,8 +176,8 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-            def new_forward(self, input, task, residual=False):
-                task_idx = adaptation_map[task] if task else None
+            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:
@@ -207,8 +204,8 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-            def new_forward(self, input, task):
-                task_idx = adaptation_map[task] if task else None
+            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:
@@ -244,6 +241,16 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             raise ValueError(
                 f'`lora_adaptations` must be a list and contain at least one element'
             )
+        self._lora_prompts = config.lora_prompts
+        if (
+            not isinstance(self._lora_prompts, dict)
+            or len(self._lora_prompts) != len(self._lora_adaptations)
+            or not all([v in self._lora_adaptations for v in self._lora_prompts.keys()])
+        ):
+            raise ValueError(
+                f'`lora_prompts` must be a dict and contain the same number of elements '
+                f'as `lora_adaptations` with all keys in `lora_prompts` present in `lora_adaptations`.'
+        )
         self._adaptation_map = {
             name: idx for idx, name in enumerate(self._lora_adaptations)
         }
@@ -335,25 +342,22 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     def encode(
         self,
         *args,
-        task: Optional[str] = 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 `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.
+        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 and task not in self._lora_adaptations:
+        if task_type and task_type not in self._lora_adaptations:
             raise ValueError(
-                f"Unsupported task '{task}'. "
+                f"Unsupported task '{task_type}'. "
                 f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
-                f"Alternatively, don't pass the `task` argument to disable LoRA."
+                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

@@ -21,7 +21,7 @@ import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from einops import rearrange
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, AutoTokenizer
 from transformers.modeling_utils import PreTrainedModel
 from transformers.modeling_outputs import MaskedLMOutput,SequenceClassifierOutput
 from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaLMHead
@@ -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, task=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,7 +215,7 @@ class XLMRobertaEncoder(nn.Module):
                 if key_padding_mask is not None
                 else None
             )
-            mixer_kwargs['task'] = task
+            mixer_kwargs['task_type'] = task_type
             for layer in self.layers:
                 if self._grad_checkpointing:
                     hidden_states = torch.utils.checkpoint.checkpoint(
@@ -233,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, "task": task}
+            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:
@@ -310,10 +310,10 @@ 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, task=None):
+    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': task} if task is not None else {}
+        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, **lora_kwargs)
@@ -443,7 +443,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         self.pooler = XLMRobertaPooler(config) if add_pooling_layer else None
 
         self.apply(partial(_init_weights, initializer_range=config.initializer_range))
-
+        self.tokenizer = AutoTokenizer.from_pretrained(self.name_or_path, trust_remote_code=True)
 
     @torch.inference_mode()
     def encode(
@@ -457,7 +457,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
         truncate_dim: Optional[int] = None,
-        task: Optional[str] = None,
+        task_type: Optional[str] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -496,12 +496,6 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             If convert_to_tensor, a stacked tensor is returned.
             If convert_to_numpy, a numpy matrix is returned.
         """
-        from transformers import AutoTokenizer
-
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            self.name_or_path, trust_remote_code=True
-        )
-
         is_training = self.training
         self.eval()
 
@@ -548,7 +542,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             )
         else:
             range_iter = range(0, len(sentences), batch_size)
-        lora_kwargs = {'task': task} if task is not None else {}
+        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],
@@ -643,7 +637,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         layer output for these tokens.
         masked_tokens_mask: (batch, seqlen), dtype=torch.bool
         """
-        task = kwargs.pop('task', None)
+        task_type = kwargs.pop('task_type', None)
         if kwargs:
             for key, value in kwargs.items():
                 if value is not None:
@@ -657,7 +651,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         )
 
         hidden_states = self.embeddings(
-            input_ids, position_ids=position_ids, token_type_ids=token_type_ids, task=task
+            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.
@@ -681,12 +675,12 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             subset_mask = None
 
         sequence_output = self.encoder(
-            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask, task=task
+            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, task=task) 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.
@@ -700,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, task=task) 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:
@@ -1282,4 +1276,4 @@ class XLMRobertaForSequenceClassification(XLMRobertaPreTrainedModel):
             logits=logits,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-        )
+        )

rotary.py

@@ -6,11 +6,13 @@ from typing import Optional, Tuple, Union
 
 import torch
 from einops import rearrange, repeat
-try:
-    from flash_attn.ops.triton.rotary import apply_rotary
-except ImportError:
-    def apply_rotary(*args, **kwargs):
-        raise RuntimeError('RoPE requires flash-attention to be installed')
+
+if torch.cuda.is_available():
+    try:
+        from flash_attn.ops.triton.rotary import apply_rotary
+    except ImportError:
+        def apply_rotary(*args, **kwargs):
+            raise RuntimeError('RoPE requires flash-attention to be installed')
 
 
 def rotate_half(x, interleaved=False):
@@ -29,6 +31,10 @@ def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
     """
     ro_dim = cos.shape[-1] * 2
     assert ro_dim <= x.shape[-1]
+    cos, sin = (
+        cos[:x.shape[1]],
+        sin[:x.shape[1]],
+    )
     cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
     sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
     return torch.cat(
@@ -60,6 +66,7 @@ class ApplyRotaryEmb(torch.autograd.Function):
             interleaved=interleaved,
             inplace=inplace,
         )
+
         if isinstance(seqlen_offsets, int):
             ctx.save_for_backward(cos, sin, cu_seqlens)  # Can't save int with save_for_backward
             ctx.seqlen_offsets = seqlen_offsets
@@ -82,6 +89,7 @@ class ApplyRotaryEmb(torch.autograd.Function):
         # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works.
         if not ctx.interleaved and not ctx.inplace:
             do = do.clone()
+
         dx = apply_rotary(
             do,
             cos,
@@ -150,21 +158,37 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
         # batch, seqlen, three, nheads, headdim = qkv.shape
         assert qkv.shape[-3] == 3
         if cos_k is None and sin_k is None and qkv.is_contiguous():
-            # Call 1 kernel instead of 2 kernels
-            # We need qkv to be contiguous so that when we reshape to combine (3, nheads)
-            # dimensions, we get the same tensor
-            qk = rearrange(qkv[..., :2, :, :], "... t h d -> ... (t h) d")
-            # qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim)
-            apply_rotary(
-                qk,
-                cos,
-                sin,
-                seqlen_offsets=seqlen_offsets,
-                interleaved=interleaved,
-                inplace=True,
-                cu_seqlens=cu_seqlens,
-                max_seqlen=max_seqlen,
-            )
+
+            if torch.cuda.is_available():
+                # Call 1 kernel instead of 2 kernels
+                # We need qkv to be contiguous so that when we reshape to combine (3, nheads)
+                # dimensions, we get the same tensor
+                qk = rearrange(qkv[..., :2, :, :], "... t h d -> ... (t h) d")
+                # qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim)
+                apply_rotary(
+                    qk,
+                    cos,
+                    sin,
+                    seqlen_offsets=seqlen_offsets,
+                    interleaved=interleaved,
+                    inplace=True,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                )
+            else:
+                q_rot = apply_rotary_emb_torch(
+                    qkv[:, :, 0],
+                    cos,
+                    sin,
+                    interleaved=interleaved,
+                )
+                k_rot = apply_rotary_emb_torch(
+                    qkv[:, :, 1],
+                    cos,
+                    sin,
+                    interleaved=interleaved,
+                )
+                qkv = torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
         else:
             cos_k = cos if cos_k is None else cos_k
             sin_k = sin if sin_k is None else sin_k
@@ -228,7 +252,6 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
             sin_k = sin if sin_k is None else sin_k
             dq, dk = dqkv[..., 0, :, :], dqkv[..., 1, :, :]
             apply_rotary(
-
                 dq,
                 cos,
                 sin,