File size: 2,522 Bytes
3b2d368 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import torch
import torch.nn as nn
from torchtune.modules import RotaryPositionalEmbeddings
from torch.nn.attention.flex_attention import flex_attention
from torch.nn.attention import sdpa_kernel, SDPBackend
class Attention(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
# Linear layers...
self.W_q = nn.Linear(config.embed_dim, config.hidden_dim, bias=False)
self.W_k = nn.Linear(config.embed_dim, config.hidden_dim, bias=False)
self.W_v = nn.Linear(config.embed_dim, config.hidden_dim, bias=True)
self.W_o = nn.Linear(config.hidden_dim, config.embed_dim, bias=True)
# RoPE
d_k = config.hidden_dim // config.n_heads
self.rotary_embeddings = RotaryPositionalEmbeddings(d_k, max_seq_len=config.max_seq_len + 10)
self.drop_resid = nn.Dropout(0.1)
def forward(self, x):
B, S, E = x.shape
H = self.config.n_heads
D_h = self.config.hidden_dim // H
# 1. Project and Reshape (B, S, E) -> (B, S, H, D_h)
q = self.W_q(x).view(B, S, H, D_h)
k = self.W_k(x).view(B, S, H, D_h)
v = self.W_v(x).view(B, S, H, D_h)
# 2. Apply RoPE and Transpose to (B, H, S, D_h)
q = self.rotary_embeddings(q).transpose(1, 2).contiguous()
k = self.rotary_embeddings(k).transpose(1, 2).contiguous()
v = v.transpose(1, 2).contiguous() # V is often transposed without RoPE
# 3. Custom Bidirectional Masking Function (No-op)
def _score_mod(scores, b, h, i, j):
# This implements bidirectional attention by applying NO mask.
# All tokens are visible to all other tokens.
return scores
# 4. Attention with Explicit Backend Control
# sdpa_kernel controls the backend used for the attention operation within the block
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
attn_output = flex_attention(
q, k, v,
score_mod=_score_mod, # <-- Apply the custom bidirectional (no-op) mask
scale=None,
)
# 5. Combine Heads and Project Back
# (B, H, S, D_h) -> (B, S, H*D_h)
attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, self.config.hidden_dim)
attn_output = self.W_o(attn_output)
attn_output = self.drop_resid(attn_output)
return attn_output |