| 51 | |
| 52 | |
| 53 | class SelfAttention(nn.Module): |
| 54 | |
| 55 | def __init__(self, |
| 56 | dim, |
| 57 | num_heads, |
| 58 | causal=False, |
| 59 | attn_dropout=0.0, |
| 60 | proj_dropout=0.0): |
| 61 | assert dim % num_heads == 0 |
| 62 | super().__init__() |
| 63 | self.dim = dim |
| 64 | self.num_heads = num_heads |
| 65 | self.head_dim = dim // num_heads |
| 66 | self.causal = causal |
| 67 | self.attn_dropout = attn_dropout |
| 68 | self.proj_dropout = proj_dropout |
| 69 | |
| 70 | # layers |
| 71 | self.to_qkv = nn.Linear(dim, dim * 3) |
| 72 | self.proj = nn.Linear(dim, dim) |
| 73 | |
| 74 | def forward(self, x): |
| 75 | """ |
| 76 | x: [B, L, C]. |
| 77 | """ |
| 78 | b, s, c, n, d = *x.size(), self.num_heads, self.head_dim |
| 79 | |
| 80 | # compute query, key, value |
| 81 | q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2) |
| 82 | |
| 83 | # compute attention |
| 84 | p = self.attn_dropout if self.training else 0.0 |
| 85 | x = flash_attention(q, k, v, dropout_p=p, causal=self.causal, version=2) |
| 86 | x = x.reshape(b, s, c) |
| 87 | |
| 88 | # output |
| 89 | x = self.proj(x) |
| 90 | x = F.dropout(x, self.proj_dropout, self.training) |
| 91 | return x |
| 92 | |
| 93 | |
| 94 | class SwiGLU(nn.Module): |