Repository: OpenNMT/OpenNMT-py

Issue:
Implement Flash Attention
Now that PyTorch 2.0 is out and supports native Flash Attention, perhaps OpenNMT-py should support it? This can help in creating document-level models with the longer context length instead of sentence-based models.

https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html

![image](https://user-images.githubusercontent.com/1308761/225953208-f01522a1-b932-4699-b739-a119fdac203d.png)



Discussion:
FYI we just committed a doc level transform. https://github.com/OpenNMT/OpenNMT-py/blob/master/docs/source/FAQ.md#context--doc-aware-transform
> 

Yes, I'm aware. But I'm looking more at like >1024 context length, preferably 4k or 8k long. Flash Attention should be helpful in the training efficiency.
PR are welcome :)
In our implementation it is not as straight forward as it may seem because we manage the caching and relative position embeddings in the `MultiHeadedAttention` class itself.
But one can try.

Is this something planned to be implemented still? :)
I'm not sure what's the actual max input length I can finetune some LLM on.
are you interested in sending a PR ? this is not very difficult to implement.

Resolving patch:
diff --git a/onmt/decoders/transformer.py b/onmt/decoders/transformer.py
--- a/onmt/decoders/transformer.py
+++ b/onmt/decoders/transformer.py
@@ -119,8 +119,9 @@ def forward(self, *args, **kwargs):
         full context alignement, :cite:`garg2019jointly`.
 
         Args:
-            * All arguments of _forward.
-            with_align (bool): whether return alignment attention.
+            * All arguments of _forward, of which
+            with_align (bool): needed to compute attn_align
+            return_attn (bool): to force MHA to return attns
 
         Returns:
             (FloatTensor, FloatTensor, FloatTensor or None):
@@ -131,7 +132,7 @@ def forward(self, *args, **kwargs):
         """
         with_align = kwargs.pop("with_align", False)
         layer_out, attns = self._forward(*args, **kwargs)
-        top_attn = attns[:, 0, :, :].contiguous()
+        top_attn = None if attns is None else attns[:, 0, :, :].contiguous()
         attn_align = None
         if with_align:
             if self.full_context_alignment:
@@ -174,10 +175,15 @@ def _compute_dec_mask(self, tgt_pad_mask, future):
             dec_mask = tgt_pad_mask
         return dec_mask
 
-    def _forward_self_attn(self, norm_layer_in, dec_mask, step):
+    def _forward_self_attn(self, norm_layer_in, dec_mask, step, return_attn=False):
         if self.self_attn_type == "scaled-dot":
             return self.self_attn(
-                norm_layer_in, norm_layer_in, norm_layer_in, mask=dec_mask, step=step
+                norm_layer_in,
+                norm_layer_in,
+                norm_layer_in,
+                mask=dec_mask,
+                step=step,
+                return_attn=return_attn,
             )
         elif self.self_attn_type == "average":
             return self.self_attn(norm_layer_in, mask=dec_mask, step=step)
@@ -269,6 +275,7 @@ def _forward(
         tgt_pad_mask,
         step=None,
         future=False,
+        return_attn=False,
     ):
         """A naive forward pass for transformer decoder.
 
@@ -281,6 +288,7 @@ def _forward(
             tgt_pad_mask (bool): ``(batch_size, 1, T)``
             step (int or None): stepwise decoding counter
             future (bool): If set True, do not apply future_mask.
+            return_attn (bool) : if set True requires attns output
 
         Returns:
             (FloatTensor, FloatTensor):
@@ -307,7 +315,11 @@ def _forward(
 
         if self.parallel_residual:
             ctx_attn, attns = self.context_attn(
-                enc_out, enc_out, norm_layer_in, mask=src_pad_mask
+                enc_out,
+                enc_out,
+                norm_layer_in,
+                mask=src_pad_mask,
+                return_attn=return_attn,
             )
             # feed_forward applies residual, so we remove and apply residual with un-normed
             layer_out = (
@@ -321,7 +333,7 @@ def _forward(
             query = self.dropout(self_attn) + layer_in
             norm_query = self.layer_norm_2(query)
             ctx_attn, attns = self.context_attn(
-                enc_out, enc_out, norm_query, mask=src_pad_mask
+                enc_out, enc_out, norm_query, mask=src_pad_mask, return_attn=return_attn
             )
             layer_out = self.feed_forward(self.dropout(ctx_attn) + query)
 
@@ -542,6 +554,7 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
         tgt_pad_mask = tgt[:, :, 0].eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]
 
         with_align = kwargs.pop("with_align", False)
+        return_attn = with_align or self._copy
         attn_aligns = []
 
         for layer in self.transformer_layers:
@@ -552,6 +565,7 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
                 tgt_pad_mask,
                 step=step,
                 with_align=with_align,
+                return_attn=return_attn,
             )
             if attn_align is not None:
                 attn_aligns.append(attn_align)
@@ -604,7 +618,9 @@ class TransformerLMDecoderLayer(TransformerDecoderLayerBase):
          See TransformerDecoderLayerBase
     """
 
-    def _forward(self, layer_in, tgt_pad_mask, step=None, future=False):
+    def _forward(
+        self, layer_in, tgt_pad_mask, step=None, future=False, return_attn=False
+    ):
         """A naive forward pass for transformer decoder.
 
         # T: could be 1 in the case of stepwise decoding or tgt_len
@@ -615,6 +631,7 @@ def _forward(self, layer_in, tgt_pad_mask, step=None, future=False):
             layer_cache (dict or None): cached layer info when stepwise decode
             step (int or None): stepwise decoding counter
             future (bool): If set True, do not apply future_mask.
+            return_attn (bool): If set True return attn
 
         Returns:
             (FloatTensor, FloatTensor):
@@ -635,7 +652,9 @@ def _forward(self, layer_in, tgt_pad_mask, step=None, future=False):
 
         norm_layer_in = self.layer_norm_1(layer_in)
 
-        attn_output, attns = self._forward_self_attn(norm_layer_in, dec_mask, step)
+        attn_output, attns = self._forward_self_attn(
+            norm_layer_in, dec_mask, step, return_attn=return_attn
+        )
 
         if self.parallel_residual:
             # feed_forward applies residual, so we remove and apply residual with un-normed
@@ -756,6 +775,7 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
         tgt_pad_mask = tgt[:, :, 0].eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]
 
         with_align = kwargs.pop("with_align", False)
+        return_attn = with_align or self._copy
         assert not with_align, "TransformerLMDecoder does not support align"
 
         for layer in self.transformer_layers:
@@ -764,6 +784,7 @@ def forward(self, tgt, enc_out=None, step=None, **kwargs):
                 tgt_pad_mask,
                 step=step,
                 with_align=with_align,
+                return_attn=return_attn,
             )
 
         dec_out = self.layer_norm(dec_out)
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
--- a/onmt/modules/multi_headed_attn.py
+++ b/onmt/modules/multi_headed_attn.py
@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 from typing import Optional, Tuple
+from torch.nn import functional as F
 import torch.nn as nn
 from torch.utils.checkpoint import checkpoint
 from torch.nn.utils import skip_init
@@ -26,6 +27,8 @@ def rotaryembeddings(dim: int, maxseqlen=8192, base=10000):
 
 
 def apply_rotary_emb(query, key, rope):
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
     query_ = query.float().reshape(*query.shape[:-1], -1, 2)
     query_ = torch.view_as_complex(query_)
     key_ = key.float().reshape(*key.shape[:-1], -1, 2)
@@ -33,7 +36,9 @@ def apply_rotary_emb(query, key, rope):
     rope = rope.view(1, query_.size(1), 1, query_.size(3))
     query_out = torch.view_as_real(query_ * rope).flatten(3)
     key_out = torch.view_as_real(key_ * rope).flatten(3)
-    return query_out.type_as(query), key_out.type_as(key)
+    return query_out.transpose(1, 2).type_as(query), key_out.transpose(1, 2).type_as(
+        key
+    )
 
 
 # Help functions for max_relative positions
@@ -327,6 +332,7 @@ def forward(
         query: Tensor,
         mask: Optional[Tensor] = None,
         step: Optional[int] = 0,
+        return_attn: Optional[bool] = False,
     ) -> Tuple[Tensor, Tensor]:
         """
         Compute the context vector and the attention vectors.
@@ -364,12 +370,7 @@ def forward(
                     start_pos = step
                     seqlen = query.size(2)
                     rope = self.rope[start_pos : start_pos + seqlen].to(query.device)
-
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
                     query, key = apply_rotary_emb(query, key, rope=rope)
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
 
                 if self.layer_cache[1]["keys"].numel() != 0:
                     key = torch.cat((self.layer_cache[1]["keys"], key), dim=2)
@@ -404,84 +405,95 @@ def forward(
                 start_pos = 0
                 seqlen = query.size(2)
                 rope = self.rope[start_pos : start_pos + seqlen].to(query.device)
-
-                query = query.transpose(1, 2)
-                key = key.transpose(1, 2)
                 query, key = apply_rotary_emb(query, key, rope=rope)
-                query = query.transpose(1, 2)
-                key = key.transpose(1, 2)
-        # 2) Calculate and scale scores.
-        query /= math.sqrt(self.dim_per_head)
+
         # expand key on heads dimension when it's less than query heads (multi-query variant)
         key = key.view(key.size(0), -1, 1, key.size(2), key.size(3)).repeat(
             1, 1, query.size(1) // key.size(1), 1, 1
         )
         key = key.view(key.size(0), query.size(1), key.size(3), key.size(4))
-        # batch x num_heads x query_len x key_len
-        scores = torch.matmul(query, key.transpose(2, 3))
-
-        if self.relative_attention_bias is not None:
-            q_len = key.size(2) if self.layer_cache[0] else query.size(2)
-            relative_position_bucket = compute_bias(
-                q_len,
-                key.size(2),
-                self.is_decoder,
-                self.max_relative_positions,
-                self.relative_positions_buckets,
-                device=key.device,
-            )
-            values = self.relative_attention_bias(
-                relative_position_bucket
-            )  # shape (query_length, key_length, num_heads)
-            position_bias = values.permute([2, 0, 1]).unsqueeze(
-                0
-            )  # shape (1, num_heads, query_length, key_length)
-            if self.layer_cache[0]:
-                position_bias = position_bias[:, :, -query.size(2) :, :]
-            scores.add_(position_bias)
-
-        elif self.relative_positions_embeddings is not None:
-            key_len = key.size(2)
-            # 1 or key_len x key_len
-            relative_positions_matrix = gen_relative_positions(
-                key_len,
-                self.max_relative_positions,
-                cache=self.layer_cache[0],
-                device=key.device,
+
+        # 2) When standard pos. enc. or rotary, use flash attention
+        if self.max_relative_positions in [-1, 0] and not return_attn:
+
+            attn_output = F.scaled_dot_product_attention(
+                query, key, value, None, 0.0, is_causal=mask is not None
             )
-            #  1 or key_len x key_len x dim_per_head
-            relations_keys = self.relative_positions_embeddings(
-                relative_positions_matrix
+            x = unshape(attn_output)
+            attn_output = self.maybe_ckpt(self.final_linear, x)
+            return attn_output, None
+
+        else:
+            query /= math.sqrt(self.dim_per_head)
+            # batch x num_heads x query_len x key_len
+            scores = torch.matmul(query, key.transpose(2, 3))
+
+            if self.relative_attention_bias is not None:
+                q_len = key.size(2) if self.layer_cache[0] else query.size(2)
+                relative_position_bucket = compute_bias(
+                    q_len,
+                    key.size(2),
+                    self.is_decoder,
+                    self.max_relative_positions,
+                    self.relative_positions_buckets,
+                    device=key.device,
+                )
+                values = self.relative_attention_bias(
+                    relative_position_bucket
+                )  # shape (query_length, key_length, num_heads)
+                position_bias = values.permute([2, 0, 1]).unsqueeze(
+                    0
+                )  # shape (1, num_heads, query_length, key_length)
+                if self.layer_cache[0]:
+                    position_bias = position_bias[:, :, -query.size(2) :, :]
+                scores.add_(position_bias)
+
+            elif self.relative_positions_embeddings is not None:
+                key_len = key.size(2)
+                # 1 or key_len x key_len
+                relative_positions_matrix = gen_relative_positions(
+                    key_len,
+                    self.max_relative_positions,
+                    cache=self.layer_cache[0],
+                    device=key.device,
+                )
+                #  1 or key_len x key_len x dim_per_head
+                relations_keys = self.relative_positions_embeddings(
+                    relative_positions_matrix
+                )
+                scores.add_(relative_matmul(query, relations_keys, True))
+            elif self.max_relative_positions == -2:  # Alibi
+                scores = self.alibi(scores)
+
+            scores = scores.float()
+
+            if mask is not None:
+                # not 100% necessary but expand to nb of heads
+                mask = mask.expand(-1, self.head_count, -1, -1)
+                # now mask and scores have the same shape
+                scores = scores.masked_fill(mask, -1e18)
+
+            # 3) Apply attention dropout and compute context vectors.
+            attn = self.softmax(scores).to(query.dtype)
+            drop_attn = self.dropout(attn)
+            # expand value on heads dimension when it's less than query heads (multi-query variant)
+            value = value.view(
+                value.size(0), -1, 1, value.size(2), value.size(3)
+            ).repeat(1, 1, query.size(1) // value.size(1), 1, 1)
+            value = value.view(
+                value.size(0), query.size(1), value.size(3), value.size(4)
             )
-            scores.add_(relative_matmul(query, relations_keys, True))
-        elif self.max_relative_positions == -2:  # Alibi
-            scores = self.alibi(scores)
-
-        scores = scores.float()
-
-        if mask is not None:
-            # not 100% necessary but expand to nb of heads
-            mask = mask.expand(-1, self.head_count, -1, -1)
-            # now mask and scores have the same shape
-            scores = scores.masked_fill(mask, -1e18)
-
-        # 3) Apply attention dropout and compute context vectors.
-        attn = self.softmax(scores).to(query.dtype)
-        drop_attn = self.dropout(attn)
-        # expand value on heads dimension when it's less than query heads (multi-query variant)
-        value = value.view(value.size(0), -1, 1, value.size(2), value.size(3)).repeat(
-            1, 1, query.size(1) // value.size(1), 1, 1
-        )
-        value = value.view(value.size(0), query.size(1), value.size(3), value.size(4))
-        context_original = torch.matmul(drop_attn, value)
+            context_original = torch.matmul(drop_attn, value)
 
-        if self.relative_positions_embeddings is not None:
-            # We use the same embeddings for key and value
-            relations_values = relations_keys
-            context_original.add_(relative_matmul(drop_attn, relations_values, False))
+            if self.relative_positions_embeddings is not None:
+                # We use the same embeddings for key and value
+                relations_values = relations_keys
+                context_original.add_(
+                    relative_matmul(drop_attn, relations_values, False)
+                )
 
-        context = unshape(context_original)
+            context = unshape(context_original)
 
-        attn_output = self.maybe_ckpt(self.final_linear, context)
+            attn_output = self.maybe_ckpt(self.final_linear, context)
 
-        return attn_output, attn
+            return attn_output, attn


Test patch: