removed SDPA

modeling_gptbert.py

@@ -3,6 +3,7 @@ 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
 
@@ -37,17 +38,11 @@ try:
         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."
         )
-        torch.backends.cuda.enable_flash_sdp(False)
-        torch.backends.cuda.enable_mem_efficient_sdp(False)
-        torch.backends.cuda.enable_math_sdp(True)
 except ImportError:
     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."
     )
-    torch.backends.cuda.enable_flash_sdp(False)
-    torch.backends.cuda.enable_mem_efficient_sdp(False)
-    torch.backends.cuda.enable_math_sdp(True)
 
 
 # from https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modeling_modernbert.py
@@ -318,6 +313,25 @@ class RotaryPositionalEmbeddings(nn.Module):
         return out.type_as(x)
 
 
+class MaskedSoftmax(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, mask: torch.BoolTensor, dim: int) -> torch.Tensor:
+        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
+
+
 class SelfAttention(nn.Module):
     def __init__(self, config: GptBertConfig, layer_idx: int):
         super().__init__()
@@ -347,6 +361,7 @@ class SelfAttention(nn.Module):
         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))
 
+        self.attention_dropout = nn.Dropout(config.attention_dropout)
         self.dropout = nn.Dropout(config.hidden_dropout)
 
         theta = 160_000 if (layer_idx + 1) % config.local_global_ratio == 0 else 10_000
@@ -390,14 +405,15 @@ class SelfAttention(nn.Module):
             else:
                 attention_mask = window_mask
 
-        output = F.scaled_dot_product_attention(
-            query=query,
-            key=key,
-            value=value,
-            attn_mask=attention_mask if not self.is_causal else None,
-            dropout_p=self.config.attention_dropout if self.training else 0.0,
-            is_causal=self.is_causal
-        )
+        attention_scores = torch.bmm(query.flatten(0, 1), key.transpose(-1, -2).flatten(0, 1)) * self.scale  # shape: [B*H, Q_T, K_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.attention_dropout(attention_probabilities)
+
+        output = torch.bmm(attention_probabilities.flatten(0, 1), value.flatten(0, 1))
+        output = output.view(batch_size, self.num_attention_heads, query_length, self.d_v)
+
         return output
 
     def forward(self, hidden_layer: torch.Tensor, qk_layer: torch.Tensor, v1: torch.Tensor | None, padding_info):