| 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 |
| |
| |
| use_lora = getattr(config, 'use_lora_icl_attention', False) |
| |
| if use_lora: |
| from .lora import LoRALinear |
| lora_rank = getattr(config, 'lora_rank', 8) |
| lora_alpha = getattr(config, 'lora_alpha', 16) |
| lora_dropout = getattr(config, 'lora_dropout', 0.0) |
| |
| self.W_q = LoRALinear( |
| config.embed_dim_phi, config.hidden_dim_f, bias=False, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| self.W_k = LoRALinear( |
| config.embed_dim_phi, config.hidden_dim_f, bias=False, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| self.W_v = LoRALinear( |
| config.embed_dim_f, config.hidden_dim_f, bias=True, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| self.W_o = LoRALinear( |
| config.hidden_dim_f, config.embed_dim_f, bias=True, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| else: |
| self.W_q = nn.Linear(config.embed_dim_phi, config.hidden_dim_f, bias=False) |
| 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).transpose(1, 2).contiguous() |
| k = self.W_k(k).view(B, S, self.config.n_heads_f, self.config.hidden_dim_f // self.config.n_heads_f).transpose(1, 2).contiguous() |
| 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) |
| k = self.rotary_embeddings(k) |
| |
| def _score_mod(scores, b, h, i, j): |
| keep = (j < i) | ((i == 0) & (j == 0)) |
| return torch.where(keep, scores, torch.full_like(scores, float("-inf"))) |
|
|
| with sdpa_kernel(SDPBackend.FLASH_ATTENTION): |
| attn_output = flex_attention( |
| q, k, v, |
| score_mod=_score_mod, |
| 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 |
|
|
|
|
| class PhiAttention(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| |
| self.config = config |
| |
| |
| use_lora = getattr(config, 'use_lora_phi_attention', False) |
| |
| if use_lora: |
| from .lora import LoRALinear |
| lora_rank = getattr(config, 'lora_rank', 8) |
| lora_alpha = getattr(config, 'lora_alpha', 16) |
| lora_dropout = getattr(config, 'lora_dropout', 0.0) |
| |
| self.W_q = LoRALinear( |
| config.embed_dim_phi, config.hidden_dim_phi, bias=False, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| self.W_k = LoRALinear( |
| config.embed_dim_phi, config.hidden_dim_phi, bias=False, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| self.W_v = LoRALinear( |
| config.embed_dim_phi, config.hidden_dim_phi, bias=True, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| self.W_o = LoRALinear( |
| config.hidden_dim_phi, config.embed_dim_phi, bias=True, |
| lora_rank=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout |
| ) |
| else: |
| 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).transpose(1, 2).contiguous() |
| k = self.W_k(x).view(B, S, self.config.n_heads_phi, self.config.hidden_dim_phi // self.config.n_heads_phi).transpose(1, 2).contiguous() |
| 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) |
| k = self.rotary_embeddings(k) |
|
|
| def _score_mod(scores, b, h, i, j): |
| keep = j <= i |
| return torch.where(keep, scores, torch.full_like(scores, float("-inf"))) |
|
|
| with sdpa_kernel(SDPBackend.FLASH_ATTENTION): |
| attn_output = flex_attention( |
| q, k, v, |
| score_mod=_score_mod, |
| 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 |
