code/gpu/models.py

"""Minimal GPT-2-style model (nanoGPT-like) shared by the GPU experiments."""

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class GPTConfig:
    vocab_size: int = 50304  # GPT-2 BPE padded (nanoGPT convention)
    ctx: int = 1024
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768


SIZES = {
    "gpt-30m": GPTConfig(n_layer=6, n_head=6, n_embd=384, ctx=1024),
    "gpt-82m": GPTConfig(n_layer=6, n_head=12, n_embd=768, ctx=1024),
    "gpt-124m": GPTConfig(n_layer=12, n_head=12, n_embd=768, ctx=1024),
    "gpt-355m": GPTConfig(n_layer=24, n_head=16, n_embd=1024, ctx=1024),
}


class Block(nn.Module):
    def __init__(self, c: GPTConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(c.n_embd)
        self.attn_qkv = nn.Linear(c.n_embd, 3 * c.n_embd)
        self.attn_proj = nn.Linear(c.n_embd, c.n_embd)
        self.ln_2 = nn.LayerNorm(c.n_embd)
        self.mlp_fc = nn.Linear(c.n_embd, 4 * c.n_embd)
        self.mlp_proj = nn.Linear(4 * c.n_embd, c.n_embd)
        self.n_head = c.n_head

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.attn_qkv(self.ln_1(x)).split(C, dim=2)
        q = q.view(B, T, self.n_head, -1).transpose(1, 2)
        k = k.view(B, T, self.n_head, -1).transpose(1, 2)
        v = v.view(B, T, self.n_head, -1).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        x = x + self.attn_proj(y)
        h = self.mlp_proj(F.gelu(self.mlp_fc(self.ln_2(x))))
        return x + h


class GPT(nn.Module):
    def __init__(self, c: GPTConfig):
        super().__init__()
        self.c = c
        self.wte = nn.Embedding(c.vocab_size, c.n_embd)
        self.wpe = nn.Embedding(c.ctx, c.n_embd)
        self.h = nn.ModuleList(Block(c) for _ in range(c.n_layer))
        self.ln_f = nn.LayerNorm(c.n_embd)
        self.lm_head = nn.Linear(c.n_embd, c.vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight  # tied
        self.apply(self._init)
        for n, p in self.named_parameters():  # GPT-2 residual-proj scaling
            if n.endswith("proj.weight"):
                nn.init.normal_(p, std=0.02 / math.sqrt(2 * c.n_layer))

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.wte(idx) + self.wpe(pos)
        for b in self.h:
            x = b(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        if targets is None:
            return logits, None
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss


def param_groups(model, weight_decay=0.1):
    """Weight decay on >=2D params only (paper App. B.2 convention)."""
    decay = [p for p in model.parameters() if p.requires_grad and p.dim() >= 2]
    nodecay = [p for p in model.parameters() if p.requires_grad and p.dim() < 2]
    return [{"params": decay, "weight_decay": weight_decay},
            {"params": nodecay, "weight_decay": 0.0}]