Upload encoder.py

encoder.py

@@ -6,13 +6,16 @@ from typing import List, Optional, Tuple, Union
 import torch
 from torch import Tensor, nn
 
-try:
-    from flash_attn import flash_attn_func
+# try:
+#     from flash_attn import flash_attn_func
 
-    IMPORT_FLASH = True
-except Exception as err:
-    IMPORT_FLASH = False
-    IMPORT_FLASH_ERR = err
+#     IMPORT_FLASH = True
+# except Exception as err:
+#     IMPORT_FLASH = False
+#     IMPORT_FLASH_ERR = err
+
+IMPORT_FLASH = False
+IMPORT_FLASH_ERR = "Flash Attention not installed."
 
 # from .utils import apply_masked_flash_attn, apply_rotary_pos_emb
 
@@ -35,59 +38,59 @@ def apply_rotary_pos_emb(
     return (q * cos) + (rtt_half(q) * sin), (k * cos) + (rtt_half(k) * sin)
 
 
-def apply_masked_flash_attn(
-    q: Tensor,
-    k: Tensor,
-    v: Tensor,
-    mask: Tensor,
-    h: int,
-    d_k: int,
-) -> Tensor:
-    """
-    Applies Flash Attention with padding masks.
-    """
-
-    from einops import rearrange
-    from flash_attn import flash_attn_varlen_func
-    from flash_attn.bert_padding import pad_input, unpad_input
-
-    pad_mask = ~mask[:, 0, :]
-    b, t = pad_mask.shape
-    q = q.view(b, t, h * d_k)
-    k = k.view(b, t, h * d_k)
-    v = v.view(b, t, h * d_k)
-
-    q_unpad, indices_q, _, max_seqlen_q = unpad_input(q, pad_mask)[:4]
-    q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
-
-    k_unpad = unpad_input(k, pad_mask)[0]
-    k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
-
-    v_unpad = unpad_input(v, pad_mask)[0]
-    v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
-
-    lengths_q = pad_mask.sum(1).to(torch.int32).to(q.device)
-    cu_seqlens_q = F.pad(lengths_q.cumsum(0), (1, 0), value=0).to(torch.int32)
-    max_seqlen_q = torch.max(lengths_q)
-
-    output_unpad = flash_attn_varlen_func(
-        q_unpad,
-        k_unpad,
-        v_unpad,
-        cu_seqlens_q,
-        cu_seqlens_q,
-        max_seqlen_q,
-        max_seqlen_q,
-    )
-
-    scores = pad_input(
-        rearrange(output_unpad, "nnz h d -> nnz (h d)"),
-        indices_q,
-        b,
-        t,
-    )
-
-    return scores
+# def apply_masked_flash_attn(
+#     q: Tensor,
+#     k: Tensor,
+#     v: Tensor,
+#     mask: Tensor,
+#     h: int,
+#     d_k: int,
+# ) -> Tensor:
+#     """
+#     Applies Flash Attention with padding masks.
+#     """
+
+#     from einops import rearrange
+#     from flash_attn import flash_attn_varlen_func
+#     from flash_attn.bert_padding import pad_input, unpad_input
+
+#     pad_mask = ~mask[:, 0, :]
+#     b, t = pad_mask.shape
+#     q = q.view(b, t, h * d_k)
+#     k = k.view(b, t, h * d_k)
+#     v = v.view(b, t, h * d_k)
+
+#     q_unpad, indices_q, _, max_seqlen_q = unpad_input(q, pad_mask)[:4]
+#     q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
+
+#     k_unpad = unpad_input(k, pad_mask)[0]
+#     k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
+
+#     v_unpad = unpad_input(v, pad_mask)[0]
+#     v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
+
+#     lengths_q = pad_mask.sum(1).to(torch.int32).to(q.device)
+#     cu_seqlens_q = F.pad(lengths_q.cumsum(0), (1, 0), value=0).to(torch.int32)
+#     max_seqlen_q = torch.max(lengths_q)
+
+#     output_unpad = flash_attn_varlen_func(
+#         q_unpad,
+#         k_unpad,
+#         v_unpad,
+#         cu_seqlens_q,
+#         cu_seqlens_q,
+#         max_seqlen_q,
+#         max_seqlen_q,
+#     )
+
+#     scores = pad_input(
+#         rearrange(output_unpad, "nnz h d -> nnz (h d)"),
+#         indices_q,
+#         b,
+#         t,
+#     )
+
+#     return scores
 
 
 class StridingSubsampling(nn.Module):
@@ -266,17 +269,17 @@ class RotaryPositionMultiHeadAttention(MultiHeadAttention):
             value.view(t, b, self.h * self.d_k).transpose(0, 1),
         )
 
-        if not self.flash_attn:
-            scores = torch.matmul(q, k.transpose(-2, -1) / math.sqrt(self.d_k))
-            out = self.forward_attention(v, scores, mask)
-        else:
-            if mask is None:
-                scores = flash_attn_func(q, k, v)
-            else:
-                scores = apply_masked_flash_attn(q, k, v, mask, self.h, self.d_k)
-
-            scores = scores.view(b, -1, self.h * self.d_k)
-            out = self.linear_out(scores)
+        # if not self.flash_attn:
+        scores = torch.matmul(q, k.transpose(-2, -1) / math.sqrt(self.d_k))
+        out = self.forward_attention(v, scores, mask)
+        # else:
+        #     if mask is None:
+        #         scores = flash_attn_func(q, k, v)
+        #     else:
+        #         scores = apply_masked_flash_attn(q, k, v, mask, self.h, self.d_k)
+
+        #     scores = scores.view(b, -1, self.h * self.d_k)
+        #     out = self.linear_out(scores)
 
         return out