model.py

"""MicroMe transformer — single readable file (plan §4).

Modern-Llama skeleton with the cheap speedrun upgrades:
RMSNorm pre-norm, RoPE, QK-norm, ReLU^2 MLP (no gate), untied embeddings,
logit softcap, zero-init residual projections. SDPA -> FlashAttention kernel.

Kept deliberately hackable: no framework, no config files yet. Presets live
in GPTConfig. Everything else in the plan (yaml runs, Muon, data shards) is a
later milestone and is NOT scaffolded here on purpose.
"""
from dataclasses import dataclass
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class GPTConfig:
    vocab_size: int = 512
    n_layer: int = 4
    n_head: int = 4
    d_model: int = 128
    ctx: int = 64
    softcap: float = 15.0
    rope_base: float = 10000.0

    @property
    def head_dim(self) -> int:
        assert self.d_model % self.n_head == 0
        hd = self.d_model // self.n_head
        assert hd % 2 == 0, "head_dim must be even for RoPE"
        return hd


# Presets from the plan (§4). micro_1m is the M0 smoke model (~1M params).
PRESETS = {
    "micro_1m":   GPTConfig(vocab_size=512,   n_layer=4,  n_head=4,  d_model=128, ctx=64),
    "micro_30m":  GPTConfig(vocab_size=24576, n_layer=8,  n_head=7,  d_model=448, ctx=512),
    "micro_125m": GPTConfig(vocab_size=24576, n_layer=20, n_head=10, d_model=640, ctx=1024),
}


def rmsnorm(x, weight=None, eps=1e-6):
    x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
    return x if weight is None else x * weight


class RMSNorm(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return rmsnorm(x, self.weight)


def apply_rope(x, cos, sin):
    # x: (B, nh, T, hd). Interleaved rotary.
    T = x.size(-2)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    cos, sin = cos[:T], sin[:T]  # (T, hd/2)
    xr1 = x1 * cos - x2 * sin
    xr2 = x1 * sin + x2 * cos
    return torch.stack((xr1, xr2), dim=-1).flatten(-2)


class Attention(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.nh, self.hd = cfg.n_head, cfg.head_dim
        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)

    def forward(self, x, cos, sin):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q = q.view(B, T, self.nh, self.hd).transpose(1, 2)
        k = k.view(B, T, self.nh, self.hd).transpose(1, 2)
        v = v.view(B, T, self.nh, self.hd).transpose(1, 2)
        q, k = rmsnorm(q), rmsnorm(k)               # QK-norm (weightless)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).reshape(B, T, C)
        return self.proj(y)


class MLP(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        hidden = 4 * cfg.d_model
        self.fc = nn.Linear(cfg.d_model, hidden, bias=False)
        self.proj = nn.Linear(hidden, cfg.d_model, bias=False)

    def forward(self, x):
        return self.proj(F.relu(self.fc(x)) ** 2)   # ReLU^2, no gate


class Block(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.n1, self.attn = RMSNorm(cfg.d_model), Attention(cfg)
        self.n2, self.mlp = RMSNorm(cfg.d_model), MLP(cfg)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.n1(x), cos, sin)
        x = x + self.mlp(self.n2(x))
        return x


class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg
        self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer))
        self.norm_f = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)  # untied

        # RoPE tables (buffers, not params).
        hd = cfg.head_dim
        inv_freq = 1.0 / (cfg.rope_base ** (torch.arange(0, hd, 2).float() / hd))
        t = torch.arange(cfg.ctx).float()
        freqs = torch.outer(t, inv_freq)            # (ctx, hd/2)
        self.register_buffer("cos", freqs.cos(), persistent=False)
        self.register_buffer("sin", freqs.sin(), persistent=False)

        self.apply(self._init)
        # zero-init residual output projections (stable start)
        for blk in self.blocks:
            nn.init.zeros_(blk.attn.proj.weight)
            nn.init.zeros_(blk.mlp.proj.weight)

    def _init(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 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)
        assert T <= self.cfg.ctx, f"seq len {T} > ctx {self.cfg.ctx}"
        x = self.wte(idx)
        for blk in self.blocks:
            x = blk(x, self.cos, self.sin)
        logits = self.lm_head(self.norm_f(x))
        cap = self.cfg.softcap
        logits = cap * torch.tanh(logits / cap)     # logit softcap
        loss = None
        if targets is not None:
            if loss_mask is None:                    # pretraining: mean over all tokens
                loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            else:                                    # SFT: mean over assistant tokens only
                ce = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), reduction="none")
                m = loss_mask.reshape(-1).to(ce.dtype)
                loss = (ce * m).sum() / m.sum().clamp(min=1)
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=None, rep_penalty=1.0):
        # no KV cache yet (that's M7); fine for short eval samples.
        start = idx.size(1)                      # penalize only tokens WE generate, not the prompt/context
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -self.cfg.ctx:])
            logits = logits[:, -1, :]
            if rep_penalty != 1.0 and idx.size(1) > start:
                for b in range(idx.size(0)):     # divide logits of already-generated tokens (anti-loop)
                    logits[b, idx[b, start:]] /= rep_penalty
            logits = logits / max(temperature, 1e-5)
            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)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx


if __name__ == "__main__":
    m = GPT(PRESETS["micro_1m"])
    print(f"micro_1m params: {m.num_params()/1e6:.2f}M")
    x = torch.randint(0, 512, (2, 64))
    logits, loss = m(x, x)
    print("logits", tuple(logits.shape), "loss", float(loss))