merge recent changes

Signed-off-by: Meow <ongjackm@gmail.com>

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,6 +28,7 @@ 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,
         task_instructions: Optional[Dict[str, str]] = None,
@@ -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,6 +105,7 @@ 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

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,11 +11,9 @@ from torch.nn import Parameter
 from torch.nn import functional as F
 from transformers import PretrainedConfig
 
-from .modeling_xlm_roberta import (
-    XLMRobertaFlashConfig,
-    XLMRobertaModel,
-    XLMRobertaPreTrainedModel,
-)
+from .rotary import RotaryEmbedding
+from .modeling_xlm_roberta import (XLMRobertaFlashConfig, XLMRobertaModel,
+                                   XLMRobertaPreTrainedModel)
 
 
 def initialized_weights(
@@ -328,16 +326,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
             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
             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)
 
@@ -399,14 +394,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             adapter_mask = torch.full(
                 (num_examples,), task_id, dtype=torch.int32, device=self.device
             )
-        if task_type in ["query", "passage"]:
-            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