"""Sprout: a tiny GPT, all of it in one file. A decoder-only transformer in the style of today's open models: - pre-norm blocks with RMSNorm, - rotary position embeddings (RoPE), - multi-head causal self-attention with QK-norm, - a SwiGLU feed-forward layer, - input and output embeddings shared (weight tying), - no biases anywhere. """ import math from dataclasses import dataclass, asdict import torch import torch.nn as nn import torch.nn.functional as F @dataclass class Config: vocab_size: int = 8192 # how many different tokens the model knows context: int = 512 # how many tokens it can look back at n_layer: int = 8 # transformer blocks stacked on top of each other n_head: int = 8 # attention heads per block d_model: int = 512 # width of the residual stream d_ff: int = 1408 # hidden width of the feed-forward layer rope_base: float = 10000.0 def to_dict(self): return asdict(self) class RMSNorm(nn.Module): """Rescale a vector to unit root-mean-square, then apply a learned gain.""" def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): return F.rms_norm(x.float(), (x.size(-1),), self.weight.float(), self.eps).type_as(x) def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) # (B, T, C) -> (B, heads, T, head_dim) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y) class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x class GPT(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)) self.norm = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.head.weight = self.embed.weight # weight tying cos, sin = rope_tables(cfg.d_model // cfg.n_head, cfg.context, cfg.rope_base) self.register_buffer('cos', cos, persistent=False) self.register_buffer('sin', sin, persistent=False) self.apply(self._init) # residual projections start small so every block begins close to "do nothing" for name, p in self.named_parameters(): if name.endswith('proj.weight') or name.endswith('w2.weight'): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) @staticmethod def _init(m): if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, mean=0.0, std=0.02) def num_params(self): return sum(p.numel() for p in self.parameters()) def forward(self, idx, targets=None, loss_mask=None): T = idx.size(1) x = self.embed(idx) cos, sin = self.cos[:T], self.sin[:T] for block in self.blocks: x = block(x, cos, sin) logits = self.head(self.norm(x)) if targets is None: return logits, None loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), reduction='none') if loss_mask is None: return logits, loss.mean() mask = loss_mask.reshape(-1).float() return logits, (loss * mask).sum() / mask.sum().clamp(min=1) @torch.no_grad() def generate(self, idx, max_new, temperature=0.8, top_k=None, top_p=0.95, stop=None): for _ in range(max_new): logits, _ = self(idx[:, -self.cfg.context:]) logits = logits[:, -1, :].float() / max(temperature, 1e-5) if top_k: kth = torch.topk(logits, top_k).values[:, -1, None] logits[logits < kth] = -float('inf') probs = F.softmax(logits, dim=-1) if top_p and top_p < 1.0: sorted_p, order = probs.sort(descending=True) drop = sorted_p.cumsum(-1) - sorted_p > top_p sorted_p[drop] = 0 probs = torch.zeros_like(probs).scatter(-1, order, sorted_p) probs /= probs.sum(-1, keepdim=True) nxt = torch.multinomial(probs, 1) idx = torch.cat((idx, nxt), dim=1) if stop is not None and nxt.item() in stop: break return idx