spb/wpai0001_uqo

1"""Tier-0 toy decoder-only transformer with explicit hook points (PyTorch, CPU/MPS/CUDA).23Pre-LN GPT-style block. Attention is written explicitly (no fused kernel) so that per-head outputs4are exposed at ``attn_heads.{l}`` before the output projection, enabling head-level patching.5"""6from __future__ import annotations78import math910import torch11import torch.nn as nn12import torch.nn.functional as F1314from wp_ai_0001.config import ToyModelConfig15from wp_ai_0001.models.hooks import HookContext, HookManager161718class Attention(nn.Module):19    def __init__(self, cfg: ToyModelConfig, layer: int, hooks: HookManager) -> None:20        super().__init__()21        self.h, self.dh = cfg.n_heads, cfg.d_model // cfg.n_heads22        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)23        self.o = nn.Linear(cfg.d_model, cfg.d_model, bias=False)24        self.layer, self.hooks = layer, hooks2526    def forward(self, x: torch.Tensor, ctx: HookContext) -> torch.Tensor:27        B, T, D = x.shape28        q, k, v = self.qkv(x).split(D, dim=-1)29        q = q.view(B, T, self.h, self.dh).transpose(1, 2)30        k = k.view(B, T, self.h, self.dh).transpose(1, 2)31        v = v.view(B, T, self.h, self.dh).transpose(1, 2)32        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)33        causal = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()34        att = att.masked_fill(~causal, float("-inf")).softmax(-1)35        heads = (att @ v).transpose(1, 2)  # [B,T,H,dh]36        heads = self.hooks(f"attn_heads.{self.layer}", heads, ctx)37        return self.o(heads.reshape(B, T, D))383940class Block(nn.Module):41    def __init__(self, cfg: ToyModelConfig, layer: int, hooks: HookManager) -> None:42        super().__init__()43        self.ln1, self.ln2 = nn.LayerNorm(cfg.d_model), nn.LayerNorm(cfg.d_model)44        self.attn = Attention(cfg, layer, hooks)45        self.mlp = nn.Sequential(nn.Linear(cfg.d_model, cfg.d_ff), nn.GELU(), nn.Linear(cfg.d_ff, cfg.d_model))46        self.layer, self.hooks = layer, hooks4748    def forward(self, x: torch.Tensor, ctx: HookContext) -> torch.Tensor:49        l, hk = self.layer, self.hooks50        x = hk(f"resid_pre.{l}", x, ctx)51        a = hk(f"attn_out.{l}", self.attn(self.ln1(x), ctx), ctx)52        x = hk(f"resid_mid.{l}", x + a, ctx)53        m = hk(f"mlp_out.{l}", self.mlp(self.ln2(x)), ctx)54        return hk(f"resid_post.{l}", x + m, ctx)555657class ToyTransformer(nn.Module):58    def __init__(self, cfg: ToyModelConfig, vocab_size: int) -> None:59        super().__init__()60        self.cfg, self.vocab_size = cfg, vocab_size61        self.hooks = HookManager()62        self.tok = nn.Embedding(vocab_size, cfg.d_model)63        self.pos = nn.Embedding(cfg.max_seq_len, cfg.d_model)64        self.blocks = nn.ModuleList([Block(cfg, i, self.hooks) for i in range(cfg.n_layers)])65        self.ln_f = nn.LayerNorm(cfg.d_model)66        self.head = nn.Linear(cfg.d_model, vocab_size, bias=False)67        if cfg.tie_embeddings:68            self.head.weight = self.tok.weight69        self.apply(self._init)7071    @staticmethod72    def _init(m: nn.Module) -> None:73        if isinstance(m, nn.Linear):74            nn.init.normal_(m.weight, std=0.02)75            if m.bias is not None:76                nn.init.zeros_(m.bias)77        elif isinstance(m, nn.Embedding):78            nn.init.normal_(m.weight, std=0.02)7980    @property81    def n_params(self) -> int:82        return sum(p.numel() for p in self.parameters())8384    def forward(self, ids: torch.Tensor, positions: torch.Tensor | None = None) -> torch.Tensor:85        B, T = ids.shape86        ctx = HookContext(name="", positions=positions, extra={"ids": ids})87        x = self.tok(ids) + self.pos(torch.arange(T, device=ids.device))[None]88        x = self.hooks("embed", x, ctx)89        for blk in self.blocks:90            x = blk(x, ctx)91        x = self.hooks("final_norm", self.ln_f(x), ctx)92        return self.head(x)  # [B,T,V]9394    def answer_logits(self, ids: torch.Tensor, answer_pos: torch.Tensor) -> torch.Tensor:95        logits = self.forward(ids, answer_pos)96        return logits[torch.arange(ids.shape[0], device=ids.device), answer_pos]9798    @torch.no_grad()99    def predict(self, ids: torch.Tensor, answer_pos: torch.Tensor) -> torch.Tensor:100        return self.answer_logits(ids, answer_pos).argmax(-1)101102    def hook_names(self) -> list[str]:103        names = ["embed"]104        for l in range(self.cfg.n_layers):105            names += [f"resid_pre.{l}", f"attn_heads.{l}", f"attn_out.{l}", f"resid_mid.{l}", f"mlp_out.{l}", f"resid_post.{l}"]106        return names + ["final_norm"]107108109def resolve_device(name: str = "auto") -> torch.device:110    if name != "auto":111        return torch.device(name)112    if torch.cuda.is_available():113        return torch.device("cuda")114    if torch.backends.mps.is_available():115        return torch.device("mps")116    return torch.device("cpu")117