lora-instructions

configuration_xlm_roberta.py

@@ -5,6 +5,9 @@ from transformers import PretrainedConfig
 
 
 class XLMRobertaFlashConfig(PretrainedConfig):
+
+    model_type = "xlm-roberta"
+
     def __init__(
         self,
         vocab_size: int = 250002,
@@ -25,9 +28,10 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         position_embedding_type: str = "rotary",
         rotary_emb_base: float = 10000.0,
         use_cache: bool = True,
+        use_reentrant: bool = False,
         classifier_dropout: Optional[float] = None,
         lora_adaptations: Optional[List[str]] = None,
-        lora_prompts: Optional[Dict[str, str]] = None,
+        task_instructions: Optional[Dict[str, str]] = None,
         lora_rank: int = 4,
         lora_dropout_p: float = 0.0,
         lora_alpha: int = 1,
@@ -62,6 +66,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             position_embedding_type (str): Type of position embeddings. Options are 'absolute', 'alibi', or 'rotary'.
             rotary_emb_base (float): Base for rotary embeddings.
             use_cache (bool): Whether or not the model should return the last key/values attentions (not used by all models).
+            use_reentrant (bool): Whether or not the model should enable the 'use_reentrant' flag in gradient checkpointing.
             classifier_dropout (Optional[float]): The dropout ratio for the classification head.
             lora_adaptations (Optional[List[str]]): LoRA adaptations configuration.
             lora_prompts (Optional[Dict[str, str]]): LoRA prompts configuration.
@@ -100,10 +105,11 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.position_embedding_type = position_embedding_type
         self.rotary_emb_base = rotary_emb_base
         self.use_cache = use_cache
+        self.use_reentrant = use_reentrant
         self.classifier_dropout = classifier_dropout
         self.load_trained_adapters = load_trained_adapters
         self.lora_adaptations = lora_adaptations
-        self.lora_prompts = lora_prompts
+        self.task_instructions = task_instructions
         self.lora_rank = lora_rank
         self.lora_dropout_p = lora_dropout_p
         self.lora_alpha = lora_alpha

mha.py

@@ -463,6 +463,7 @@ class MHA(nn.Module):
                 scale_base=rotary_emb_scale_base,
                 interleaved=rotary_emb_interleaved,
                 device=device,
+                use_flash_attn=use_flash_attn,
             )
 
         if fused_bias_fc and FusedDense is None:

modeling_lora.py

@@ -11,6 +11,7 @@ from torch.nn import Parameter
 from torch.nn import functional as F
 from transformers import PretrainedConfig
 
+from .rotary import RotaryEmbedding
 from .modeling_xlm_roberta import (XLMRobertaFlashConfig, XLMRobertaModel,
                                    XLMRobertaPreTrainedModel)
 
@@ -164,7 +165,6 @@ class LoRAParametrization(nn.Module):
     ):
         """
         Registering LoRA adapters to all embedding and linear layers.
-
         Additionally, we implement a custom forward function for LoRA parametrization.
         This function modifies the layer's forward pass to optionally use task-specific
         parameters. When a `task_id` is provided, it employs a LoRA parametrization
@@ -241,6 +241,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     """
     A wrapper class around the Jina XLM-RoBERTa model that integrates LoRA (Low-Rank Adaptation) adapters.
     """
+
     def __init__(
         self, config: XLMRobertaFlashConfig, roberta: Optional[XLMRobertaModel] = None
     ):
@@ -258,15 +259,17 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             raise ValueError(
                 f"`lora_adaptations` must be a list and contain at least one element"
             )
-        self._lora_prompts = config.lora_prompts
+        self._task_instructions = config.task_instructions
         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()])
+            not isinstance(self._task_instructions, dict)
+            or len(self._task_instructions) != len(self._lora_adaptations)
+            or not all(
+                [v in self._lora_adaptations for v in self._task_instructions.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`."
+                f"`task_instructions` must be a dict and contain the same number of elements "
+                f"as `lora_adaptations` with all keys in `task_instructions` present in `lora_adaptations`."
             )
         self._adaptation_map = {
             name: idx for idx, name in enumerate(self._lora_adaptations)
@@ -322,16 +325,13 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         use_safetensors: bool = None,
         **kwargs,
     ):
-        config = XLMRobertaFlashConfig.from_pretrained(
-            pretrained_model_name_or_path, *model_args, **kwargs
-        )
-        if config.load_trained_adapters: # checkpoint already contains LoRA adapters
+        if config.load_trained_adapters:  # checkpoint already contains LoRA adapters
             return super().from_pretrained(
-                pretrained_model_name_or_path, *model_args, **kwargs
+                pretrained_model_name_or_path, *model_args, use_flash_attn=config.use_flash_attn, **kwargs
             )
-        else: # initializing new adapters
+        else:  # initializing new adapters
             roberta = XLMRobertaModel.from_pretrained(
-                pretrained_model_name_or_path, *model_args, **kwargs
+                pretrained_model_name_or_path, *model_args, use_flash_attn=config.use_flash_attn, **kwargs
             )
             return cls(config, roberta=roberta)
 
@@ -372,7 +372,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
         Computes sentence embeddings.
-
         sentences(`str` or `List[str]`):
             Sentence or sentences to be encoded
         task_type(`str`, *optional*, defaults to `None`):
@@ -393,6 +392,10 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             adapter_mask = torch.full(
                 (num_examples,), task_id, dtype=torch.int32, device=self.device
             )
+            if isinstance(sentences, str):
+                sentences = self._task_instructions[task_type] + sentences
+            else:
+                sentences = [self._task_instructions[task_type] + sentence for sentence in sentences]
         return self.roberta.encode(
             sentences, *args, adapter_mask=adapter_mask, **kwargs
         )

modeling_xlm_roberta.py

@@ -30,6 +30,7 @@ from transformers.models.bert.modeling_bert import (
 from transformers.models.xlm_roberta.modeling_xlm_roberta import \
     XLMRobertaLMHead
 
+from .rotary import RotaryEmbedding
 from .block import Block
 from .configuration_xlm_roberta import XLMRobertaFlashConfig
 from .embedding import XLMRobertaEmbeddings
@@ -63,9 +64,7 @@ logger = logging.getLogger(__name__)
 
 
 def get_use_flash_attn(config: XLMRobertaFlashConfig):
-    if not getattr(config, "use_flash_attn", False):
-        return False
-    if not torch.cuda.is_available():
+    if not getattr(config, "use_flash_attn", False) or not torch.cuda.is_available():
         return False
     if importlib.util.find_spec("flash_attn") is None:
         logger.warning(
@@ -181,6 +180,7 @@ class XLMRobertaEncoder(nn.Module):
     def __init__(self, config: XLMRobertaFlashConfig):
         super().__init__()
         self.use_flash_attn = get_use_flash_attn(config)
+        self.use_reentrant = config.use_reentrant
         self.layers = nn.ModuleList(
             [create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
         )
@@ -210,7 +210,7 @@ class XLMRobertaEncoder(nn.Module):
                     hidden_states = torch.utils.checkpoint.checkpoint(
                         layer,
                         hidden_states,
-                        use_reentrant=False,
+                        use_reentrant=self.use_reentrant,
                         mixer_kwargs=mixer_kwargs,
                     )
                 else:
@@ -234,7 +234,7 @@ class XLMRobertaEncoder(nn.Module):
                         hidden_states = torch.utils.checkpoint.checkpoint(
                             layer,
                             hidden_states,
-                            use_reentrant=False,
+                            use_reentrant=self.use_reentrant,
                             mixer_kwargs=mixer_kwargs,
                         )
                     else:
@@ -246,7 +246,7 @@ class XLMRobertaEncoder(nn.Module):
                         hidden_states = torch.utils.checkpoint.checkpoint(
                             layer,
                             hidden_states,
-                            use_reentrant=False,
+                            use_reentrant=self.use_reentrant,
                             mixer_kwargs=mixer_kwargs,
                         )
                     else:
@@ -284,7 +284,7 @@ class XLMRobertaEncoder(nn.Module):
                     torch.utils.checkpoint.checkpoint(
                         self.layers[-1],
                         hidden_states_subset,
-                        use_reentrant=False,
+                        use_reentrant=self.use_reentrant,
                         mixer_kwargs=mixer_kwargs,
                     )
                 else:

rotary.py

@@ -4,7 +4,6 @@
 
 # Copyright (c) 2023, Tri Dao.
 
-import math
 from typing import Optional, Tuple, Union
 
 import torch
@@ -16,7 +15,10 @@ if torch.cuda.is_available():
     except ImportError:
 
         def apply_rotary(*args, **kwargs):
-            raise RuntimeError("RoPE requires flash-attention to be installed")
+            raise RuntimeError(
+                "FlashAttention is not installed. To proceed with training, please install FlashAttention. "
+                "For inference, you have two options: either install FlashAttention or disable it by setting use_flash_attn=False when loading the model."
+            )
 
 
 def rotate_half(x, interleaved=False):
@@ -169,12 +171,13 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
         seqlen_offsets: Union[int, torch.Tensor] = 0,
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seqlen: Optional[int] = None,
+        use_flash_attn: bool = True,
     ):
         # 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():
 
-            if torch.cuda.is_available():
+            if use_flash_attn:
                 # 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
@@ -288,7 +291,7 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
                 cu_seqlens=cu_seqlens,
                 max_seqlen=ctx.max_seqlen,
             )
-        return dqkv, None, None, None, None, None, None, None, None
+        return dqkv, None, None, None, None, None, None, None, None, None
 
 
 def apply_rotary_emb_qkv_(
@@ -301,6 +304,7 @@ def apply_rotary_emb_qkv_(
     seqlen_offsets: Union[int, torch.Tensor] = 0,
     cu_seqlens: Optional[torch.Tensor] = None,
     max_seqlen: Optional[int] = None,
+    use_flash_attn=True,
 ):
     """
     Arguments:
@@ -321,7 +325,7 @@ def apply_rotary_emb_qkv_(
     Apply rotary embedding *inplace* to the first rotary_dim of Q and K.
     """
     return ApplyRotaryEmbQKV_.apply(
-        qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, cu_seqlens, max_seqlen
+        qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets, cu_seqlens, max_seqlen, use_flash_attn,
     )
 
 
@@ -443,6 +447,7 @@ class RotaryEmbedding(torch.nn.Module):
         scale_base=None,
         pos_idx_in_fp32=True,
         device=None,
+        use_flash_attn=True,
     ):
         """
         interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
@@ -462,6 +467,7 @@ class RotaryEmbedding(torch.nn.Module):
         self.dim = dim
         self._base = float(base)
         self.pos_idx_in_fp32 = pos_idx_in_fp32
+        self.use_flash_attn = use_flash_attn
         # Generate and save the inverse frequency buffer (non trainable)
         inv_freq = self._compute_inv_freq(device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -588,6 +594,7 @@ class RotaryEmbedding(torch.nn.Module):
                     seqlen_offsets=seqlen_offset,
                     cu_seqlens=cu_seqlens,
                     max_seqlen=max_seqlen,
+                    use_flash_attn=self.use_flash_attn,
                 )
             else:
                 return apply_rotary_emb_qkv_(
@@ -600,6 +607,7 @@ class RotaryEmbedding(torch.nn.Module):
                     seqlen_offsets=seqlen_offset,
                     cu_seqlens=cu_seqlens,
                     max_seqlen=max_seqlen,
+                    use_flash_attn=self.use_flash_attn,
                 )
         else:
             q = qkv