fixes and optimizations

modeling_gptbert.py

@@ -3,14 +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, 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
+from transformers.utils import is_flash_attn_2_available, logging
 from transformers.modeling_outputs import (
     MaskedLMOutput,
     MultipleChoiceModelOutput,
@@ -23,6 +22,9 @@ 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:
@@ -31,13 +33,15 @@ try:
         from flash_attn.layers.rotary import RotaryEmbedding
         from flash_attn.ops.triton.rotary import apply_rotary
     else:
-        flash_attn_varlen_qkvpacked_func = None
-        RotaryEmbedding = object
-        apply_rotary = None
+        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:
-    flash_attn_varlen_qkvpacked_func = None
-    RotaryEmbedding = object
-    apply_rotary = None
+    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."
+    )
 
 
 # from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
@@ -89,34 +93,6 @@ class CastedLinearIn(nn.Linear):
         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 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)
-
-    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__()
@@ -138,77 +114,40 @@ class MultiCastedLinearOrthoIn(nn.Module):
         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)))
-
-    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)
         return x * gelu_new(gate)
 
 
-class MaskedSoftmax(torch.autograd.Function):
-    @staticmethod
-    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)
-        x.masked_fill_(mask, 0.0)
-        ctx.save_for_backward(x)
-        return x
-
-    @staticmethod
-    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: GptBertConfig):
         super().__init__()
-
         self.layers = nn.ModuleList([Layer(config, i) for i in range(config.num_layers)])
         self.short_long_ratio = config.short_long_ratio
 
     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)
+            if (i + 1) % self.local_global_ratio == 0:
+                layer.set_window_length(config.global_window_length)
             else:
-                layer.set_window_length(256)
+                layer.set_window_length(config.local_window_length)
 
-    def forward(self, hidden_layer: torch.Tensor, padding_info):
-        hidden_states = []
-        attention_probs = []
+    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
 
         for layer in self.layers:
-            hidden_layer, v1, attention_p = layer(hidden_layer, embeddings, v1, padding_info)
-            hidden_states.append(hidden_layer)
-            attention_probs.append(attention_p)
+            if checkpoint_activations:
+                hidden_layer, v1 = torch.utils.checkpoint.checkpoint(layers, hidden_layer, embeddings, v1, padding_info, use_reentrant=True)
+            else:
+                hidden_layer, v1 = layer(hidden_layer, embeddings, v1, padding_info)
 
-        return hidden_states, attention_probs
+            if output_hidden_states:
+                hidden_layers.append(hidden_layer)
+
+        return hidden_layer, hidden_layers
 
 
 class Layer(nn.Module):
@@ -227,12 +166,12 @@ class Layer(nn.Module):
         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, padding_info)
+        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, attention_p
+        return output, v1
 
 
 class Embedding(nn.Module):
@@ -257,21 +196,22 @@ class Embedding(nn.Module):
         return self.dropout(word_embedding)
 
 
-class MaskClassifier(nn.Module):
-    def __init__(self, config: GptBertConfig, embedding_weights: nn.Parameter):
+class Classifier(nn.Module):
+    def __init__(self, config: GptBertConfig, n_labels: int):
         super().__init__()
 
         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.emb2vocab = CastedLinearIn(config.hidden_size, n_labels, bias=True)
 
     def forward(self, x: torch.Tensor):
-        x = self.pre_norm(x)
+        x = self.pre_norm(x.float()).type_as(x)
         x = self.projection(x)
         x = gelu_new(x)
-        x = self.post_norm(x)
-        return self.emb2vocab(x)
+        x = self.post_norm(x.float()).type_as(x)
+        x = self.emb2vocab(x)
+        return x
 
 
 def flash_attention_forward(
@@ -354,14 +294,8 @@ class SelfAttention(nn.Module):
         theta = 160_000 if (layer_idx + 1) % config.short_long_ratio == 0 else 10_000
 
         # 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
-            )
+        if self.config._attn_implementation == "flash_attention_2":
+            self.rope_embedding = UnpaddedRotaryEmbedding(dim=config.d_qk, base=theta, max_seqlen=config.max_sequence_length)
         else:
             self.rope_embedding = RotaryPositionalEmbeddings(config, theta)
 
@@ -380,10 +314,10 @@ class SelfAttention(nn.Module):
         """Create and cache window attention mask."""
         if self.is_causal:
             mask = torch.ones(query_length, key_length, dtype=torch.bool, device=device)
-            mask = ~mask.tril().triu(diagonal=-self.window_length)
+            mask = mask.tril().triu(diagonal=-self.window_length)
         else:
             mask = torch.ones(query_length, key_length, dtype=torch.bool, device=device)
-            mask = ~mask.tril(diagonal=self.window_length).triu(diagonal=-self.window_length)
+            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]:
@@ -394,26 +328,24 @@ class SelfAttention(nn.Module):
         # 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
+                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)
-
-        attention_probabilities = MaskedSoftmax.apply(attention_scores, attention_mask, -1)
-        attention_probabilities = self.dropout(attention_probabilities)
-
-        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)
-
-        return value, attention_probabilities.detach()
+        output = F.scaled_dot_product_attention(
+            query=query,
+            key=key,
+            value=value,
+            attn_mask=attention_mask,
+            dropout_p=self.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() and isinstance(padding_info, tuple):
+        if self.config._attn_implementation == "flash_attention_2":
             # Unpadded case
             indices, cu_seqlens, max_seqlen = padding_info
             total_seqlen = hidden_layer.size(0)
@@ -421,16 +353,14 @@ class SelfAttention(nn.Module):
         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)
+        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)
 
         query, key = self.qk_proj(qk_layer).tensor_split([self.q_out_dim], dim=-1)
         value = self.v_proj(hidden_layer)
 
-        if is_flash_attn_2_available() and isinstance(padding_info, tuple):
+        if self.config._attn_implementation == "flash_attention_2":
             # 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)
@@ -445,15 +375,7 @@ class SelfAttention(nn.Module):
             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)
+            qkv = torch.stack([query, key, value], dim=1)  # (total_seqlen, 3, num_heads, head_dim)
 
             # Determine window size for local attention
             if self.window_length is not None and self.window_length > 0:
@@ -478,16 +400,15 @@ class SelfAttention(nn.Module):
 
             # Reshape output back
             output = output.view(total_seqlen, self.d_v * self.num_attention_heads)
-            attention_probabilities = None
 
         else:
             # Standard attention path
             query_length = hidden_layer.size(0)
             key_length = hidden_layer.size(0)
 
-            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)
+            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.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)
@@ -500,27 +421,15 @@ class SelfAttention(nn.Module):
             query = self.rope_embedding(query)
             key = self.rope_embedding(key)
 
-            # 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)
+            output = self.attention_operation(query, key, value, padding_info)
+            output = output.transpose(1, 2).flatten(2, 3)  # shape: [B, T, H*D]
 
-            output, attention_probabilities = self.attention_operation(query, key, value, padding_info)
-            output = output.permute(2, 0, 1, 3).flatten(2, 3)  # shape: [T, B, H*D]
-
-        output = self.inter_norm(output)
+        output = self.inter_norm(output.float()).type_as(output)
         output = self.out_proj(output)
+        output = self.dropout(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 output, v1
 
-        return self.dropout(output), v1, attention_probabilities
 
 class FeedForward(nn.Module):    
     def __init__(self, config: GptBertConfig):
@@ -533,12 +442,13 @@ class FeedForward(nn.Module):
         self.dropout = nn.Dropout(config.feed_forward_dropout_p)
         
     def forward(self, x: torch.Tensor):
-        x = self.pre_norm(x)
+        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)
-        return self.dropout(x)
+        x = self.dropout(x)
+        return x
 
 
 class ApplyRotaryEmbUnpad(torch.autograd.Function):
@@ -596,23 +506,17 @@ class ApplyRotaryEmbUnpad(torch.autograd.Function):
         return do, None, None, None, None, None, None
 
 
-def apply_rotary_unpadded(
-    qkv,
-    cos,
-    sin,
-    cu_seqlens: Optional[torch.Tensor] = None,
-    max_seqlen: Optional[int] = None,
-):
+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)
 
 
 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)
+    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
 
         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)
+            self._update_cos_sin_cache(max_seqlen, device=device, dtype=None)
 
     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:
@@ -696,7 +600,7 @@ class GptBertPreTrainedModel(PreTrainedModel):
     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_()
@@ -706,6 +610,33 @@ class GptBertPreTrainedModel(PreTrainedModel):
             module.bias.data.zero_()
             module.weight.data.fill_(1.0)
 
+    @classmethod
+    def _autoset_attn_implementation(
+        cls,
+        config,
+        torch_dtype: Optional[torch.dtype] = None,
+        device_map: Optional[Union[str, Dict[str, int]]] = None,
+        check_device_map: bool = True,
+    ):
+        if config._attn_implementation_internal is None:
+            config._attn_implementation_internal = "flash_attention_2"
+            try:
+                return cls._check_and_enable_flash_attn_2(
+                    config,
+                    torch_dtype=torch.float16,
+                    device_map=device_map,
+                    hard_check_only=False,
+                    check_device_map=check_device_map,
+                )
+            except (ValueError, ImportError):
+                config._attn_implementation_internal = None
+        return super()._autoset_attn_implementation(
+            config,
+            torch_dtype=torch_dtype,
+            device_map=device_map,
+            check_device_map=check_device_map,
+        )
+
 
 class GptBertModel(GptBertPreTrainedModel):
     def __init__(self, config: GptBertConfig, add_mlm_layer=False, **kwargs):
@@ -715,8 +646,10 @@ class GptBertModel(GptBertPreTrainedModel):
 
         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 = Classifier(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)
@@ -732,7 +665,7 @@ class GptBertModel(GptBertPreTrainedModel):
         input_ids: 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()
         else:
@@ -741,39 +674,50 @@ class GptBertModel(GptBertPreTrainedModel):
         batch_size, seq_length = input_shape
         device = input_ids.device
 
-        if is_flash_attn_2_available():
-            if attention_mask is None:
-                attention_mask = torch.ones(batch_size, seq_length, dtype=torch.bool, device=device)
-            elif attention_mask is not None and len(attention_mask.size()) != 2:
-                raise ValueError("Only attention mask with two dimensions is supported now.")
-            input_ids, indices, cu_seqlens, max_seqlen_in_batch = _unpad_input(input_ids, attention_mask)
+        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 self.config._attn_implementation == "flash_attention_2":
+            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 attention_mask is None:
-                attention_mask = torch.ones(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 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(1)
+            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)
-        contextualized_embeddings, attention_probs = self.encoder(static_embeddings, padding_info)
-        last_layer = contextualized_embeddings[-1]
+
+        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]
 
         # Pad output if using FlashAttention
-        if is_flash_attn_2_available():
+        if self.config._attn_implementation == "flash_attention_2":
             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
+        return last_layer, contextualized_embeddings
 
     def forward(
         self,
@@ -786,26 +730,22 @@ 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, output_hidden_states
-        )
+        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: GptBertConfig, **kwargs):
         super().__init__(config, add_mlm_layer=True, **kwargs)
@@ -820,17 +760,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)
 
@@ -847,60 +784,19 @@ class GptBertForMaskedLM(GptBertModel):
         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: GptBertConfig, 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 = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
-        # 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 = nn.Linear(config.hidden_size, num_labels, bias=True)
-        # 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 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"]
+    _tied_weights_keys = ["classifier.emb2vocab.weight"]
 
     def __init__(self, config: GptBertConfig, **kwargs):
         config.is_decoder = True
@@ -947,29 +843,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(
@@ -1025,21 +919,19 @@ class GptBertForCausalLM(GptBertModel):
 
 
 class GptBertForSequenceClassification(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["classifier"]
+    _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab"]
 
     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,
@@ -1047,8 +939,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:
@@ -1076,35 +968,31 @@ 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_unexpected = ["classifier.emb2vocab"]
 
     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,
@@ -1112,8 +1000,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:
@@ -1137,21 +1025,19 @@ class GptBertForTokenClassification(GptBertModel):
 
 
 class GptBertForQuestionAnswering(GptBertModel):
-    _keys_to_ignore_on_load_unexpected = ["classifier"]
+    _keys_to_ignore_on_load_unexpected = ["classifier.emb2vocab"]
 
     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,
@@ -1160,8 +1046,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()
@@ -1189,8 +1075,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
 
@@ -1198,28 +1083,25 @@ 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_unexpected = ["classifier.emb2vocab"]
 
     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
@@ -1230,8 +1112,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
@@ -1242,14 +1124,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
         )