FlashAttention support

modeling_gptbert.py

@@ -3,13 +3,13 @@ from __future__ import annotations
 import torch
 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, logging
 from transformers.modeling_outputs import (
     MaskedLMOutput,
     MultipleChoiceModelOutput,
@@ -22,111 +22,82 @@ from transformers.modeling_outputs import (
 import math
 from typing import TYPE_CHECKING, Optional, Union, Tuple, List
 
+logger = logging.get_logger(__name__)
+
+
+# Workaround for transformers < 4.36.0 check_imports issue
+# See: https://github.com/huggingface/transformers/issues/28459
 try:
-    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+    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
+    else:
+        flash_attn_varlen_qkvpacked_func, RotaryEmbedding, apply_rotary = None, object, None
+        logger.warning_once(
+            "NorBERT4 støtter FlashAttention, men det er ikke funnet i miljøet ditt. Du bør vurdere å oppdatere miljøet ditt for å få raskere og mindre minnekrevende behandling."
+        )
 except ImportError:
-    pass
+    flash_attn_varlen_qkvpacked_func, RotaryEmbedding, apply_rotary = None, object, None
+    logger.warning_once(
+        "NorBERT4 støtter FlashAttention, men det er ikke funnet i miljøet ditt. Du bør vurdere å oppdatere miljøet ditt for å få raskere og mindre minnekrevende behandling."
+    )
 
 
-class ModelOutput:
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+@torch.compiler.disable()
+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
-        self.out_features = out_features
-
-        self.weights = nn.ParameterList()
-        for out_feature in out_features:
-            self.weights.append(nn.Parameter(torch.empty((out_feature, in_features))))
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(sum(out_features)))
-        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
         self.out_features = out_features
 
@@ -141,244 +112,217 @@ 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
-
-        self.weights = nn.ParameterList()
-        for out_feature in out_features:
-            self.weights.append(nn.Parameter(torch.empty((out_feature, in_features))))
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(sum(out_features)))
-        else:
-            self.bias = self.register_parameter("bias", None)
-
-        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)
-
-
 class GeGLU(nn.Module):
     def forward(self, x):
         x, gate = x.chunk(2, dim=-1)
-        x = x * gelu_new(gate)
-        return x
-
-
-class MaskedSoftmax(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x: torch.Tensor, mask: torch.BoolTensor, dim: int) -> torch.Tensor:
-        ctx.dim: int
-
-        ctx.dim = dim
-        x.masked_fill_(mask, float('-inf'))
-        x = torch.softmax(x, ctx.dim)
-        x.masked_fill_(mask, 0.0)
-        ctx.save_for_backward(x)
-        return x
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
-        output: torch.Tensor
-
-        output, = ctx.saved_tensors
-        inputGrad: torch.Tensor = _softmax_backward_data(grad_output, output, ctx.dim, output.dtype)
-        return inputGrad, None, None
-
+        return x * gelu_new(gate)
 
-class Encoder(nn.Module):
 
-    def __init__(self, config) -> None:
+class Embedding(nn.Module):
+    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:
-        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)
-            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]
-
-        hidden_states = []
-        attention_probs = []
-        v1 = None
-
-        for layer in self.layers:
-            hidden_layer, v1, attention_p = layer(hidden_layer, embeddings, v1, mask)
-            hidden_states.append(hidden_layer)
-            attention_probs.append(attention_p)
+        self.word_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.word_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False, bias=False)
+        self.word_scale = nn.Parameter(torch.zeros(config.hidden_size))
+        self.dropout = nn.Dropout(config.embedding_dropout)
 
-        return hidden_states, attention_probs
+    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)
 
+        return self.dropout(word_embedding)
 
-class Layer(nn.Module):
 
-    def __init__(self, config, layer_idx: int) -> None:
+class LMClassifier(nn.Module):
+    def __init__(self, config: GptBertConfig, n_labels: 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
-
-        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)
-        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)
-
-        return output, v1, attention_p
-
+        self.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.projection = CastedLinearIn(config.hidden_size, config.hidden_size, bias=False)
+        self.post_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.emb2vocab = CastedLinearIn(config.hidden_size, n_labels, bias=True)
+
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x.float()).type_as(x)
+        x = self.projection(x)
+        x = gelu_new(x)
+        x = self.post_norm(x.float()).type_as(x)
+        x = self.emb2vocab(x)
+        return x
 
-class Embedding(nn.Module):
 
-    def __init__(self, config) -> None:
+class Classifier(nn.Module):
+    def __init__(self, config: GptBertConfig, n_labels: int):
         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.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.projection = CastedLinearIn(config.hidden_size, config.hidden_size, bias=False)
+        self.post_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.dropout = nn.Dropout(config.classifier_dropout)
+        self.output_projection = CastedLinearIn(config.hidden_size, n_labels, bias=True)
+
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x.float()).type_as(x)
+        x = self.projection(x)
+        x = gelu_new(x)
+        x = self.post_norm(x.float()).type_as(x)
+        x = self.dropout(x)
+        x = self.output_projection(x)
+        return x
 
-        self.initialize(config.hidden_size, config.vocab_size)
 
-    @torch.no_grad()
-    def initialize(self, hidden_size: int, vocab_size: int) -> None:
-        std: float
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+def flash_attention_forward(qkv: torch.Tensor, rotary_emb: UnpaddedRotaryEmbedding, cu_seqlens: torch.Tensor, max_seqlen: int, causal: bool, 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,
+            causal=False
+        )
+        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,
+            causal=False
+        )
+    return attn
 
-        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
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+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)
+
+        ctx.save_for_backward(cos, sin, cu_seqlens)
+        ctx.max_seqlen = max_seqlen
+        return qkv
 
-        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))
+    @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,
+        )
 
-        return self.dropout(word_embedding)
+        return do, None, None, None, None, None, None
 
 
-class MaskClassifier(nn.Module):
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+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 __init__(self, config, embedding_weights: nn.Parameter) -> None:
-        super().__init__()
 
-        self.projection: CastedLinear
-        self.emb2vocab: CastedLinear
-        self.pre_norm: nn.LayerNorm
-        self.post_norm: nn.LayerNorm
+# from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
+class UnpaddedRotaryEmbedding(RotaryEmbedding):
+    def __init__(self, dim: int, base: float = 10000.0, max_seqlen: Optional[int] = None):
+        super().__init__(dim=dim, base=base, pos_idx_in_fp32=True, device=None, interleaved=False)
+        self.max_seqlen = max_seqlen
 
-        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)
+    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)
 
-        self.initialize(config.hidden_size, config.vocab_size, embedding_weights)
+        qkv = apply_rotary_unpadded(
+            qkv,
+            self._cos_cached,
+            self._sin_cached,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
 
-    @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))
+        return qkv
 
-        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
+class RotaryPositionalEmbeddings(nn.Module):
+    def __init__(self, config, theta: int):
+        super().__init__()
 
-        projection = self.projection(hidden_layer)
-        projection = gelu_new(projection)
-        projection = self.post_norm(projection)
+        head_size = config.query_key_head_size
+        assert head_size % 2 == 0
+        max_seq_len = config.max_sequence_length
 
-        return projection
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_size, 2, dtype=torch.float32) / head_size))
+        pos = torch.arange(max_seq_len, dtype=torch.float32)
+        embedding = torch.einsum('n, d -> nd', pos, inv_freq)
+        embedding = torch.cat([embedding, embedding], dim=-1).unsqueeze(0)
+        self.register_buffer("cos_matrix", embedding.cos(), persistent=False)
+        self.register_buffer("sin_matrix", embedding.sin(), persistent=False)
 
-    def calculate_output(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        return self.emb2vocab(hidden_layer)
+    def forward(self, x: torch.Tensor):
+        hidden_layer = x.float()
 
-    def forward(self, hidden_layer: torch.Tensor, labels: torch.Tensor | None = None) -> torch.Tensor:
-        output: torch.Tensor
+        seq_len = x.shape[2]
 
-        if labels is not None:
-            hidden_layer = torch.index_select(hidden_layer.flatten(0, 1), 0, torch.nonzero(labels.flatten() != -100).squeeze())
+        cos_matrix = self.cos_matrix[:, None, :seq_len, :]
+        sin_matrix = self.sin_matrix[:, None, :seq_len, :]
 
-        hidden_layer = self.pre_norm(hidden_layer)
-        hidden_layer = self.project(hidden_layer)
-        output = self.calculate_output(hidden_layer)
+        x_rotate_half = torch.cat(
+            [
+                -hidden_layer[:, :, :, x.size(-1) // 2:],
+                hidden_layer[:, :, :, :x.size(-1) // 2]
+            ],
+            dim=-1
+        )
 
-        return output
+        out = hidden_layer * cos_matrix + x_rotate_half * sin_matrix
+        return out.type_as(x)
 
 
 class SelfAttention(nn.Module):
-
-    def __init__(self, config, layer_idx) -> None:
+    def __init__(self, config: GptBertConfig, layer_idx: int):
         super().__init__()
-        self.d_qk = config.d_qk
-        self.d_v = config.d_v
+
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.d_qk = config.query_key_head_size
+        self.d_v = config.value_head_size
         self.num_attention_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_kv_heads
+        self.num_kv_heads = config.num_attention_heads
         self.hidden_size = config.hidden_size
 
         self.q_out_dim = self.d_qk * self.num_attention_heads
@@ -389,256 +333,228 @@ class SelfAttention(nn.Module):
         self.v_proj = CastedLinearIn(self.hidden_size, self.v_out_dim, bias=False)
         self.out_proj = CastedLinearIn(self.d_v*self.num_attention_heads, self.hidden_size, bias=False)
 
-        self.pre_v_norm = nn.LayerNorm(config.hidden_size, eps=config.attention_pre_norm_eps, elementwise_affine=config.attention_pre_norm_affine)
-        self.pre_qk_norm = nn.LayerNorm(config.hidden_size, eps=config.attention_pre_norm_eps, elementwise_affine=config.attention_pre_norm_affine)
-        self.inter_norm = nn.LayerNorm(self.d_v * self.num_attention_heads, eps=config.attention_inter_norm_eps, elementwise_affine=config.attention_inter_norm_affine)
-        self.q_norm = nn.LayerNorm(config.d_qk, eps=config.attention_pre_norm_eps, elementwise_affine=False, bias=False)
-        self.k_norm = nn.LayerNorm(config.d_qk, eps=config.attention_pre_norm_eps, elementwise_affine=False, bias=False)
-        self.k_scale = nn.Parameter(torch.ones(self.num_kv_heads, config.d_qk))
-        self.q_scale = nn.Parameter(torch.ones(self.num_attention_heads, config.d_qk))
-
-        self.dropout = nn.Dropout(config.attention_output_dropout_p)
+        self.pre_v_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.pre_qk_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.inter_norm = nn.LayerNorm(self.d_v * self.num_attention_heads, eps=config.layer_norm_eps, elementwise_affine=False)
+        self.q_norm = nn.LayerNorm(self.d_qk, eps=config.layer_norm_eps, elementwise_affine=False, bias=False)
+        self.k_norm = nn.LayerNorm(self.d_qk, eps=config.layer_norm_eps, elementwise_affine=False, bias=False)
+        self.k_scale = nn.Parameter(torch.ones(self.num_kv_heads, self.d_qk))
+        self.q_scale = nn.Parameter(torch.ones(self.num_attention_heads, self.d_qk))
 
-        theta = 160_000 if (layer_idx + 1) % config.short_long_ratio == 0 else 10_000
+        self.dropout = nn.Dropout(config.hidden_dropout)
 
-        self.rope_embedding = RotaryPositionalEmbeddings(config, theta)
-        self.scale: float = 1.0 / math.sqrt(self.d_qk)
+        theta = 160_000 if (layer_idx + 1) % config.local_global_ratio == 0 else 10_000
 
-        self.dropout = nn.Dropout(config.attention_dropout if hasattr(config, "attention_dropout") else 0.0)
+        # Initialize rotary embeddings based on whether FlashAttention is available
+        if is_flash_attn_2_available():
+            self.rope_embedding = UnpaddedRotaryEmbedding(dim=self.d_qk, base=theta, max_seqlen=config.max_sequence_length)
+        else:
+            self.rope_embedding = RotaryPositionalEmbeddings(config, theta)
 
+        self.scale = 1.0 / math.sqrt(self.d_qk)
         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_()
-
-    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)
+        self.window_length = None
 
-    def causal_mask_mode(self, window_length, b, _, q_idx, kv_idx):
-        return (q_idx >= kv_idx) & ((q_idx - kv_idx) < window_length)
+    def set_window_length(self, window_length: int):
+        self.window_length = 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:
+    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
+        # 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
+
+        output = F.scaled_dot_product_attention(
+            query=query,
+            key=key,
+            value=value,
+            attn_mask=attention_mask,
+            dropout_p=self.config.attention_dropout if self.training else 0.0,
+            is_causal=self.is_causal
+        )
+        return output
+
+    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():
+            # Unpadded case
+            indices, cu_seqlens, max_seqlen = padding_info
+            total_seqlen = hidden_layer.size(0)
+            batch_size = cu_seqlens.size(0) - 1
         else:
-            attention_mask = window_mask
+            # Padded case
+            batch_size, seq_length = hidden_layer.size(0), hidden_layer.size(1)
 
-        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)
+        hidden_layer = self.pre_v_norm(hidden_layer.float()).type_as(hidden_layer)
+        qk_layer = self.pre_qk_norm(qk_layer.float()).type_as(qk_layer)
 
-        attention_probabilities = MaskedSoftmax.apply(attention_scores, attention_mask, -1)
-        attention_probabilities = self.dropout(attention_probabilities)
+        query, key = self.qk_proj(qk_layer).tensor_split([self.q_out_dim], dim=-1)
+        value = self.v_proj(hidden_layer)
 
-        value = torch.bmm(attention_probabilities.flatten(0, 1), value.flatten(0, 1))
-        value = value.view(batch_size, self.num_attention_heads, query_length, self.d_v)
+        if is_flash_attn_2_available():
+            # 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)
 
-        return value, attention_probabilities.detach()
+            # 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)
 
-    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]:
-        hidden_layer = self.pre_v_norm(hidden_layer)
-        qk_layer = self.pre_qk_norm(qk_layer)
+            if v1 is None:
+                v1 = value
+            value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
 
-        query, key = self.qk_proj(qk_layer).tensor_split([self.q_out_dim], dim=-1)
-        value = self.v_proj(hidden_layer)
+            # Prepare qkv for FlashAttention
+            qkv = torch.stack([query, key, value], dim=1)  # (total_seqlen, 3, num_heads, head_dim)
 
-        query_length: int = hidden_layer.size(0)
-        key_length: int = hidden_layer.size(0)
-        batch_size: int = hidden_layer.size(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,
+                self.is_causal,
+                local_attention,
+                self.config.attention_dropout if self.training else 0.0,
+                self.config.deterministic_flash_attn
+            )
 
-        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]
+            # Reshape output back
+            output = output.view(total_seqlen, self.d_v * self.num_attention_heads)
 
-        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)
+        else:
+            # Standard attention path
+            query_length = query.size(1)
+            key_length = key.size(1)
 
-        if v1 is None:
-            v1 = value
-        value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
+            query = query.reshape(batch_size, query_length, self.num_attention_heads, self.d_qk).transpose(1, 2)
+            key = key.reshape(batch_size, key_length, self.num_kv_heads, self.d_qk).transpose(1, 2)
+            value = value.reshape(batch_size, key_length, self.num_kv_heads, self.d_v).transpose(1, 2)
 
-        query = self.rope_embedding(query)
-        key = self.rope_embedding(key)
+            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)
 
-        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
-                )
+            if v1 is None:
+                v1 = value
             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
-                )
+                value = (1 - self.lambdas[0]) * value + self.lambdas[0] * v1
 
-            output = flex_attention(
-                query, key, value, block_mask=block_mask, enable_gqa=True
-            )
-            attention_probabilities = None
+            # Apply rotary embeddings
+            query = self.rope_embedding(query)
+            key = self.rope_embedding(key)
 
-        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)
+            output = self.attention_operation(query, key, value, padding_info)
+            output = output.transpose(1, 2).flatten(2, 3)  # shape: [B, T, H*D]
 
-        return self.dropout(output), v1, attention_probabilities
+        output = self.inter_norm(output.float()).type_as(output)
+        output = self.out_proj(output)
+        output = self.dropout(output)
 
+        return output, v1
 
-class FeedForward(nn.Module):
 
-    def __init__(self, config) -> None:
+class FeedForward(nn.Module):    
+    def __init__(self, config: GptBertConfig):
         super().__init__()
-
-        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
-
-        self.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.feed_forward_pre_norm_eps, elementwise_affine=config.feed_forward_pre_norm_affine)
+        self.pre_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
         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.inter_norm = nn.LayerNorm(config.intermediate_size, eps=config.layer_norm_eps, elementwise_affine=False)
         self.down_proj = CastedLinearIn(config.intermediate_size, config.hidden_size, bias=False)
-        self.dropout = nn.Dropout(config.feed_forward_dropout_p)
-
-        self.initialize(config.hidden_size)
-
-    @torch.no_grad()
-    def initialize(self, hidden_size: int) -> None:
-        std: float = math.sqrt(2.0 / (5*hidden_size))
-
-        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_()
-
-    def up_project(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        hidden_layer = self.pre_norm(hidden_layer)
-        return self.up_proj(hidden_layer)
-
-    def activate(self, projection: torch.Tensor) -> torch.Tensor:
-        activated_projection: torch.Tensor
-
-        activated_projection = self.activation(projection)
-        activated_projection = self.inter_norm(activated_projection.float()).type_as(projection)
-
-        return activated_projection
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        
+    def forward(self, x: torch.Tensor):
+        x = self.pre_norm(x.float()).type_as(x)
+        x = self.up_proj(x)
+        x = self.activation(x)
+        x = self.inter_norm(x.float()).type_as(x)
+        x = self.down_proj(x)
+        x = self.dropout(x)
+        return x
 
-    def down_project(self, activated_projection: torch.Tensor) -> torch.Tensor:
-        output: torch.Tensor
 
-        output = self.down_proj(activated_projection)
+class Layer(nn.Module):
+    def __init__(self, config: GptBertConfig, layer_idx: int):
+        super().__init__()
 
-        return self.dropout(output)
+        self.attention = SelfAttention(config, layer_idx)
+        self.mlp = FeedForward(config)
+        self.lambdas = nn.Parameter(torch.tensor([0., 0., 1., 0., 1., 0.]))
 
-    def forward(self, hidden_layer: torch.Tensor) -> torch.Tensor:
-        output: torch.Tensor
+    def set_window_length(self, window_length: int):
+        self.attention.set_window_length(window_length)
 
-        output = self.up_project(hidden_layer)
-        output = self.activate(output)
-        output = self.down_project(output)
+    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)
 
-        return output
+        attention_output, v1 = 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)
 
+        return output, v1
 
-class RotaryPositionalEmbeddings(nn.Module):
 
-    def __init__(self, config, theta: int) -> None:
+class Encoder(nn.Module):
+    def __init__(self, config: GptBertConfig):
         super().__init__()
+        self.layers = nn.ModuleList([Layer(config, i) for i in range(config.num_layers)])
+        self.local_global_ratio = config.local_global_ratio
 
-        assert hasattr(config, "d_qk"), "The config must have a d_qk attribute!"
-        assert hasattr(config, "max_sequence_length"), "The config must have a max_sequence_length attribute!"
-
-        self.inv_freq: torch.Tensor
-        self.cos_matrix: torch.Tensor
-        self.sin_matrix: torch.Tensor
-        head_size: int
-        max_seq_len: int
-        inv_freq: torch.Tensor
-        pos: torch.Tensor
-        embedding: torch.Tensor
-
-        head_size = config.d_qk
-        assert head_size % 2 == 0
-        max_seq_len = config.max_sequence_length
+    def set_window_length(self, config: GptBertConfig):
+        for i, layer in enumerate(self.layers):
+            if (i + 1) % self.local_global_ratio == 0:
+                layer.set_window_length(config.global_window_length)
+            else:
+                layer.set_window_length(config.local_window_length)
 
-        inv_freq = 1.0 / (theta ** (torch.arange(0, head_size, 2, dtype=torch.float32) / head_size))
-        pos = torch.arange(max_seq_len, dtype=torch.float32)
-        embedding = torch.einsum('n, d -> nd', pos, inv_freq)
-        embedding = torch.cat([embedding, embedding], dim=-1).unsqueeze(0)
-        self.register_buffer("cos_matrix", embedding.cos(), persistent=False)
-        self.register_buffer("sin_matrix", embedding.sin(), persistent=False)
+    def forward(self, hidden_layer: torch.Tensor, padding_info, output_hidden_states=False, checkpoint_activations=False):
+        hidden_layers = [hidden_layer] if output_hidden_states else None
+        v1 = None
+        embeddings = hidden_layer
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        seq_len: int
-        cos_matrix: torch.Tensor
-        sin_matrix: torch.Tensor
-        x_rotate_half: torch.Tensor
-        out: torch.Tensor
+        for layer in self.layers:
+            if checkpoint_activations:
+                hidden_layer, v1 = torch.utils.checkpoint.checkpoint(layer, hidden_layer, embeddings, v1, padding_info, use_reentrant=True)
+            else:
+                hidden_layer, v1 = layer(hidden_layer, embeddings, v1, padding_info)
 
-        hidden_layer = x.float()
+            if output_hidden_states:
+                hidden_layers.append(hidden_layer)
 
-        seq_len = x.shape[2]
-
-        cos_matrix = self.cos_matrix[:, None, :seq_len, :]
-        sin_matrix = self.sin_matrix[:, None, :seq_len, :]
-
-        x_rotate_half = torch.cat(
-            [
-                -hidden_layer[:, :, :, x.size(-1) // 2:],
-                hidden_layer[:, :, :, :x.size(-1) // 2]
-            ],
-            dim=-1
-        )
-
-        out = hidden_layer * cos_matrix + x_rotate_half * sin_matrix
-        return out.type_as(x)
+        return hidden_layer, hidden_layers
 
 
 #
@@ -647,15 +563,15 @@ 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):
         std = math.sqrt(2.0 / (5.0 * self.hidden_size))
 
-        if isinstance(module, nn.Linear):
+        if isinstance(module, nn.Linear) or isinstance(module, CastedLinearIn):
             nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
             if module.bias is not None:
                 module.bias.data.zero_()
@@ -667,16 +583,17 @@ class GptBertPreTrainedModel(PreTrainedModel):
 
 
 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
 
         self.embedding = Embedding(config)
         self.encoder = Encoder(config)
-        self.classifier = MaskClassifier(config, self.embedding.word_embedding.weight) if add_mlm_layer else None
+        self.classifier = LMClassifier(config, config.vocab_size) if add_mlm_layer else None
         self.set_window_length(config)
+        self.gradient_checkpointing = False
+        self.post_init()
 
     def set_window_length(self, config) -> None:
         self.encoder.set_window_length(config)
@@ -690,8 +607,9 @@ class GptBertModel(GptBertPreTrainedModel):
     def get_contextualized_embeddings(
         self,
         input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None
-    ) -> List[torch.Tensor]:
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None
+    ):
         if input_ids is not None:
             input_shape = input_ids.size()
         else:
@@ -700,35 +618,55 @@ 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 attention_mask is None:
+            attention_mask = torch.ones(batch_size, seq_length, dtype=torch.bool, device=device)
+        else:
+            attention_mask = attention_mask.bool()
 
+        if is_flash_attn_2_available():
+            if len(attention_mask.size()) != 2:
+                raise ValueError("Bare `attention_mask` med to dimensjoner støttes nå for FlashAttention.")
+            with torch.no_grad():
+                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:
             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)
+            padding_info = attention_mask
+
+        static_embeddings = self.embedding(input_ids)
+
+        original_dtype = static_embeddings.dtype
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported() and static_embeddings.dtype == torch.float32:
+            static_embeddings = static_embeddings.bfloat16()
+
+        last_layer, contextualized_embeddings = self.encoder(
+            static_embeddings,
+            padding_info,
+            output_hidden_states=output_hidden_states,
+            checkpoint_activations=self.gradient_checkpointing and self.training
+        )
+
+        last_layer = last_layer.to(original_dtype)
+        if output_hidden_states:
+            contextualized_embeddings = [layer.to(original_dtype) for layer in contextualized_embeddings]
 
-            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)
+        # 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
 
-        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]
-        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))
-        ]
-        return last_layer, contextualized_embeddings, attention_probs
+        return last_layer, contextualized_embeddings
 
     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,
@@ -736,26 +674,24 @@ 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 = self.get_contextualized_embeddings(input_ids, attention_mask, output_hidden_states)
 
         if not return_dict:
             return (
                 sequence_output,
-                *([contextualized_embeddings] if output_hidden_states else []),
-                *([attention_probs] if output_attentions else [])
+                *([contextualized_embeddings] if output_hidden_states else [])
             )
 
         return BaseModelOutput(
             last_hidden_state=sequence_output,
-            hidden_states=contextualized_embeddings if output_hidden_states else None,
-            attentions=attention_probs if output_attentions else None
+            hidden_states=contextualized_embeddings if output_hidden_states else None
         )
 
 
 class GptBertForMaskedLM(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["head"]
+    _tied_weights_keys = ["classifier.emb2vocab.weight"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=True, **kwargs)
 
     def get_output_embeddings(self):
@@ -768,17 +704,14 @@ class GptBertForMaskedLM(GptBertModel):
         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,
         labels: Optional[torch.LongTensor] = None,
         **kwargs
     ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
         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 = self.get_contextualized_embeddings(input_ids, attention_mask, output_hidden_states)
         subword_prediction = self.classifier(sequence_output)
         subword_prediction = 30 * torch.sigmoid(subword_prediction / 7.5)
 
@@ -788,78 +721,28 @@ class GptBertForMaskedLM(GptBertModel):
             subword_prediction_flatten = subword_prediction[:, :-1].flatten(0, 1)
             masked_lm_loss = F.cross_entropy(subword_prediction_flatten, labels_flatten)
 
+        bos_logits = torch.zeros(subword_prediction.size(0), 1, self.config.vocab_size, dtype=subword_prediction.dtype, device=subword_prediction.device)
+        bos_logits[:, :, self.config.bos_token_id] = 1.0
+        subword_prediction = torch.cat([bos_logits, subword_prediction[:, :-1]], dim=1)
+
         if not return_dict:
             output = (
                 subword_prediction,
-                *([contextualized_embeddings] if output_hidden_states else []),
-                *([attention_probs] if output_attentions else [])
+                *([contextualized_embeddings] if output_hidden_states else [])
             )
             return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
 
         return MaskedLMOutput(
             loss=masked_lm_loss,
             logits=subword_prediction,
-            hidden_states=contextualized_embeddings if output_hidden_states else None,
-            attentions=attention_probs if output_attentions else None
+            hidden_states=contextualized_embeddings if output_hidden_states else None
         )
 
 
-class Classifier(nn.Module):
-    def __init__(self, config, num_labels: int):
-        super().__init__()
-
-        drop_out = getattr(config, "cls_dropout", None)
-        drop_out = config.hidden_dropout_prob if drop_out is None else drop_out
-
-        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 = CastedLinear(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 = CastedLinear(config.hidden_size, num_labels, bias=True)
-        self.dropout = nn.Dropout(drop_out)
-
-        self.initialize(config.hidden_size, config.intermediate_size, num_labels)
-
-    @torch.no_grad()
-    def initialize(self, hidden_size: int, intermediate_size: int, vocab_size: int) -> None:
-        proj_std: float = math.sqrt(2.0 / (hidden_size + intermediate_size))
-
-        nn.init.trunc_normal_(self.projection.weight, mean=0.0, std=proj_std, a=-2*proj_std, b=2*proj_std)
-        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
-
-
 class GptBertForCausalLM(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["head"]
+    _tied_weights_keys = ["classifier.emb2vocab.weight"]
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: GptBertConfig, **kwargs):
         config.is_decoder = True
         super().__init__(config, add_mlm_layer=True, **kwargs)
 
@@ -904,29 +787,27 @@ class GptBertForCausalLM(GptBertModel):
         assert past_key_values is None, "past_key_values is not supported for now"
         assert not use_cache, "use_cache is not supported for now"
 
-        sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
+        sequence_output, contextualized_embeddings = self.get_contextualized_embeddings(input_ids, attention_mask, output_hidden_states)
         subword_prediction = self.classifier(sequence_output)
         subword_prediction = 30 * torch.sigmoid(subword_prediction / 7.5)
 
-        masked_lm_loss = None
+        causal_lm_loss = None
         if labels is not None:
             labels_flatten = labels[:, 1:].flatten()
             subword_prediction_flatten = subword_prediction[:, :-1].flatten(0, 1)
-            masked_lm_loss = F.cross_entropy(subword_prediction_flatten, labels_flatten)
+            causal_lm_loss = F.cross_entropy(subword_prediction_flatten, labels_flatten)
 
         if not return_dict:
             output = (
                 subword_prediction,
-                *([contextualized_embeddings] if output_hidden_states else []),
-                *([attention_probs] if output_attentions else [])
+                *([contextualized_embeddings] if output_hidden_states else [])
             )
-            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+            return ((causal_lm_loss,) + output) if masked_lm_loss is not None else output
 
         return CausalLMOutput(
-            loss=masked_lm_loss,
+            loss=causal_lm_loss,
             logits=subword_prediction,
-            hidden_states=contextualized_embeddings if output_hidden_states else None,
-            attentions=attention_probs if output_attentions else None
+            hidden_states=contextualized_embeddings if output_hidden_states else None
         )
 
     def prepare_inputs_for_generation(
@@ -982,21 +863,20 @@ class GptBertForCausalLM(GptBertModel):
 
 
 class GptBertForSequenceClassification(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["classifier"]
+    _keys_to_ignore_on_load_missing = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
+    _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
-    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
-        self.head = Classifier(config, self.num_labels)
+        self.classifier = Classifier(config, self.num_labels)
+        self.post_init()
 
     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_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         labels: Optional[torch.LongTensor] = None,
@@ -1004,8 +884,8 @@ class GptBertForSequenceClassification(GptBertModel):
     ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
         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)
-        logits = self.head(sequence_output[:, 0, :])
+        sequence_output, contextualized_embeddings = self.get_contextualized_embeddings(input_ids, attention_mask, output_hidden_states)
+        logits = self.classifier(sequence_output[:, 0, :])
 
         loss = None
         if labels is not None:
@@ -1033,35 +913,32 @@ class GptBertForSequenceClassification(GptBertModel):
         if not return_dict:
             output = (
                 logits,
-                *([contextualized_embeddings] if output_hidden_states else []),
-                *([attention_probs] if output_attentions else [])
+                *([contextualized_embeddings] if output_hidden_states else [])
             )
             return ((loss,) + output) if loss is not None else output
 
         return SequenceClassifierOutput(
             loss=loss,
             logits=logits,
-            hidden_states=contextualized_embeddings if output_hidden_states else None,
-            attentions=attention_probs if output_attentions else None
+            hidden_states=contextualized_embeddings if output_hidden_states else None
         )
 
 
 class GptBertForTokenClassification(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["classifier"]
+    _keys_to_ignore_on_load_missing = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
+    _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
-    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
-        self.head = Classifier(config, self.num_labels)
+        self.classifier = Classifier(config, self.num_labels)
+        self.post_init()
 
     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_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         labels: Optional[torch.LongTensor] = None,
@@ -1069,8 +946,8 @@ class GptBertForTokenClassification(GptBertModel):
     ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
         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)
-        logits = self.head(sequence_output)
+        sequence_output, contextualized_embeddings = self.get_contextualized_embeddings(input_ids, attention_mask, output_hidden_states)
+        logits = self.classifier(sequence_output)
 
         loss = None
         if labels is not None:
@@ -1094,21 +971,20 @@ class GptBertForTokenClassification(GptBertModel):
 
 
 class GptBertForQuestionAnswering(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["classifier"]
-    
-    def __init__(self, config, **kwargs):
+    _keys_to_ignore_on_load_missing = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
+    _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
+
+    def __init__(self, config: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=False, **kwargs)
 
         self.num_labels = config.num_labels
-        self.head = Classifier(config, self.num_labels)
+        self.classifier = Classifier(config, self.num_labels)
+        self.post_init()
 
     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_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         start_positions: Optional[torch.Tensor] = None,
@@ -1117,8 +993,8 @@ class GptBertForQuestionAnswering(GptBertModel):
     ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
         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)
-        logits = self.head(sequence_output)
+        sequence_output, contextualized_embeddings = self.get_contextualized_embeddings(input_ids, attention_mask, output_hidden_states)
+        logits = self.classifier(sequence_output)
 
         start_logits, end_logits = logits.split(1, dim=-1)
         start_logits = start_logits.squeeze(-1).contiguous()
@@ -1146,8 +1022,7 @@ class GptBertForQuestionAnswering(GptBertModel):
             output = (
                 start_logits,
                 end_logits,
-                *([contextualized_embeddings] if output_hidden_states else []),
-                *([attention_probs] if output_attentions else [])
+                *([contextualized_embeddings] if output_hidden_states else [])
             )
             return ((total_loss,) + output) if total_loss is not None else output
 
@@ -1155,28 +1030,26 @@ class GptBertForQuestionAnswering(GptBertModel):
             loss=total_loss,
             start_logits=start_logits,
             end_logits=end_logits,
-            hidden_states=contextualized_embeddings if output_hidden_states else None,
-            attentions=attention_probs if output_attentions else None
+            hidden_states=contextualized_embeddings if output_hidden_states else None
         )
 
 
 class GptBertForMultipleChoice(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["classifier"]
+    _keys_to_ignore_on_load_missing = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
+    _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab.weight", "classifier.emb2vocab.bias"]
 
-    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)
-        self.head = Classifier(config, self.num_labels)
+        self.classifier = Classifier(config, self.num_labels)
+        self.post_init()
 
     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,
         labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         **kwargs
@@ -1187,8 +1060,8 @@ class GptBertForMultipleChoice(GptBertModel):
         flat_input_ids = input_ids.view(-1, input_ids.size(-1))
         flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
 
-        sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(flat_input_ids, flat_attention_mask)
-        logits = self.head(sequence_output)
+        sequence_output, contextualized_embeddings = self.get_contextualized_embeddings(flat_input_ids, flat_attention_mask, output_hidden_states)
+        logits = self.classifier(sequence_output)
         reshaped_logits = logits.view(-1, num_choices)
 
         loss = None
@@ -1199,14 +1072,12 @@ class GptBertForMultipleChoice(GptBertModel):
         if not return_dict:
             output = (
                 reshaped_logits,
-                *([contextualized_embeddings] if output_hidden_states else []),
-                *([attention_probs] if output_attentions else [])
+                *([contextualized_embeddings] if output_hidden_states else [])
             )
             return ((loss,) + output) if loss is not None else output
 
         return MultipleChoiceModelOutput(
             loss=loss,
             logits=reshaped_logits,
-            hidden_states=contextualized_embeddings if output_hidden_states else None,
-            attentions=attention_probs if output_attentions else None
+            hidden_states=contextualized_embeddings if output_hidden_states else None
         )