"""A modern decoder-only transformer, from scratch, readable. The same building blocks as Llama / DeepSeek — at mini scale: - RoPE : tells attention WHERE each token sits (rotary positions) - RMSNorm : keeps numbers stable between blocks (cheap normalization) - SwiGLU : the "thinking" feed-forward block, gated for quality - weight tying: input embedding and output head share one matrix Contract R2: every block commented in plain English. """ import math from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F @dataclass class ModelConfig: vocab_size: int = 4096 dim: int = 128 # width of every token's vector n_layers: int = 4 # how many attention+thinking floors n_heads: int = 4 # parallel attention "viewpoints" max_seq_len: int = 256 # longest context the model can see dropout: float = 0.0 # tiny data? keep 0; large runs may use 0.1 class RMSNorm(nn.Module): """Normalize a vector to a stable length, then let a learned gain rescale. Why: after many layers, numbers drift huge or tiny; training explodes. RMSNorm is LayerNorm's cheaper cousin (no mean subtraction) — the choice of Llama/DeepSeek. """ def __init__(self, dim, eps=1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return x * rms * self.weight def build_rope_cache(dim, max_seq_len, base=10000.0): """Precompute rotation angles for RoPE. Idea in one line: rotate each pair of vector coordinates by an angle that depends on the token's POSITION — attention then "feels" distance between words without any extra position embeddings. """ half = dim // 2 freqs = 1.0 / (base ** (torch.arange(0, half).float() / half)) t = torch.arange(max_seq_len).float() angles = torch.outer(t, freqs) # [seq, half] return torch.cos(angles), torch.sin(angles) def apply_rope(x, cos, sin): """Rotate query/key pairs. x: [batch, heads, seq, head_dim].""" seq = x.shape[2] cos = cos[:seq].view(1, 1, seq, -1) sin = sin[:seq].view(1, 1, seq, -1) x1, x2 = x.chunk(2, dim=-1) # split coords into pairs return torch.cat([x1 * cos - x2 * sin, # classic 2D rotation x2 * cos + x1 * sin], dim=-1) class Attention(nn.Module): """Causal self-attention: every token looks at all PREVIOUS tokens and decides whom to listen to. "Causal" = no peeking at the future, because the model's job is to predict it. """ def __init__(self, cfg: ModelConfig): super().__init__() assert cfg.dim % cfg.n_heads == 0 self.n_heads = cfg.n_heads self.head_dim = cfg.dim // cfg.n_heads self.wq = nn.Linear(cfg.dim, cfg.dim, bias=False) self.wk = nn.Linear(cfg.dim, cfg.dim, bias=False) self.wv = nn.Linear(cfg.dim, cfg.dim, bias=False) self.wo = nn.Linear(cfg.dim, cfg.dim, bias=False) self.dropout = cfg.dropout def forward(self, x, cos, sin): B, T, C = x.shape # each token asks a Question, offers a Key, carries a Value q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # fused attention with the causal (no-future) mask 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.wo(y) class SwiGLU(nn.Module): """The "thinking" block. Gated: one path proposes, another path decides how much of it to let through. Empirically beats plain ReLU MLPs — used by Llama, PaLM, DeepSeek. """ def __init__(self, cfg: ModelConfig): super().__init__() hidden = int(8 * cfg.dim / 3 / 64) * 64 or 2 * cfg.dim # round to 64 self.w_gate = nn.Linear(cfg.dim, hidden, bias=False) self.w_up = nn.Linear(cfg.dim, hidden, bias=False) self.w_down = nn.Linear(hidden, cfg.dim, bias=False) def forward(self, x): return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) class Block(nn.Module): """One floor of the tower: attention (talk to context), then SwiGLU (think it over). Residual "+x" paths let information skip floors — that is what makes deep towers trainable. """ def __init__(self, cfg: ModelConfig): super().__init__() self.norm1 = RMSNorm(cfg.dim) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.dim) self.ffn = SwiGLU(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) x = x + self.ffn(self.norm2(x)) return x class TinyLLM(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.cfg = cfg self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.dim) self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers)) self.norm_out = RMSNorm(cfg.dim) self.head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False) # weight tying: reading and writing tokens share one matrix — # fewer parameters, better tiny-model quality self.head.weight = self.tok_emb.weight cos, sin = build_rope_cache(cfg.dim // cfg.n_heads, cfg.max_seq_len) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init) @staticmethod def _init(module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): x = self.tok_emb(idx) # tokens -> vectors for block in self.blocks: # up the tower x = block(x, self.rope_cos, self.rope_sin) x = self.norm_out(x) logits = self.head(x) # vectors -> vocab scores loss = None if targets is not None: # cross-entropy: how surprised was the model by the true next token loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=50): """Write text: predict, sample, append, repeat.""" self.eval() for _ in range(max_new_tokens): ctx = idx[:, -self.cfg.max_seq_len:] logits, _ = self(ctx) logits = logits[:, -1, :] / max(temperature, 1e-6) if top_k: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") probs = F.softmax(logits, dim=-1) nxt = torch.multinomial(probs, num_samples=1) idx = torch.cat([idx, nxt], dim=1) return idx def num_params(self): return sum(p.numel() for p in self.parameters()) if __name__ == "__main__": # self-check: forward + backward on random tokens, and a generation call cfg = ModelConfig(vocab_size=4096, dim=128, n_layers=4, n_heads=4) model = TinyLLM(cfg) print(f"model built: {model.num_params():,} parameters") x = torch.randint(0, cfg.vocab_size, (2, 32)) y = torch.randint(0, cfg.vocab_size, (2, 32)) _, loss = model(x, y) loss.backward() print(f"forward+backward OK, random-guess loss = {loss.item():.3f} " f"(expected ~{math.log(cfg.vocab_size):.3f} = ln(vocab))") out = model.generate(x[:1, :4], max_new_tokens=8) print(f"generate OK: {out.shape[1]} tokens")