| """Sprout: a small GPT trained from scratch on TinyStories. |
| |
| Modern-but-small transformer: RoPE, RMSNorm (pre-norm), SwiGLU MLP, |
| flash attention via F.scaled_dot_product_attention, untied heads. |
| """ |
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def rope_cache(seq_len, head_dim, device, base=10000.0): |
| """Precompute cos/sin tables for rotary embeddings [seq_len, head_dim/2].""" |
| inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) |
| t = torch.arange(seq_len, device=device).float() |
| freqs = torch.outer(t, inv_freq) |
| return torch.cos(freqs), torch.sin(freqs) |
|
|
|
|
| def apply_rope(x, cos, sin): |
| """x: [batch, heads, seq, head_dim]; cos/sin: [seq, head_dim/2].""" |
| half = x.shape[-1] // 2 |
| x1, x2 = x[..., :half], x[..., half:] |
| cos = cos[None, None, :, :].to(x.dtype) |
| sin = sin[None, None, :, :].to(x.dtype) |
| |
| out1 = x1 * cos - x2 * sin |
| out2 = x2 * cos + x1 * sin |
| return torch.cat([out1, out2], dim=-1) |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, dim, n_heads, dropout=0.0): |
| super().__init__() |
| assert dim % n_heads == 0 |
| self.n_heads = n_heads |
| self.head_dim = dim // n_heads |
| self.qkv = nn.Linear(dim, 3 * dim, bias=False) |
| self.proj = nn.Linear(dim, dim, bias=False) |
| self.dropout = dropout |
|
|
| def forward(self, x, cos, sin): |
| B, T, C = x.shape |
| q, k, v = self.qkv(x).split(C, dim=2) |
| shape = (B, T, self.n_heads, self.head_dim) |
| q = q.view(shape).transpose(1, 2) |
| k = k.view(shape).transpose(1, 2) |
| v = v.view(shape).transpose(1, 2) |
| q = apply_rope(q, cos, sin) |
| k = apply_rope(k, cos, sin) |
| y = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.proj(y) |
|
|
|
|
| class SwiGLU(nn.Module): |
| """SwiGLU MLP; hidden dim set to ~8/3 of model dim rounded to multiple of 256.""" |
|
|
| def __init__(self, dim, dropout=0.0): |
| super().__init__() |
| hidden = max(256, int(dim * 8 / 3 / 256 + 0.5) * 256) |
| self.w_gate = nn.Linear(dim, hidden, bias=False) |
| self.w_up = nn.Linear(dim, hidden, bias=False) |
| self.w_down = nn.Linear(hidden, dim, bias=False) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| return self.drop(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, dim, n_heads, dropout=0.0): |
| super().__init__() |
| self.n1 = nn.RMSNorm(dim, eps=1e-5) |
| self.attn = CausalSelfAttention(dim, n_heads, dropout) |
| self.n2 = nn.RMSNorm(dim, eps=1e-5) |
| self.mlp = SwiGLU(dim, dropout) |
|
|
| def forward(self, x, cos, sin): |
| x = x + self.attn(self.n1(x), cos, sin) |
| x = x + self.mlp(self.n2(x)) |
| return x |
|
|
|
|
| class Sprout(nn.Module): |
| def __init__(self, vocab_size, dim=512, n_layers=8, n_heads=8, |
| block_size=1024, dropout=0.0): |
| super().__init__() |
| self.block_size = block_size |
| self.tok_emb = nn.Embedding(vocab_size, dim) |
| self.drop = nn.Dropout(dropout) |
| self.blocks = nn.ModuleList(Block(dim, n_heads, dropout) for _ in range(n_layers)) |
| self.norm_f = nn.RMSNorm(dim, eps=1e-5) |
| self.head = nn.Linear(dim, vocab_size, bias=False) |
| |
| self.apply(self._init_weights) |
| for pn, p in self.named_parameters(): |
| if pn.endswith("proj.weight") or pn.endswith("w_down.weight"): |
| nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layers)) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, nn.Linear): |
| nn.init.normal_(m.weight, mean=0.0, std=0.02) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| elif isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, mean=0.0, std=0.02) |
|
|
| def forward(self, idx, targets=None): |
| B, T = idx.shape |
| assert T <= self.block_size, f"sequence {T} exceeds block size {self.block_size}" |
| cos, sin = rope_cache(T, self.blocks[0].attn.head_dim, idx.device) |
| x = self.drop(self.tok_emb(idx)) |
| for block in self.blocks: |
| x = block(x, cos, sin) |
| x = self.norm_f(x) |
| if targets is not None: |
| logits = self.head(x) |
| return F.cross_entropy( |
| logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1 |
| ) |
| |
| return self.head(x[:, [-1], :]) |
|
|
