Build a GPT

import torch, torch.nn as nn

class Model1(nn.Module):
    """Predicts the next token from the current token alone. No context."""
    def __init__(self, vocab, dim=64):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, dim)   # ids   -> vectors
        self.head    = nn.Linear(dim, vocab)      # vector -> scores

    def forward(self, idx):        # idx: [B, T] integers
        x = self.tok_emb(idx)      #      [B, T, dim]
        return self.head(x)        #      [B, T, vocab]

# This is a learned bigram model. It cannot use context, so its loss floors out
# early — which is exactly the point: everything added later exists to fix that.

import numpy as np, torch, torch.nn.functional as F

def softmax(logits, temperature=1.0):
    z = np.asarray(logits, dtype=np.float64) / temperature
    z = z - z.max()                 # shift for stability; result unchanged
    e = np.exp(z)
    return e / e.sum()

logits = [2.0, 1.0, 0.1]
print(softmax(logits, 1.0))    # [0.659 0.242 0.099]
print(softmax(logits, 0.5))    # sharper
print(softmax(logits, 2.0))    # flatter

# Without the shift, large logits overflow:
np.exp(np.array([1000.0]))     # inf -> nan after division

# PyTorch, on the last dimension:
print(F.softmax(torch.tensor(logits) / 0.7, dim=-1))

import torch

def sample_next(logits, temperature=1.0, top_k=None, top_p=None):
    logits = logits / max(temperature, 1e-8)

    if top_k is not None:
        kth = torch.topk(logits, top_k).values[..., -1, None]
        logits = logits.masked_fill(logits < kth, float("-inf"))

    if top_p is not None:
        srt, idx = torch.sort(logits, descending=True)
        cum = torch.softmax(srt, dim=-1).cumsum(dim=-1)
        drop = cum - torch.softmax(srt, dim=-1) > top_p   # keep the first one always
        srt = srt.masked_fill(drop, float("-inf"))
        logits = torch.empty_like(logits).scatter_(-1, idx, srt)

    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)   # samples proportionally

@torch.no_grad()
def generate(model, idx, max_new=50, **kw):
    for _ in range(max_new):
        logits = model(idx[:, -model.block_size:])   # crop to context length
        nxt = sample_next(logits[:, -1, :], **kw)    # only the last position
        idx = torch.cat([idx, nxt], dim=1)           # append and repeat
    return idx

import torch, torch.nn as nn

class Model2(nn.Module):
    def __init__(self, vocab, dim=64, block_size=32):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab, dim)
        self.pos_emb = nn.Embedding(block_size, dim)   # learned, like GPT-2
        self.ln      = nn.LayerNorm(dim)
        self.head    = nn.Linear(dim, vocab, bias=False)
        # Weight tying: the unembedding IS the transposed embedding. Saves
        # vocab*dim parameters and usually improves quality.
        self.head.weight = self.tok_emb.weight

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)      # broadcast over batch
        return self.head(self.ln(x))

import torch

T = 4
ones = torch.tril(torch.ones(T, T))          # lower triangular
w = ones / ones.sum(dim=1, keepdim=True)     # each row averages its prefix
print(w)
# [[1.00, 0,    0,    0   ],
#  [0.50, 0.50, 0,    0   ],
#  [0.33, 0.33, 0.33, 0   ],
#  [0.25, 0.25, 0.25, 0.25]]

x = torch.randn(T, 8)
print(torch.allclose(w @ x, torch.stack([x[:i+1].mean(0) for i in range(T)])))

# Equivalent, and the form attention actually uses: mask with -inf, then softmax.
scores = torch.zeros(T, T).masked_fill(ones == 0, float("-inf"))
print(torch.softmax(scores, dim=-1))         # same uniform prefix weights
# -inf is used because exp(-inf) = 0 exactly. Masking after softmax would
# leave the removed positions contributing to the denominator.

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

class Head(nn.Module):
    def __init__(self, dim, head_dim, block_size):
        super().__init__()
        self.key   = nn.Linear(dim, head_dim, bias=False)
        self.query = nn.Linear(dim, head_dim, bias=False)
        self.value = nn.Linear(dim, head_dim, bias=False)
        # A buffer moves with .to(device) but is not a parameter.
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):                      # x: [B, T, dim]
        B, T, _ = x.shape
        k, q, v = self.key(x), self.query(x), self.value(x)

        att = q @ k.transpose(-2, -1) * k.size(-1) ** -0.5    # [B, T, T]
        att = att.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        return att @ v                                        # [B, T, head_dim]

# PyTorch's fused equivalent — same maths, far less memory:
# F.scaled_dot_product_attention(q, k, v, is_causal=True)

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

class MultiHeadAttention(nn.Module):
    def __init__(self, dim, n_heads, block_size, dropout=0.1):
        super().__init__()
        assert dim % n_heads == 0
        self.n_heads, self.head_dim = n_heads, dim // n_heads
        self.qkv  = nn.Linear(dim, 3 * dim, bias=False)   # all three at once
        self.proj = nn.Linear(dim, dim)
        self.drop = nn.Dropout(dropout)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # [B, T, C] -> [B, heads, T, head_dim]; heads become a batch dimension.
        shape = lambda t: t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        q, k, v = shape(q), shape(k), shape(v)

        att = (q @ k.transpose(-2, -1)) * self.head_dim ** -0.5
        att = att.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
        att = self.drop(F.softmax(att, dim=-1))

        out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)   # re-merge
        return self.drop(self.proj(out))

import torch, torch.nn as nn

class Block(nn.Module):
    """Pre-norm: normalise, sublayer, add back to the residual stream."""
    def __init__(self, dim, n_heads, block_size, dropout=0.1):
        super().__init__()
        self.ln1  = nn.LayerNorm(dim)
        self.attn = MultiHeadAttention(dim, n_heads, block_size, dropout)
        self.ln2  = nn.LayerNorm(dim)
        self.mlp  = nn.Sequential(
            nn.Linear(dim, 4 * dim),   # the 4x expansion is conventional
            nn.GELU(),
            nn.Linear(4 * dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # communicate across positions
        x = x + self.mlp(self.ln2(x))    # compute within each position
        return x

class GPT(nn.Module):
    def __init__(self, vocab, dim=384, n_heads=6, n_layers=6, block_size=256, dropout=0.1):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab, dim)
        self.pos_emb = nn.Embedding(block_size, dim)
        self.drop    = nn.Dropout(dropout)
        self.blocks  = nn.ModuleList(
            [Block(dim, n_heads, block_size, dropout) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, vocab, bias=False)
        self.head.weight = self.tok_emb.weight        # tied

    def forward(self, idx):
        B, T = idx.shape
        assert T <= self.block_size, "sequence longer than the context window"
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        for block in self.blocks:
            x = block(x)
        return self.head(self.ln_f(x))

import time, torch

def benchmark(model, x, n=20):
    device = x.device
    for _ in range(3):                       # warm-up: kernels, autotuning
        model(x)
    if device.type == "cuda":
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n):
        model(x)
    if device.type == "cuda":
        torch.cuda.synchronize()             # without this you time the queue
    return (time.perf_counter() - t0) / n

x = torch.randint(0, vocab, (16, 256))
print("cpu", benchmark(model.cpu(), x.cpu()))
if torch.cuda.is_available():
    print("gpu", benchmark(model.cuda(), x.cuda()))

from transformers import AutoModelForCausalLM
import torch

def count(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total, trainable

gpt2 = AutoModelForCausalLM.from_pretrained("gpt2")
print(count(gpt2))                       # (124439808, 124439808)

# Where they are:
for name, p in gpt2.named_parameters():
    if p.dim() > 1:
        print(f"{name:40} {tuple(p.shape)} {p.numel():>10,}")

# Trained weight distributions are informative: most are near-Gaussian and
# tightly centred, layernorm gains sit near 1, and a small number of very
# large outlier weights appear in specific dimensions.
w = gpt2.transformer.h[0].attn.c_attn.weight.detach()
print(w.mean().item(), w.std().item(), w.abs().max().item())