fix-device-issues (#38)

- fix: inference issues (5365db89d3c1ad3dabfed37554c9f10e3a5918de)


Co-authored-by: Jack Min Ong <Jackmin108@users.noreply.huggingface.co>

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)
 
@@ -328,7 +329,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             )
         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)
 

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

rotary.py

@@ -9,6 +9,17 @@ from typing import Optional, Tuple, Union
 import torch
 from einops import rearrange, repeat
 
+if torch.cuda.is_available():
+    try:
+        from flash_attn.ops.triton.rotary import apply_rotary
+    except ImportError:
+
+        def apply_rotary(*args, **kwargs):
+            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):
     if not interleaved:
@@ -60,8 +71,6 @@ class ApplyRotaryEmb(torch.autograd.Function):
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seqlen: Optional[int] = None,
     ):
-        from flash_attn.ops.triton.rotary import apply_rotary
-
         out = apply_rotary(
             x,
             cos,
@@ -88,8 +97,6 @@ class ApplyRotaryEmb(torch.autograd.Function):
 
     @staticmethod
     def backward(ctx, do):
-        from flash_attn.ops.triton.rotary import apply_rotary
-
         seqlen_offsets = ctx.seqlen_offsets
         if seqlen_offsets is None:
             cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
@@ -171,8 +178,6 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
         if cos_k is None and sin_k is None and qkv.is_contiguous():
 
             if use_flash_attn:
-                from flash_attn.ops.triton.rotary import apply_rotary
-
                 # 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
@@ -203,8 +208,6 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
                 )
                 qkv = torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
         else:
-            from flash_attn.ops.triton.rotary import apply_rotary
-
             cos_k = cos if cos_k is None else cos_k
             sin_k = sin if sin_k is None else sin_k
             q, k = qkv[..., 0, :, :], qkv[..., 1, :, :]
@@ -241,8 +244,6 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
 
     @staticmethod
     def backward(ctx, dqkv):
-        from flash_attn.ops.triton.rotary import apply_rotary
-
         seqlen_offsets = ctx.seqlen_offsets
         if seqlen_offsets is None:
             cos, sin, cos_k, sin_k, cu_seqlens, seqlen_offsets = ctx.saved_tensors
@@ -340,8 +341,6 @@ class ApplyRotaryEmbKV_(torch.autograd.Function):
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seqlen: Optional[int] = None,
     ):
-        from flash_attn.ops.triton.rotary import apply_rotary
-
         # batch, seqlen, two, nheads, headdim = kv.shape
         assert kv.shape[-3] == 2
         k = kv[..., 0, :, :]
@@ -369,8 +368,6 @@ class ApplyRotaryEmbKV_(torch.autograd.Function):
 
     @staticmethod
     def backward(ctx, dkv):
-        from flash_attn.ops.triton.rotary import apply_rotary
-
         seqlen_offsets = ctx.seqlen_offsets
         if seqlen_offsets is None:
             cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors