File size: 4,742 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchtune.modules import RotaryPositionalEmbeddings
from torch.nn.attention.flex_attention import flex_attention
from torch.nn.attention import sdpa_kernel, SDPBackend
class ICLAttention(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.W_q = nn.Linear(config.embed_dim_phi, config.hidden_dim_f, bias=False) # We set biases according to https://arxiv.org/pdf/2302.08626
self.W_k = nn.Linear(config.embed_dim_phi, config.hidden_dim_f, bias=False)
self.W_v = nn.Linear(config.embed_dim_f, config.hidden_dim_f, bias=True)
self.W_o = nn.Linear(config.hidden_dim_f, config.embed_dim_f, bias=True)
self.rotary_embeddings = RotaryPositionalEmbeddings(config.hidden_dim_f // config.n_heads_f, max_seq_len=config.max_seq_len + 10)
self.drop_resid = nn.Dropout(0.1)
def forward(self, q, k, v):
B, S, E = q.shape
q = self.W_q(q).view(B, S, self.config.n_heads_f, self.config.hidden_dim_f // self.config.n_heads_f)
k = self.W_k(k).view(B, S, self.config.n_heads_f, self.config.hidden_dim_f // self.config.n_heads_f)
v = self.W_v(v).view(B, S, self.config.n_heads_f, self.config.hidden_dim_f // self.config.n_heads_f).transpose(1, 2).contiguous()
q = self.rotary_embeddings(q).transpose(1, 2).contiguous()
k = self.rotary_embeddings(k).transpose(1, 2).contiguous()
# 🎯 关键修改点 1: 移除 score_mod，或将其设置为返回 scores 的恒等函数
# 为了实现双向（非因果）Attention，我们不应用任何掩码。
# 最简单的方法是在 flex_attention 调用中直接传入 score_mod=None
# 如果必须定义一个函数，它可以是：
# def _score_mod_bidirectional(scores, b, h, i, j):
# return scores
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
attn_output = flex_attention(
q, k, v,
# 传入 None 禁用因果掩码，实现双向/全连接 Attention
score_mod=None,
scale=None,
enable_gqa=False
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(B, S, self.config.hidden_dim_f)
attn_output = self.W_o(attn_output)
attn_output = self.drop_resid(attn_output)
return attn_output
# Vanilla Attention, but with dimensions properly aligned
class PhiAttention(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.W_q = nn.Linear(config.embed_dim_phi, config.hidden_dim_phi, bias=False)
self.W_k = nn.Linear(config.embed_dim_phi, config.hidden_dim_phi, bias=False)
self.W_v = nn.Linear(config.embed_dim_phi, config.hidden_dim_phi, bias=True)
self.W_o = nn.Linear(config.hidden_dim_phi, config.embed_dim_phi, bias=True)
self.rotary_embeddings = RotaryPositionalEmbeddings(config.hidden_dim_phi // config.n_heads_phi, max_seq_len=config.max_seq_len + 10)
self.drop_resid = nn.Dropout(0.1)
def forward(self, x):
B, S, E = x.shape
q = self.W_q(x).view(B, S, self.config.n_heads_phi, self.config.hidden_dim_phi // self.config.n_heads_phi)
k = self.W_k(x).view(B, S, self.config.n_heads_phi, self.config.hidden_dim_phi // self.config.n_heads_phi)
v = self.W_v(x).view(B, S, self.config.n_heads_phi, self.config.hidden_dim_phi // self.config.n_heads_phi).transpose(1, 2).contiguous()
q = self.rotary_embeddings(q).transpose(1, 2).contiguous()
k = self.rotary_embeddings(k).transpose(1, 2).contiguous()
# 🎯 关键修改点 2: 移除 score_mod，或将其设置为返回 scores 的恒等函数
# 原代码： keep = j <= i (因果掩码)
# 现代码： 直接在 flex_attention 中传入 score_mod=None
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
attn_output = flex_attention(
q, k, v,
# 传入 None 禁用因果掩码，实现双向/全连接 Attention
score_mod=None,
scale=None,
enable_gqa=False
)
attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, self.config.hidden_dim_phi)
attn_output = self.W_o(attn_output)
attn_output = self.drop_resid(attn_output)
return attn_output |