Add FlashAttention + unpadding support

modeling_gptbert.py

@@ -5,11 +5,12 @@ import torch.nn as nn
 from torch.nn import functional as F
 from torch import _softmax_backward_data as _softmax_backward_data
 
-from functools import partial
+from functools import partial, lru_cache
 
 from .configuration_gptbert import GptBertConfig
 from transformers.modeling_utils import PreTrainedModel
 from transformers.activations import gelu_new
+from transformers.utils import is_flash_attn_2_available, is_flax_available
 from transformers.modeling_outputs import (
     MaskedLMOutput,
     MultipleChoiceModelOutput,
@@ -22,82 +23,71 @@ from transformers.modeling_outputs import (
 import math
 from typing import TYPE_CHECKING, Optional, Union, Tuple, List
 
-try:
-    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
-except ImportError:
-    pass
+if is_flash_attn_2_available():
+    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
+    from flash_attn.layers.rotary import RotaryEmbedding
+    from flash_attn.ops.triton.rotary import apply_rotary
 
 
-class ModelOutput:
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+def _unpad_input(input_ids: torch.Tensor, attention_mask: torch.Tensor):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = int(seqlens_in_batch.max().item())
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
 
-    def __init__(
-        self,
-        logits: torch.Tensor | None = None,
-        loss: torch.Tensor | float | None = None,
-        perplexity: torch.Tensor | float | None = None,
-        accuracy: float | None = None,
-        z_loss: torch.Tensor | float | None = None,
-        **kwargs
-    ):
-        self.logits: torch.Tensor | None
-        self.loss: torch.Tensor | float | None
-        self.perplexity: torch.Tensor | float | None
-        self.accuracy: float | None
-        self.z_loss: torch.Tensor | float | None
+    if input_ids.dim() == 2:
+        unpadded_inputs = input_ids.flatten()[indices]
+    else:
+        batch_size, sequence_length, *rest = input_ids.shape
+        shape = batch_size * sequence_length
+        unpadded_inputs = input_ids.view(shape, *rest)[indices]
 
-        self.logits = logits
-        self.loss = loss
-        self.perplexity = perplexity
-        self.accuracy = accuracy
-        self.z_loss = z_loss
+    return unpadded_inputs, indices, cu_seqlens, max_seqlen_in_batch
 
-        for attr, value in kwargs.items():
-            setattr(self, attr, value)
 
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+def _pad_output(input_ids: torch.Tensor, indices: torch.Tensor, batch_size: int, sequence_length: int) -> torch.Tensor:
+    if input_ids.dim() == 1:
+        output = torch.zeros(batch_size * sequence_length, dtype=input_ids.dtype, device=input_ids.device)
+        output[indices] = input_ids
+        padded_inputs = output.view(batch_size, sequence_length)
+    else:
+        _, *rest = input_ids.shape
+        output = torch.zeros(batch_size * sequence_length, *rest, dtype=input_ids.dtype, device=input_ids.device)
+        output[indices] = input_ids
+        padded_inputs = output.view(batch_size, sequence_length, *rest)
+
+    return padded_inputs
 
-class CastedLinear(nn.Linear):
 
+class CastedLinear(nn.Linear):
     def __init__(self, in_features, out_features, bias):
         super().__init__(in_features, out_features, bias=bias)
 
-    def reset_parameters(self) -> None:
-        std: float = math.sqrt(2.0 / (self.in_features + self.out_features))
-        nn.init.trunc_normal_(self.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
     def forward(self, x):
         return F.linear(x, self.weight.type_as(x), bias=self.bias.type_as(x) if self.bias is not None else None)
 
 
 class CastedLinearIn(nn.Linear):
-
     def __init__(self, in_features, out_features, bias):
         super().__init__(in_features, out_features, bias=bias)
         self.scale = nn.Parameter(torch.ones(in_features))
 
-    def reset_parameters(self) -> None:
-        std: float = math.sqrt(2.0 / (self.in_features + self.out_features))
-        nn.init.trunc_normal_(self.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
     def forward(self, x):
         return F.linear(x, (self.weight * (self.scale + 1.0).unsqueeze(0)).type_as(x), bias=self.bias.type_as(x) if self.bias is not None else None)
 
 
 class CastedLinearOut(nn.Linear):
-
     def __init__(self, in_features, out_features, bias):
         super().__init__(in_features, out_features, bias=bias)
         self.scale = nn.Parameter(torch.ones(out_features))
 
-    def reset_parameters(self) -> None:
-        std: float = math.sqrt(2.0 / (self.in_features + self.out_features))
-        nn.init.trunc_normal_(self.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
     def forward(self, x):
         return F.linear(x, (self.scale.unsqueeze(1) * self.weight).type_as(x), bias=self.bias.type_as(x) if self.bias is not None else None)
 
 
 class MultiCastedLinearOrtho(nn.Module):
-
     def __init__(self, in_features, out_features, bias):
         super().__init__()
         self.in_features = in_features
@@ -112,19 +102,11 @@ class MultiCastedLinearOrtho(nn.Module):
         else:
             self.bias = self.register_parameter("bias", None)
 
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        for i, weight in enumerate(self.weights):
-            std: float = math.sqrt(2.0 / (self.in_features + self.out_features[i]))
-            nn.init.trunc_normal_(weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
     def forward(self, x):
         return F.linear(x, torch.cat([weight for weight in self.weights], dim=0).type_as(x), bias=self.bias.type_as(x) if self.bias is not None else None)
 
 
 class MultiCastedLinearOrthoIn(nn.Module):
-
     def __init__(self, in_features, out_features, bias):
         super().__init__()
         self.in_features = in_features
@@ -141,23 +123,14 @@ class MultiCastedLinearOrthoIn(nn.Module):
 
         self.scale = nn.Parameter(torch.ones(in_features))
 
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        for weight in self.weights:
-            std = 0.5 * (self.in_features ** -0.5)
-            bound = (3 ** 0.5) * std
-            with torch.no_grad():
-                weight.uniform_(-bound, bound)
-
     def forward(self, x):
         return F.linear(x, (torch.cat([weight for weight in self.weights], dim=0) * (self.scale + 1.0).unsqueeze(0)).type_as(x), bias=self.bias.type_as(x) if self.bias is not None else None)
 
 
 class MultiCastedLinearOrthoOut(nn.Module):
-
     def __init__(self, in_features, out_features, bias):
         super().__init__()
+
         self.in_features = in_features
         self.out_features = out_features
 
@@ -172,15 +145,6 @@ class MultiCastedLinearOrthoOut(nn.Module):
 
         self.scale = nn.Parameter(torch.ones(sum(out_features)))
 
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        for weight in self.weights:
-            std = 0.5 * (self.in_features ** -0.5)
-            bound = (3 ** 0.5) * std
-            with torch.no_grad():
-                weight.uniform_(-bound, bound)
-
     def forward(self, x):
         return F.linear(x, (self.scale.unsqueeze(1) * torch.cat([weight for weight in self.weights], dim=0)).type_as(x), bias=self.bias.type_as(x) if self.bias is not None else None)
 
@@ -188,15 +152,12 @@ class MultiCastedLinearOrthoOut(nn.Module):
 class GeGLU(nn.Module):
     def forward(self, x):
         x, gate = x.chunk(2, dim=-1)
-        x = x * gelu_new(gate)
-        return x
+        return x * gelu_new(gate)
 
 
 class MaskedSoftmax(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, x: torch.Tensor, mask: torch.BoolTensor, dim: int) -> torch.Tensor:
-        ctx.dim: int
-
+    def forward(ctx, x: torch.Tensor, mask: torch.BoolTensor, dim: int):
         ctx.dim = dim
         x.masked_fill_(mask, float('-inf'))
         x = torch.softmax(x, ctx.dim)
@@ -205,47 +166,34 @@ class MaskedSoftmax(torch.autograd.Function):
         return x
 
     @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
-        output: torch.Tensor
-
+    def backward(ctx, grad_output: torch.Tensor):
         output, = ctx.saved_tensors
         inputGrad: torch.Tensor = _softmax_backward_data(grad_output, output, ctx.dim, output.dtype)
         return inputGrad, None, None
 
 
 class Encoder(nn.Module):
-
-    def __init__(self, config) -> None:
+    def __init__(self, config: GptBertConfig):
         super().__init__()
 
-        self.layers: nn.ModuleList[Layer]
-
         self.layers = nn.ModuleList([Layer(config, i) for i in range(config.num_layers)])
-
-        for i, layer in enumerate(self.layers):
-            for weight in layer.mlp.up_proj.weights:
-                weight.data *= math.sqrt(1.0 / (2.0 * (i + 1)))
-            layer.mlp.down_proj.weight.data *= math.sqrt(1.0 / (2.0 * (i + 1)))
-
         self.short_long_ratio = config.short_long_ratio
 
-    def set_window_length(self, config) -> None:
+    def set_window_length(self, config: GptBertConfig):
         for i, layer in enumerate(self.layers):
             if (i+1) % self.short_long_ratio == 0:
-                layer.set_window_length(config.window_length, config.not_flex)
+                layer.set_window_length(config.window_length)
             else:
-                layer.set_window_length(256, config.not_flex)
-
-    def forward(self, hidden_layer: torch.Tensor, embeddings: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        hidden_layer: List[torch.Tensor]
-        attention_probs: List[torch.Tensor]
+                layer.set_window_length(256)
 
+    def forward(self, hidden_layer: torch.Tensor, padding_info):
         hidden_states = []
         attention_probs = []
         v1 = None
+        embeddings = hidden_layer
 
         for layer in self.layers:
-            hidden_layer, v1, attention_p = layer(hidden_layer, embeddings, v1, mask)
+            hidden_layer, v1, attention_p = layer(hidden_layer, embeddings, v1, padding_info)
             hidden_states.append(hidden_layer)
             attention_probs.append(attention_p)
 
@@ -253,29 +201,22 @@ class Encoder(nn.Module):
 
 
 class Layer(nn.Module):
-
-    def __init__(self, config, layer_idx: int) -> None:
+    def __init__(self, config: GptBertConfig, layer_idx: int):
         super().__init__()
 
-        self.attention: SelfAttention
-        self.mlp: FeedForward
-
         self.attention = SelfAttention(config, layer_idx)
         self.mlp = FeedForward(config)
         self.lambdas = nn.Parameter(torch.tensor([0., 0., 1., 0., 1., 0.]))
 
-    def set_window_length(self, window_length: int, not_flex: bool) -> None:
-        self.attention.set_window_length(window_length, not_flex)
-
-    def forward(self, hidden_layer: torch.Tensor, embeddings: torch.Tensor, v1: torch.Tensor | None, mask: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor]:
-        output: torch.Tensor
-        attention_p: torch.Tensor
+    def set_window_length(self, window_length: int):
+        self.attention.set_window_length(window_length)
 
+    def forward(self, hidden_layer: torch.Tensor, embeddings: torch.Tensor, v1: torch.Tensor | None, padding_info):
         attention_output = (1 - self.lambdas[0]) * hidden_layer + self.lambdas[0] * embeddings
         qk_layer = (1 - self.lambdas[1]) * hidden_layer + self.lambdas[1] * embeddings
         mlp_layer = F.softplus(self.lambdas[2]) * ((1 - self.lambdas[3]) * hidden_layer + self.lambdas[3] * embeddings)
 
-        attention_output, v1, attention_p = self.attention(attention_output, qk_layer, v1, mask)
+        attention_output, v1, attention_p = self.attention(attention_output, qk_layer, v1, padding_info)
         mlp_layer = mlp_layer + attention_output
         hidden_layer = F.softplus(self.lambdas[4]) * ((1 - self.lambdas[5]) * hidden_layer + self.lambdas[5] * embeddings)
         output = hidden_layer + attention_output + self.mlp(mlp_layer)
@@ -284,97 +225,92 @@ class Layer(nn.Module):
 
 
 class Embedding(nn.Module):
-
-    def __init__(self, config) -> None:
+    def __init__(self, config: GptBertConfig):
         super().__init__()
 
         assert hasattr(config, "vocab_size"), "The config must have a vocab_size attribute!"
         assert hasattr(config, "hidden_size"), "The config must have a hidden_size attribute!"
         assert hasattr(config, "embedding_dropout_p"), "The model must have a embedding_dropout_p attribute!"
 
-        self.word_embedding: nn.Embedding
-        self.word_norm: nn.LayerNorm
-        self.dropout: nn.Dropout
-
         self.word_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
         self.word_norm = nn.LayerNorm(config.hidden_size, eps=config.word_norm_eps, elementwise_affine=False, bias=False)
         self.word_scale = nn.Parameter(torch.zeros(config.hidden_size))
 
         self.dropout = nn.Dropout(config.embedding_dropout_p)
 
-        self.initialize(config.hidden_size, config.vocab_size)
-
-    @torch.no_grad()
-    def initialize(self, hidden_size: int, vocab_size: int) -> None:
-        std: float
-
-        std = math.sqrt(2.0 / (hidden_size + vocab_size))
-        nn.init.trunc_normal_(self.word_embedding.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
-    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        word_embedding: torch.Tensor
-
+    def forward(self, input_ids: torch.Tensor):
         word_embedding = self.word_embedding(input_ids)
         word_embedding = self.word_norm(word_embedding)
-        word_embedding = (word_embedding * (self.word_scale + 1.0).unsqueeze(0).unsqueeze(0))
+        word_embedding = word_embedding * (self.word_scale + 1.0)
 
         return self.dropout(word_embedding)
 
 
 class MaskClassifier(nn.Module):
-
-    def __init__(self, config, embedding_weights: nn.Parameter) -> None:
+    def __init__(self, config: GptBertConfig, embedding_weights: nn.Parameter):
         super().__init__()
 
-        self.projection: CastedLinear
-        self.emb2vocab: CastedLinear
-        self.pre_norm: nn.LayerNorm
-        self.post_norm: nn.LayerNorm
-
         self.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.classifier_pre_norm_eps, elementwise_affine=config.classifier_pre_norm_affine)
         self.projection = CastedLinearIn(config.hidden_size, config.hidden_size, bias=False)
         self.post_norm = nn.LayerNorm(config.hidden_size, eps=config.classifier_post_norm_eps, elementwise_affine=config.classifier_post_norm_affine)
         self.emb2vocab = CastedLinearIn(config.hidden_size, config.vocab_size, bias=True)
 
-        self.initialize(config.hidden_size, config.vocab_size, embedding_weights)
-
-    @torch.no_grad()
-    def initialize(self, hidden_size: int, vocab_size: int, embedding_weights: nn.Parameter) -> None:
-        proj_std: float = math.sqrt(2.0 / (hidden_size + 4*hidden_size))
-
-        nn.init.trunc_normal_(self.projection.weight, mean=0.0, std=proj_std, a=-2*proj_std, b=2*proj_std)
-        self.emb2vocab.weight = embedding_weights
-        self.emb2vocab.bias.zero_()
-
-    def project(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        projection: torch.Tensor
-
-        projection = self.projection(hidden_layer)
-        projection = gelu_new(projection)
-        projection = self.post_norm(projection)
-
-        return projection
-
-    def calculate_output(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        return self.emb2vocab(hidden_layer)
-
-    def forward(self, hidden_layer: torch.Tensor, labels: torch.Tensor | None = None) -> torch.Tensor:
-        output: torch.Tensor
-
-        if labels is not None:
-            hidden_layer = torch.index_select(hidden_layer.flatten(0, 1), 0, torch.nonzero(labels.flatten() != -100).squeeze())
-
-        hidden_layer = self.pre_norm(hidden_layer)
-        hidden_layer = self.project(hidden_layer)
-        output = self.calculate_output(hidden_layer)
-
-        return output
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x)
+        x = self.projection(x)
+        x = gelu_new(x)
+        x = self.post_norm(x)
+        return self.emb2vocab(x)
+
+
+def flash_attention_forward(
+    qkv: torch.Tensor,
+    rotary_emb: UnpaddedRotaryEmbedding,
+    cu_seqlens: torch.Tensor,
+    max_seqlen: int,
+    local_attention: Tuple[int, int],
+    dropout_p: float,
+    deterministic: bool,
+    target_dtype: torch.dtype = torch.bfloat16,
+    **_kwargs,
+):
+    qkv = rotary_emb(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+
+    convert_dtype = qkv.dtype not in (torch.float16, torch.bfloat16)
+    if convert_dtype:
+        # FA2 implementation only supports fp16 and bf16. If FA2 is supported,
+        # bfloat16 must be supported as of FA2 2.5.7. (Turing GPUs not supported)
+        orig_dtype = qkv.dtype
+        qkv = qkv.to(target_dtype)
+
+        attn = flash_attn_varlen_qkvpacked_func(
+            qkv,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            dropout_p=dropout_p,
+            deterministic=deterministic,
+            window_size=local_attention,
+        )
+        attn = attn.to(orig_dtype)  # type: ignore
+    else:
+        attn = flash_attn_varlen_qkvpacked_func(
+            qkv,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            dropout_p=dropout_p,
+            deterministic=deterministic,
+            window_size=local_attention,
+        )
+    return attn
 
 
 class SelfAttention(nn.Module):
-
-    def __init__(self, config, layer_idx) -> None:
+    def __init__(self, config: GptBertConfig, layer_idx: int):
         super().__init__()
+
+        self.config = config
+        self.layer_idx = layer_idx
+
         self.d_qk = config.d_qk
         self.d_v = config.d_v
         self.num_attention_heads = config.num_attention_heads
@@ -398,72 +334,59 @@ class SelfAttention(nn.Module):
         self.q_scale = nn.Parameter(torch.ones(self.num_attention_heads, config.d_qk))
 
         self.dropout = nn.Dropout(config.attention_output_dropout_p)
+        self.attention_dropout = config.attention_dropout if hasattr(config, "attention_dropout") else 0.0
+        self.deterministic_flash_attn = getattr(config, "deterministic_flash_attn", False)
 
         theta = 160_000 if (layer_idx + 1) % config.short_long_ratio == 0 else 10_000
 
-        self.rope_embedding = RotaryPositionalEmbeddings(config, theta)
-        self.scale: float = 1.0 / math.sqrt(self.d_qk)
+        # Initialize rotary embeddings based on whether FlashAttention is available
+        if is_flash_attn_2_available():
+            self.rope_embedding = UnpaddedRotaryEmbedding(
+                dim=config.d_qk,
+                base=theta,
+                max_seqlen=config.max_sequence_length,
+                device=None,
+                dtype=None
+            )
+        else:
+            self.rope_embedding = RotaryPositionalEmbeddings(config, theta)
 
+        self.scale = 1.0 / math.sqrt(self.d_qk)
         self.dropout = nn.Dropout(config.attention_dropout if hasattr(config, "attention_dropout") else 0.0)
 
         self.lambdas = nn.Parameter(torch.tensor([0.5]))
 
-        self.initialize()
-
         self.sequence_length = config.max_sequence_length
         self.is_causal = config.is_decoder
-        self.not_flex = config.not_flex
-
-    @torch.no_grad()
-    def initialize(self) -> None:
-        std: float = math.sqrt(2.0 / (self.hidden_size + 4*self.hidden_size))
-        for weight in self.qk_proj.weights:
-            nn.init.trunc_normal_(weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        nn.init.trunc_normal_(self.v_proj.weight, mean=0.0, std=std, a=2*std, b=2*std)
-        self.out_proj.weight.data.zero_()
+        self.window_length = None
 
-    def set_window_length(self, window_length: int, not_flex: bool) -> None:
-        self.window_length: int = window_length
-        if not not_flex:
-            self.block_mask = self.create_block_mask(window_length)
+    def set_window_length(self, window_length: int):
+        self.window_length = window_length
 
-    def causal_mask_mode(self, window_length, b, _, q_idx, kv_idx):
-        return (q_idx >= kv_idx) & ((q_idx - kv_idx) < window_length)
-
-    def bidirectional_mask_mode(self, window_length, b, _, q_idx, kv_idx):
-        return ((q_idx - kv_idx) < window_length) & ((kv_idx - q_idx) < window_length)
-
-    def create_block_mask(self, window_length: int) -> torch.Tensor:
+    @lru_cache(maxsize=32)
+    def _get_window_mask(self, query_length: int, key_length: int, device: torch.device):
+        """Create and cache window attention mask."""
         if self.is_causal:
-            return create_block_mask(
-                partial(self.causal_mask_mode, self.window_length),
-                1, 1, self.sequence_length, self.sequence_length, device=self.k_scale.device
-            )
+            mask = torch.ones(query_length, key_length, dtype=torch.bool, device=device)
+            mask = ~mask.tril().triu(diagonal=-self.window_length)
         else:
-            return create_block_mask(
-                partial(self.bidirectional_mask_mode, self.window_length),
-                1, 1, self.sequence_length, self.sequence_length, device=self.k_scale.device
-            )
-
-    def attention_operation(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, padding_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        attention_scores: torch.Tensor
-        attention_probabilities: torch.Tensor
-        batch_size: int
-        query_length: int
-        key_length: int
+            mask = torch.ones(query_length, key_length, dtype=torch.bool, device=device)
+            mask = ~mask.tril(diagonal=self.window_length).triu(diagonal=-self.window_length)
+        return mask.view(1, 1, query_length, key_length)
 
+    def attention_operation(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, padding_mask: Optional[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Standard attention computation with masking."""
         batch_size, _, query_length, _ = query.size()
         _, _, key_length, _ = key.size()
 
-        if self.is_causal:
-            window_mask = ~torch.ones(query_length, key_length, dtype=torch.bool, device=self.k_scale.device).tril().triu(diagonal=-self.window_length).view(1, 1, query_length, key_length)
-        else:
-            window_mask = ~torch.ones(query_length, key_length, dtype=torch.bool, device=self.k_scale.device).tril(diagonal=self.window_length).triu(diagonal=-self.window_length).view(1, 1, query_length, key_length)
-        
-        if padding_mask is not None:
-            attention_mask = padding_mask | window_mask
-        else:
-            attention_mask = window_mask
+        # Use cached window mask
+        with torch.no_grad():
+            window_mask = self._get_window_mask(query_length, key_length, query.device)
+
+            if padding_mask is not None:
+                attention_mask = padding_mask | window_mask
+            else:
+                attention_mask = window_mask
 
         attention_scores = torch.bmm(query.flatten(0, 1), key.transpose(-1, -2).flatten(0, 1)) * self.scale  # shape: [B*H, T, T]
         attention_scores = attention_scores.view(batch_size, self.num_attention_heads, query_length, key_length)
@@ -476,120 +399,226 @@ class SelfAttention(nn.Module):
 
         return value, attention_probabilities.detach()
 
-    def forward(self, hidden_layer: torch.Tensor, qk_layer: torch.Tensor, v1: torch.Tensor | None, mask: torch.Tensor | None = None, doc_ids:  torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, hidden_layer: torch.Tensor, qk_layer: torch.Tensor, v1: torch.Tensor | None, padding_info)]:
+        # Get original shape info
+        if is_flash_attn_2_available() and isinstance(padding_info, tuple):
+            # Unpadded case
+            indices, cu_seqlens, max_seqlen = padding_info
+            total_seqlen = hidden_layer.size(0)
+            batch_size = cu_seqlens.size(0) - 1
+        else:
+            # Padded case
+            batch_size, seq_length = hidden_layer.size(0), hidden_layer.size(1)
+            hidden_layer = hidden_layer.transpose(0, 1)  # [seq_len, batch_size, hidden_size]
+            qk_layer = qk_layer.transpose(0, 1)
+
+
         hidden_layer = self.pre_v_norm(hidden_layer)
         qk_layer = self.pre_qk_norm(qk_layer)
 
         query, key = self.qk_proj(qk_layer).tensor_split([self.q_out_dim], dim=-1)
         value = self.v_proj(hidden_layer)
 
-        query_length: int = hidden_layer.size(0)
-        key_length: int = hidden_layer.size(0)
-        batch_size: int = hidden_layer.size(1)
+        if is_flash_attn_2_available() and isinstance(padding_info, tuple):
+            # Reshape for FlashAttention: (total_seqlen, num_heads, head_dim)
+            query = query.view(total_seqlen, self.num_attention_heads, self.d_qk)
+            key = key.view(total_seqlen, self.num_kv_heads, self.d_qk)
+            value = value.view(total_seqlen, self.num_kv_heads, self.d_v)
 
-        query = query.reshape(query_length, batch_size, self.num_attention_heads, self.d_qk).permute(1, 2, 0, 3)  # shape: [B, H, T, D]
-        key = key.reshape(key_length, batch_size, self.num_kv_heads, self.d_qk).permute(1, 2, 0, 3)  # shape: [B, H, T, D]
-        value = value.reshape(key_length, batch_size, self.num_kv_heads, self.d_qk).permute(1, 2, 0, 3)  # shape: [B, H, T, D]
+            # Apply layer norm and scaling
+            query = ((self.q_scale + 1.0).unsqueeze(0) * self.q_norm(query.float())).type_as(query)
+            key = ((self.k_scale + 1.0).unsqueeze(0) * self.k_norm(key.float())).type_as(key)
 
-        query, key = ((self.q_scale + 1.0).unsqueeze(1).unsqueeze(0) * self.q_norm(query.float())).type_as(query), ((self.k_scale + 1.0).unsqueeze(1).unsqueeze(0) * self.k_norm(key.float())).type_as(key)
+            if v1 is None:
+                v1 = value
+            value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
 
-        if v1 is None:
-            v1 = value
-        value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
+            # Prepare qkv for FlashAttention
+            if self.num_kv_heads == self.num_attention_heads:
+                # Standard MHA
+                qkv = torch.stack([query, key, value], dim=1)  # (total_seqlen, 3, num_heads, head_dim)
+            else:
+                # GQA case - need to repeat k,v heads
+                num_rep = self.num_attention_heads // self.num_kv_heads
+                key = key.repeat_interleave(num_rep, dim=1)
+                value = value.repeat_interleave(num_rep, dim=1)
+                qkv = torch.stack([query, key, value], dim=1)
+
+            # Determine window size for local attention
+            if self.window_length is not None and self.window_length > 0:
+                if self.is_causal:
+                    local_attention = (self.window_length - 1, 0)
+                else:
+                    local_attention = (self.window_length - 1, self.window_length - 1)
+            else:
+                local_attention = (-1, -1)
+
+            # Apply FlashAttention
+            output = flash_attention_forward(
+                qkv,
+                self.rope_embedding,
+                cu_seqlens,
+                max_seqlen,
+                local_attention,
+                self.attention_dropout if self.training else 0.0,
+                self.deterministic_flash_attn
+            )
 
-        query = self.rope_embedding(query)
-        key = self.rope_embedding(key)
+            # Reshape output back
+            output = output.view(total_seqlen, self.d_v * self.num_attention_heads)
+            attention_probabilities = None
 
-        if self.not_flex:
-            output, attention_probabilities = self.attention_operation(query, key, value, mask)
         else:
-            def document_score_mod(score, b, _, q_idx, kv_idx):
-                return torch.where(doc_ids[q_idx] == doc_ids[kv_idx], score, -float("inf"))
-
-            if self.is_causal:
-                block_mask = create_block_mask(
-                    partial(self.causal_mask_mode, self.window_length),
-                    1, 1, query_length, key_length, device=self.k_scale.device
-                )
-            else:
-                block_mask = create_block_mask(
-                    partial(self.bidirectional_mask_mode, self.window_length),
-                    1, 1, query_length, key_length, device=self.k_scale.device
-                )
+            # Standard attention path
+            query_length = hidden_layer.size(0)
+            key_length = hidden_layer.size(0)
 
-            output = flex_attention(
-                query, key, value, block_mask=block_mask, enable_gqa=True
-            )
-            attention_probabilities = None
+            query = query.reshape(query_length, batch_size, self.num_attention_heads, self.d_qk).permute(1, 2, 0, 3)
+            key = key.reshape(key_length, batch_size, self.num_kv_heads, self.d_qk).permute(1, 2, 0, 3)
+            value = value.reshape(key_length, batch_size, self.num_kv_heads, self.d_v).permute(1, 2, 0, 3)
 
-        output = output.permute(2, 0, 1, 3).flatten(2, 3)  # shape: [T, B, H*D]
-        output = self.inter_norm(output)
-        output = self.out_proj(output)
+            query = ((self.q_scale + 1.0).unsqueeze(1).unsqueeze(0) * self.q_norm(query.float())).type_as(query)
+            key = ((self.k_scale + 1.0).unsqueeze(1).unsqueeze(0) * self.k_norm(key.float())).type_as(key)
 
-        return self.dropout(output), v1, attention_probabilities
+            if v1 is None:
+                v1 = value
+            value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
 
+            # Apply rotary embeddings
+            query = self.rope_embedding(query)
+            key = self.rope_embedding(key)
 
-class FeedForward(nn.Module):
+            # Handle GQA for standard attention
+            if self.num_kv_heads != self.num_attention_heads:
+                num_rep = self.num_attention_heads // self.num_kv_heads
+                key = key.repeat_interleave(num_rep, dim=1)
+                value = value.repeat_interleave(num_rep, dim=1)
 
-    def __init__(self, config) -> None:
-        super().__init__()
+            output, attention_probabilities = self.attention_operation(query, key, value, padding_info if not isinstance(padding_info, tuple) else None)
+            output = output.permute(2, 0, 1, 3).flatten(2, 3)  # shape: [T, B, H*D]
 
-        self.up_proj: CastedLinear
-        self.down_proj: CastedLinear
-        self.pre_norm: nn.LayerNorm
-        self.inter_norm: nn.LayerNorm
-        self.activation: GeGLU
-        self.dropout: nn.Dropout
+        output = self.inter_norm(output)
+        output = self.out_proj(output)
 
+        # Handle output padding if necessary
+        if is_flash_attn_2_available() and isinstance(padding_info, tuple):
+            # Already in correct format for unpadded
+            pass
+        else:
+            # Transpose back to [batch_size, seq_len, hidden_size]
+            output = output.transpose(0, 1)
+
+        return self.dropout(output), v1, attention_probabilities
+
+class FeedForward(nn.Module):    
+    def __init__(self, config: GptBertConfig):
+        super().__init__()
         self.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.feed_forward_pre_norm_eps, elementwise_affine=config.feed_forward_pre_norm_affine)
         self.up_proj = MultiCastedLinearOrthoIn(config.hidden_size, [config.intermediate_size, config.intermediate_size], bias=False)
         self.activation = GeGLU()
         self.inter_norm = nn.LayerNorm(config.intermediate_size, eps=config.feed_forward_inter_norm_eps, elementwise_affine=config.feed_forward_inter_norm_affine)
         self.down_proj = CastedLinearIn(config.intermediate_size, config.hidden_size, bias=False)
         self.dropout = nn.Dropout(config.feed_forward_dropout_p)
+        
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x)
+        x = self.up_proj(x)
+        x = self.activation(x)
+        x = self.inter_norm(x.float()).type_as(x)
+        x = self.down_proj(x)
+        return self.dropout(x)
 
-        self.initialize(config.hidden_size)
 
-    @torch.no_grad()
-    def initialize(self, hidden_size: int) -> None:
-        std: float = math.sqrt(2.0 / (5*hidden_size))
+class ApplyRotaryEmbUnpad(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cos,
+        sin,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
+    ):
+        # (total_nnz, 3, nheads, headdim)
+        qkv = qkv.contiguous()
+        total_nnz, _three, _nheads, headdim = qkv.shape
+        # We need qkv to be contiguous so that when we reshape to combine (3, nheads) dimensions,
+        # we get the same tensor
+        # qk = rearrange(qkv[:, :2], "b_s t h d -> b_s (t h) d")
+        qk = qkv[:, :2].view(total_nnz, -1, headdim)
+        apply_rotary(
+            qk,
+            cos,
+            sin,
+            seqlen_offsets=0,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            interleaved=False,
+            inplace=True,
+        )
 
-        for weight in self.up_proj.weights:
-            nn.init.trunc_normal_(weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        self.down_proj.weight.data.zero_()
+        ctx.save_for_backward(cos, sin, cu_seqlens)
+        ctx.max_seqlen = max_seqlen
+        return qkv
 
-    def up_project(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        hidden_layer = self.pre_norm(hidden_layer)
-        return self.up_proj(hidden_layer)
+    @staticmethod
+    def backward(ctx, do):
+        cos, sin, cu_seqlens = ctx.saved_tensors
+        do = do.contiguous()
+        total_nnz, _three, _nheads, headdim = do.shape
+        # We need dqkv to be contiguous so that when we reshape to combine (3, nheads) dimensions,
+        # we get the same tensor
+        dqk = do[:, :2].view(total_nnz, -1, headdim)
+        apply_rotary(
+            dqk,
+            cos,
+            sin,
+            seqlen_offsets=0,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=ctx.max_seqlen,
+            interleaved=False,
+            inplace=True,
+            conjugate=True,
+        )
 
-    def activate(self, projection: torch.Tensor) -> torch.Tensor:
-        activated_projection: torch.Tensor
+        return do, None, None, None, None, None, None
 
-        activated_projection = self.activation(projection)
-        activated_projection = self.inter_norm(activated_projection.float()).type_as(projection)
 
-        return activated_projection
+def apply_rotary_unpadded(
+    qkv,
+    cos,
+    sin,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
+):
+    return ApplyRotaryEmbUnpad.apply(qkv, cos, sin, cu_seqlens, max_seqlen)
 
-    def down_project(self, activated_projection: torch.Tensor) -> torch.Tensor:
-        output: torch.Tensor
 
-        output = self.down_proj(activated_projection)
+class UnpaddedRotaryEmbedding(RotaryEmbedding):
+    def __init__(self, dim: int, base: float = 10000.0, max_seqlen: Optional[int] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        super().__init__(dim=dim, base=base, pos_idx_in_fp32=True, device=device, interleaved=False)
+        self.max_seqlen = max_seqlen
 
-        return self.dropout(output)
+        if max_seqlen is not None and device is not None and dtype is not None:
+            self._update_cos_sin_cache(max_seqlen, device=device, dtype=dtype)
 
-    def forward(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        output: torch.Tensor
+    def forward(self, qkv: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: Optional[int] = None) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if max_seqlen is not None:
+            self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
 
-        output = self.up_project(hidden_layer)
-        output = self.activate(output)
-        output = self.down_project(output)
+        qkv = apply_rotary_unpadded(
+            qkv,
+            self._cos_cached,
+            self._sin_cached,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
 
-        return output
+        return qkv
 
 
 class RotaryPositionalEmbeddings(nn.Module):
-
-    def __init__(self, config, theta: int) -> None:
+    def __init__(self, config, theta: int):
         super().__init__()
 
         assert hasattr(config, "d_qk"), "The config must have a d_qk attribute!"
@@ -615,7 +644,7 @@ class RotaryPositionalEmbeddings(nn.Module):
         self.register_buffer("cos_matrix", embedding.cos(), persistent=False)
         self.register_buffer("sin_matrix", embedding.sin(), persistent=False)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor):
         seq_len: int
         cos_matrix: torch.Tensor
         sin_matrix: torch.Tensor
@@ -647,18 +676,17 @@ class RotaryPositionalEmbeddings(nn.Module):
 
 class GptBertPreTrainedModel(PreTrainedModel):
     config_class = GptBertConfig
-    supports_gradient_checkpointing = False
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        raise NotImplementedError("Gradient checkpointing is not supported by this model")
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = False
 
     def _init_weights(self, module):
         pass
 
 
 class GptBertModel(GptBertPreTrainedModel):
-
-    def __init__(self, config, add_mlm_layer=False, **kwargs):
+    def __init__(self, config: GptBertConfig, add_mlm_layer=False, **kwargs):
         super().__init__(config, **kwargs)
         self.config = config
         self.hidden_size = config.hidden_size
@@ -680,7 +708,8 @@ class GptBertModel(GptBertPreTrainedModel):
     def get_contextualized_embeddings(
         self,
         input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None
     ) -> List[torch.Tensor]:
         if input_ids is not None:
             input_shape = input_ids.size()
@@ -690,35 +719,36 @@ class GptBertModel(GptBertPreTrainedModel):
         batch_size, seq_length = input_shape
         device = input_ids.device
 
-        # if attention_mask is None:
-        #     attention_mask = torch.zeros(batch_size, seq_length, dtype=torch.bool, device=device)
-        if attention_mask is not None:
-            attention_mask = ~attention_mask.bool()
-
-            if len(attention_mask.size()) == 2:
-                attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-            elif len(attention_mask.size()) == 3:
-                attention_mask = attention_mask.unsqueeze(1)
+        if attention_mask is None:
+            attention_mask = torch.ones(batch_size, seq_length, dtype=torch.bool, device=device)
+        elif attention_mask is not None:
+            if len(attention_mask.size()) != 2:
+                raise ValueError("Only attention mask with two dimensions is supported now.")
 
-            if self.config.is_decoder:
-                attention_mask = attention_mask | torch.triu(torch.ones(seq_length, seq_length, dtype=torch.bool, device=device), 1).unsqueeze(0).unsqueeze(0)
+        if is_flash_attn_2_available():
+            input_ids, indices, cu_seqlens, max_seqlen_in_batch = _unpad_input(input_ids, attention_mask)
+            padding_info = (indices, cu_seqlens, max_seqlen_in_batch)
+        else:
+            padding_info = attention_mask
 
-        static_embeddings = self.embedding(input_ids.t())
-        contextualized_embeddings, attention_probs = self.encoder(static_embeddings, static_embeddings, attention_mask)
-        contextualized_embeddings = [e.transpose(0, 1) for e in contextualized_embeddings]
+        static_embeddings = self.embedding(input_ids)
+        contextualized_embeddings, attention_probs = self.encoder(static_embeddings, padding_info)
         last_layer = contextualized_embeddings[-1]
-        contextualized_embeddings = [contextualized_embeddings[0]] + [
-            contextualized_embeddings[i] - contextualized_embeddings[i - 1]
-            for i in range(1, len(contextualized_embeddings))
-        ]
+
+        # Pad output if using FlashAttention
+        if is_flash_attn_2_available():
+            last_layer = _pad_output(last_layer, indices, batch_size, seq_length)
+            if output_hidden_states:
+                contextualized_embeddings = [_pad_output(layer, indices, batch_size, seq_length) for layer in contextualized_embeddings]
+            else:
+                contextualized_embeddings = None
+
         return last_layer, contextualized_embeddings, attention_probs
 
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         return_dict: Optional[bool] = None,
@@ -726,7 +756,9 @@ class GptBertModel(GptBertPreTrainedModel):
     ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
+        sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(
+            input_ids, attention_mask, output_hidden_states
+        )
 
         if not return_dict:
             return (
@@ -745,7 +777,7 @@ class GptBertModel(GptBertPreTrainedModel):
 class GptBertForMaskedLM(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["head"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=True, **kwargs)
 
     def get_output_embeddings(self):
@@ -799,7 +831,7 @@ class GptBertForMaskedLM(GptBertModel):
 
 
 class Classifier(nn.Module):
-    def __init__(self, config, num_labels: int):
+    def __init__(self, config: GptBertConfig, num_labels: int):
         super().__init__()
 
         drop_out = getattr(config, "cls_dropout", None)
@@ -826,34 +858,19 @@ class Classifier(nn.Module):
         nn.init.trunc_normal_(self.emb2vocab.weight, mean=0.0, std=proj_std, a=-2*proj_std, b=2*proj_std)
         self.emb2vocab.bias.zero_()
 
-    def project(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        projection: torch.Tensor
-
-        projection = self.pre_norm(hidden_layer)
-        projection = self.dropout(projection)
-        projection = self.projection(projection)
-        projection = gelu_new(projection)
-        projection = self.post_norm(projection)
-
-        return projection
-
-    def calculate_output(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        return self.emb2vocab(hidden_layer)
-
-    def forward(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        output: torch.Tensor
-        projection: torch.Tensor
-
-        projection = self.project(hidden_layer)
-        output = self.calculate_output(projection)
-
-        return output
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x)
+        x = self.dropout(x)
+        x = self.projection(x)
+        x = gelu_new(x)
+        x = self.post_norm(x)
+        return self.emb2vocab(x)
 
 
 class GptBertForCausalLM(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["head"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         config.is_decoder = True
         super().__init__(config, add_mlm_layer=True, **kwargs)
 
@@ -978,7 +995,7 @@ class GptBertForCausalLM(GptBertModel):
 class GptBertForSequenceClassification(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
@@ -1043,7 +1060,7 @@ class GptBertForSequenceClassification(GptBertModel):
 class GptBertForTokenClassification(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
@@ -1090,7 +1107,7 @@ class GptBertForTokenClassification(GptBertModel):
 class GptBertForQuestionAnswering(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
@@ -1157,7 +1174,7 @@ class GptBertForQuestionAnswering(GptBertModel):
 class GptBertForMultipleChoice(GptBertModel):
     _keys_to_ignore_on_load_unexpected = ["classifier"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = getattr(config, "num_labels", 2)