File size: 3,302 Bytes
cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 cd1ec5e 1a2dfe2 | 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 | # file: architecture.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class MultiHeadAttention(nn.Module):
"""Scaled dot-product self-attention with a fused QKV projection."""
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
# One matmul for Q, K and V is cheaper than three separate projections
self.qkv_projection = nn.Linear(d_model, 3 * d_model, bias=False)
self.out_projection = nn.Linear(d_model, d_model, bias=False)
def forward(self, x, is_causal=True):
batch_size, seq_len, d_model = x.shape
qkv = self.qkv_projection(x)
q, k, v = qkv.chunk(3, dim=-1)
# (B, T, d_model) -> (B, heads, T, d_k) so each head runs independently
q = q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
k = k.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
v = v.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
# Scale by sqrt(d_k) — keeps the scores from blowing up and the softmax stable
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
if is_causal:
# Block attention to future tokens — required for next-token prediction,
# otherwise the model would just read the answer.
mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)).view(1, 1, seq_len, seq_len)
scores = scores.masked_fill(mask == 0, float('-inf'))
attention_weights = F.softmax(scores, dim=-1)
context = torch.matmul(attention_weights, v)
# Merge heads back: (B, heads, T, d_k) -> (B, T, d_model)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
return self.out_projection(context)
class FeedForwardNetwork(nn.Module):
"""Per-token MLP: expand to d_ff, non-linearity, project back."""
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.dropout = nn.Dropout(dropout)
self.linear2 = nn.Linear(d_ff, d_model)
def forward(self, x):
return self.linear2(self.dropout(F.gelu(self.linear1(x))))
class TransformerBlock(nn.Module):
"""One decoder block: attention + feed-forward, pre-norm with residuals."""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
self.attention = MultiHeadAttention(d_model, n_heads)
self.feed_forward = FeedForwardNetwork(d_model, d_ff=4 * d_model, dropout=dropout) # 4x is the usual ratio
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, is_causal=True):
# Pre-norm (LayerNorm before each sub-layer) trains more stably than post-norm.
# The `x + ...` residuals let gradients flow and each layer refine rather than replace.
x = x + self.dropout(self.attention(self.ln1(x), is_causal=is_causal))
x = x + self.dropout(self.feed_forward(self.ln2(x)))
return x
|