"""A GPT-style decoder-only transformer, written from scratch in PyTorch. Architecture: token embeddings -> N x [RMSNorm -> causal self-attention (rotary position embeddings) -> RMSNorm -> SwiGLU MLP] -> RMSNorm -> LM head (weights tied to the embeddings). Generation uses a KV cache. Version 2 adds grouped-query attention (fewer key/value heads: a smaller KV cache and faster generation) and QK-norm (normalised queries/keys, which keeps attention logits bounded and allows higher learning rates). Version-1 checkpoints load unchanged: the new fields default to v1 behaviour. """ from __future__ import annotations import math from dataclasses import asdict, dataclass import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint @dataclass class GPTConfig: vocab_size: int = 4096 block_size: int = 512 n_layer: int = 6 n_head: int = 6 n_embd: int = 384 dropout: float = 0.0 n_kv_head: int | None = None # v2: grouped-query attention (None = one KV head per query head) qk_norm: bool = False # v2: RMSNorm on queries and keys arch_version: int = 1 @property def kv_heads(self) -> int: return self.n_kv_head or self.n_head def to_dict(self) -> dict: return asdict(self) _V2 = dict(qk_norm=True, arch_version=2) # Named sizes. Parameter counts are for an 8k vocabulary. PRESETS: dict[str, dict] = { # version 1 "tiny": dict(n_layer=2, n_head=2, n_embd=64, block_size=128), # ~0.6M, for tests "small": dict(n_layer=4, n_head=4, n_embd=256, block_size=256), # ~5M, CPU-friendly "base": dict(n_layer=6, n_head=6, n_embd=384, block_size=512), # ~14M "medium": dict(n_layer=8, n_head=8, n_embd=512, block_size=1024), # ~30M, wants a GPU "large": dict(n_layer=12, n_head=12, n_embd=768, block_size=1024), # ~92M, GPU # version 2 (GQA + QK-norm) "v2-tiny": dict(n_layer=2, n_head=4, n_kv_head=2, n_embd=64, block_size=128, **_V2), "v2-small": dict(n_layer=8, n_head=8, n_kv_head=2, n_embd=256, block_size=512, **_V2), "v2-base": dict(n_layer=12, n_head=8, n_kv_head=2, n_embd=512, block_size=1024, **_V2), "v2-medium": dict(n_layer=16, n_head=12, n_kv_head=4, n_embd=768, block_size=1024, **_V2), "v2-large": dict(n_layer=24, n_head=16, n_kv_head=4, n_embd=1024, block_size=2048, **_V2), } # version 3: the ~40M-parameter model (v2 architecture, scaled up, trained with Muon) _V3 = dict(qk_norm=True, arch_version=3) PRESETS.update({ "v3-tiny": dict(n_layer=3, n_head=4, n_kv_head=2, n_embd=64, block_size=128, **_V3), "v3-40m": dict(n_layer=13, n_head=8, n_kv_head=2, n_embd=512, block_size=1024, **_V3), }) # version 4: ~100M parameters. Deep and narrow (34 x 512), which suits small models and lets # v4 be grown from a trained v3 (13 x 512) without changing its width. _V4 = dict(qk_norm=True, arch_version=4) PRESETS.update({ "v4-tiny": dict(n_layer=5, n_head=4, n_kv_head=2, n_embd=64, block_size=128, **_V4), "v4-100m": dict(n_layer=34, n_head=8, n_kv_head=2, n_embd=512, block_size=1024, **_V4), }) # version 5: grown from v4 (34 x 512) in width only, so the depth stays 34. # v5-385m 34 x 1024 (16 q / 4 kv heads) ~385M -> free Kaggle / Colab GPUs (T4) # v5-1.5b 34 x 2048 (32 q / 8 kv heads) ~1.5B -> a rented A100 / H100 _V5 = dict(qk_norm=True, arch_version=5) PRESETS.update({ "v5-tiny": dict(n_layer=6, n_head=8, n_kv_head=4, n_embd=128, block_size=128, **_V5), # grows from v4-tiny "v5-385m": dict(n_layer=34, n_head=16, n_kv_head=4, n_embd=1024, block_size=1024, **_V5), "v5-1.5b": dict(n_layer=34, n_head=32, n_kv_head=8, n_embd=2048, block_size=1024, **_V5), }) DEFAULT_PRESET = "v2-small" class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6) -> None: super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() return (x.float() * norm).type_as(x) * self.weight def rope_tables(head_dim: int, length: int, device=None, base: float = 10000.0): inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) t = torch.arange(length, device=device).float() freqs = torch.outer(t, inv_freq) # (T, hd/2) return freqs.cos(), freqs.sin() def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: # x: (B, H, T, hd); cos/sin: (T, hd/2) x1, x2 = x[..., 0::2], x[..., 1::2] cos, sin = cos[None, None], sin[None, None] out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) return out.flatten(-2).type_as(x) class Attention(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() assert cfg.n_embd % cfg.n_head == 0, "n_embd must be divisible by n_head" assert cfg.n_head % cfg.kv_heads == 0, "n_head must be a multiple of n_kv_head" self.n_head = cfg.n_head self.n_kv = cfg.kv_heads self.head_dim = cfg.n_embd // cfg.n_head assert self.head_dim % 2 == 0, "head dimension must be even for rotary embeddings" self.kv_dim = self.n_kv * self.head_dim # One fused projection; with n_kv == n_head this is exactly the v1 layout. self.qkv = nn.Linear(cfg.n_embd, cfg.n_embd + 2 * self.kv_dim, bias=False) self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False) if cfg.qk_norm: self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) else: self.q_norm = self.k_norm = None self.dropout = cfg.dropout def _sdpa(self, q, k, v, **kwargs): if self.n_kv != self.n_head: repeat = self.n_head // self.n_kv k = k.repeat_interleave(repeat, dim=1) v = v.repeat_interleave(repeat, dim=1) return F.scaled_dot_product_attention(q, k, v, **kwargs) def forward(self, x, cos, sin, cache=None): B, T, C = x.shape q, k, v = self.qkv(x).split([C, self.kv_dim, self.kv_dim], dim=2) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_kv, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_kv, self.head_dim).transpose(1, 2) if self.q_norm is not None: q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) if cache is not None: if cache.get("k") is not None: k = torch.cat([cache["k"], k], dim=2) v = torch.cat([cache["v"], v], dim=2) cache["k"], cache["v"] = k, v # cached at n_kv heads: GQA keeps the cache small # Causal masking: with a cache, the T new queries see all past keys. if cache is not None and k.size(2) != T: if T == 1: y = self._sdpa(q, k, v) else: past = k.size(2) - T mask = torch.ones(T, k.size(2), dtype=torch.bool, device=x.device).tril(diagonal=past) y = self._sdpa(q, k, v, attn_mask=mask) else: y = self._sdpa(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 MLP(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() hidden = int(8 * cfg.n_embd / 3) hidden = 64 * ((hidden + 63) // 64) self.gate = nn.Linear(cfg.n_embd, hidden, bias=False) self.up = nn.Linear(cfg.n_embd, hidden, bias=False) self.down = nn.Linear(hidden, cfg.n_embd, bias=False) self.drop = nn.Dropout(cfg.dropout) def forward(self, x): return self.drop(self.down(F.silu(self.gate(x)) * self.up(x))) class Block(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() self.norm1 = RMSNorm(cfg.n_embd) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.n_embd) self.mlp = MLP(cfg) def forward(self, x, cos, sin, cache=None): x = x + self.attn(self.norm1(x), cos, sin, cache) return x + self.mlp(self.norm2(x)) class GPT(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() self.cfg = cfg self.grad_checkpoint = False self.embed = nn.Embedding(cfg.vocab_size, cfg.n_embd) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)) self.norm = RMSNorm(cfg.n_embd) self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) self.head.weight = self.embed.weight # weight tying cos, sin = rope_tables(cfg.n_embd // cfg.n_head, cfg.block_size) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init_weights) for name, p in self.named_parameters(): if name.endswith("proj.weight") or name.endswith("down.weight"): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) @staticmethod def _init_weights(module: nn.Module) -> None: 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 num_params(self) -> int: return sum(p.numel() for p in self.parameters()) def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None, loss_mask: torch.Tensor | None = None, caches=None, start_pos: int = 0): B, T = idx.shape if start_pos + T > self.cfg.block_size: raise ValueError(f"sequence length {start_pos + T} exceeds block_size {self.cfg.block_size}") cos = self.rope_cos[start_pos: start_pos + T] sin = self.rope_sin[start_pos: start_pos + T] # fp32 residual stream, even when sharded training gathers the weights in fp16/bf16. x = self.drop(self.embed(idx).float()) checkpointing = self.grad_checkpoint and self.training and caches is None for i, block in enumerate(self.blocks): if checkpointing: # Recompute this block's activations in the backward pass instead of storing them: # far less memory (fits bigger models/batches on a 16 GB GPU) for ~30% more compute. x = torch.utils.checkpoint.checkpoint(block, x, cos, sin, None, use_reentrant=False) else: x = block(x, cos, sin, caches[i] if caches is not None else None) x = self.norm(x) if targets is None: return self.head(x[:, -1:, :]), None logits = self.head(x) losses = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1), reduction="none") if loss_mask is not None: mask = loss_mask.reshape(-1).float() loss = (losses * mask).sum() / mask.sum().clamp(min=1.0) else: loss = losses.mean() return logits, loss @torch.no_grad() def generate(self, idx: torch.Tensor, max_new_tokens: int, *, temperature: float = 0.8, top_k: int | None = 50, top_p: float | None = 0.95, stop_ids: set[int] | None = None, on_token=None, no_repeat_ngram: int = 0) -> torch.Tensor: """Sample tokens after ``idx`` (shape (1, T)). Uses a KV cache. ``no_repeat_ngram=n`` forbids any token that would repeat an n-token span already generated; this breaks degenerate loops without penalising ordinary repetition such as indentation or a variable name used twice.""" self.eval() seen: dict[tuple[int, ...], set[int]] = {} generated: list[int] = [] stop_ids = stop_ids or set() # Keep the most recent context that leaves room for generation. budget = self.cfg.block_size - max_new_tokens if budget < 1: max_new_tokens = self.cfg.block_size // 2 budget = self.cfg.block_size - max_new_tokens idx = idx[:, -budget:] caches = [dict() for _ in self.blocks] logits, _ = self(idx, caches=caches, start_pos=0) pos = idx.size(1) out = idx for _ in range(max_new_tokens): next_logits = logits[:, -1, :].float() if no_repeat_ngram > 1 and len(generated) >= no_repeat_ngram - 1: banned = seen.get(tuple(generated[-(no_repeat_ngram - 1):])) if banned: next_logits[:, list(banned)] = float("-inf") if temperature <= 0: next_id = next_logits.argmax(dim=-1, keepdim=True) else: next_logits = next_logits / temperature if top_k: kth = torch.topk(next_logits, min(top_k, next_logits.size(-1))).values[:, -1, None] next_logits = next_logits.masked_fill(next_logits < kth, float("-inf")) probs = F.softmax(next_logits, dim=-1) if top_p and top_p < 1.0: sorted_probs, sorted_idx = probs.sort(descending=True) cumulative = sorted_probs.cumsum(-1) sorted_probs[cumulative - sorted_probs > top_p] = 0.0 probs = torch.zeros_like(probs).scatter(-1, sorted_idx, sorted_probs) probs = probs / probs.sum(-1, keepdim=True) next_id = torch.multinomial(probs, 1) token = int(next_id) generated.append(token) if no_repeat_ngram > 1 and len(generated) >= no_repeat_ngram: seen.setdefault(tuple(generated[-no_repeat_ngram:-1]), set()).add(token) out = torch.cat([out, next_id], dim=1) if token in stop_ids: break if on_token is not None: on_token(token) if pos >= self.cfg.block_size: break logits, _ = self(next_id, caches=caches, start_pos=pos) pos += 1 return out