| """A small GPT-style causal decoder with hooks for probing, steering and patching. |
| |
| Residual stream index l in {0..L}: l=0 is the embedding output (input to block 1), |
| l=L is the output of the last block (before the final LayerNorm). |
| """ |
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class GPTConfig: |
| d_in: int |
| d_out: int |
| T: int = 48 |
| d_model: int = 128 |
| n_layer: int = 4 |
| n_head: int = 4 |
| d_mlp: int = 512 |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg: GPTConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.ln1 = nn.LayerNorm(cfg.d_model) |
| self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model) |
| self.proj = nn.Linear(cfg.d_model, cfg.d_model) |
| self.ln2 = nn.LayerNorm(cfg.d_model) |
| self.mlp = nn.Sequential(nn.Linear(cfg.d_model, cfg.d_mlp), nn.GELU(), nn.Linear(cfg.d_mlp, cfg.d_model)) |
|
|
| def attn(self, x, capture=None, head_edit=None, layer=0): |
| B, T, D = x.shape |
| H = self.cfg.n_head |
| dh = D // H |
| q, k, v = self.qkv(x).split(D, dim=-1) |
| q = q.view(B, T, H, dh).transpose(1, 2) |
| k = k.view(B, T, H, dh).transpose(1, 2) |
| v = v.view(B, T, H, dh).transpose(1, 2) |
| if capture is None and head_edit is None: |
| o = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| else: |
| att = (q @ k.transpose(-1, -2)) / math.sqrt(dh) |
| mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril() |
| att = att.masked_fill(~mask, float("-inf")).softmax(-1) |
| o = att @ v |
| if capture is not None: |
| capture.setdefault("attn", []).append(att.detach()) |
| capture.setdefault("heads", []).append(o.detach()) |
| if head_edit is not None: |
| o = head_edit(layer, o) |
| o = o.transpose(1, 2).reshape(B, T, D) |
| return self.proj(o) |
|
|
| def forward(self, x, capture=None, head_edit=None, layer=0): |
| x = x + self.attn(self.ln1(x), capture, head_edit, layer) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, cfg: GPTConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.embed = nn.Linear(cfg.d_in, cfg.d_model) |
| self.pos = nn.Parameter(torch.zeros(1, cfg.T, cfg.d_model)) |
| nn.init.normal_(self.pos, std=0.02) |
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]) |
| self.ln_f = nn.LayerNorm(cfg.d_model) |
| self.head = nn.Linear(cfg.d_model, cfg.d_out) |
|
|
| def forward(self, tokens, return_resid=False, edit=None, head_edit=None, capture=None): |
| """tokens: (B, T, d_in). edit(l, h) -> h is applied to residual l in {0..L}. |
| head_edit(layer, o) -> o edits per-head attention outputs (B,H,T,dh) of block `layer` (0-based). |
| """ |
| T = tokens.shape[1] |
| h = self.embed(tokens) + self.pos[:, :T] |
| resid = [] |
| for l, blk in enumerate(self.blocks): |
| if edit is not None: |
| h = edit(l, h) |
| if return_resid: |
| resid.append(h) |
| h = blk(h, capture, head_edit, l) |
| L = len(self.blocks) |
| if edit is not None: |
| h = edit(L, h) |
| if return_resid: |
| resid.append(h) |
| out = self.head(self.ln_f(h)) |
| if return_resid: |
| return out, resid |
| return out |
|
|
| def step(self, tok_t, t, cache): |
| """Incremental forward for one new token (B, d_in) at position t, with a differentiable KV cache |
| (list per layer of (k, v) lists). Equivalent to forward(tokens[:, :t+1])[:, -1].""" |
| h = self.embed(tok_t) + self.pos[:, t] |
| H = self.cfg.n_head |
| D = self.cfg.d_model |
| dh = D // H |
| for l, blk in enumerate(self.blocks): |
| x = blk.ln1(h) |
| q, k, v = blk.qkv(x).split(D, dim=-1) |
| cache[l][0].append(k.view(-1, H, 1, dh)) |
| cache[l][1].append(v.view(-1, H, 1, dh)) |
| K = torch.cat(cache[l][0], 2) |
| V = torch.cat(cache[l][1], 2) |
| att = (q.view(-1, H, 1, dh) @ K.transpose(-1, -2)) / math.sqrt(dh) |
| o = (att.softmax(-1) @ V).reshape(-1, D) |
| h = h + blk.proj(o) |
| h = h + blk.mlp(blk.ln2(h)) |
| return self.head(self.ln_f(h)) |
|
|
| def new_cache(self): |
| return [([], []) for _ in self.blocks] |
|
|
| def readout(self, h): |
| return self.head(self.ln_f(h)) |
|
|
| def forward_from(self, h, l, edit=None): |
| """Continue the forward pass from residual h at index l.""" |
| for j in range(l, len(self.blocks)): |
| if edit is not None: |
| h = edit(j, h) |
| h = self.blocks[j](h) |
| if edit is not None: |
| h = edit(len(self.blocks), h) |
| return self.readout(h) |
|
|
|
|
| def n_params(model): |
| return sum(p.numel() for p in model.parameters()) |
|
|
