Build a GPT from scratch: the uniform baseline

01 · Concept

Concept

The fastest way to make the Transformer concrete is to implement one. The program below is deliberately small but complete: it creates a character tokenizer, samples training batches, implements causal multi-head self-attention, stacks residual blocks, minimizes next-token cross-entropy, and generates text. It uses PyTorch for tensors and automatic differentiation, but no high-level Transformer module.

Call what it builds the uniform baseline: every layer is the same block, every block’s mixer is the same attention, positions come from one learned table, and the norm is classic LayerNorm. This is the pedagogical skeleton the whole track has assembled — and it is exactly the frame you need to appreciate a shipped instance like Qwen3.8-27B, which departs from uniformity on purpose at almost every joint. First build the baseline; the departures come at the end.

Save the block as minimal_gpt.py, install PyTorch with the command documented for your platform at pytorch.org, and run python minimal_gpt.py. It trains on CPU in a short educational run. The tiny repeated corpus and small network are chosen for inspectability, not language quality. The program prints measured losses; their exact values can vary by PyTorch version and hardware.

"""A complete, minimal character-level GPT. Run: python minimal_gpt.py"""

from dataclasses import dataclass

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


torch.manual_seed(42)
device = torch.device("cpu")  # Portable baseline; change deliberately if desired.

# A real training corpus, kept inside the file so the example is self-contained.
text = ("attention routes information; residual streams preserve it.\n" * 200)
chars = sorted(set(text))
stoi = {character: index for index, character in enumerate(chars)}
itos = {index: character for character, index in stoi.items()}
data = torch.tensor([stoi[character] for character in text], dtype=torch.long)


@dataclass(frozen=True)
class Config:
    vocab_size: int = len(chars)
    block_size: int = 32
    batch_size: int = 16
    n_embed: int = 64
    n_head: int = 4
    n_layer: int = 2


config = Config()
split = int(0.9 * len(data))
train_data, validation_data = data[:split], data[split:]


def get_batch(source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return prefixes x and their one-position-shifted next-token targets y."""
    starts = torch.randint(len(source) - config.block_size - 1, (config.batch_size,))
    x = torch.stack([source[i : i + config.block_size] for i in starts])
    y = torch.stack([source[i + 1 : i + config.block_size + 1] for i in starts])
    return x.to(device), y.to(device)


class AttentionHead(nn.Module):
    def __init__(self, head_size: int) -> None:
        super().__init__()
        self.key = nn.Linear(config.n_embed, head_size, bias=False)
        self.query = nn.Linear(config.n_embed, head_size, bias=False)
        self.value = nn.Linear(config.n_embed, head_size, bias=False)
        # This lower triangle moves with the model but receives no gradients.
        self.register_buffer("causal", torch.tril(torch.ones(config.block_size, config.block_size)))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        _, time, _ = x.shape
        key, query, value = self.key(x), self.query(x), self.value(x)
        # (B,T,H) @ (B,H,T) -> (B,T,T): every query scores every key.
        weights = query @ key.transpose(-2, -1) * (key.shape[-1] ** -0.5)
        weights = weights.masked_fill(self.causal[:time, :time] == 0, float("-inf"))
        weights = F.softmax(weights, dim=-1)
        return weights @ value  # (B,T,T) @ (B,T,H) -> (B,T,H)


class MultiHeadAttention(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        head_size = config.n_embed // config.n_head
        self.heads = nn.ModuleList(
            [AttentionHead(head_size) for _ in range(config.n_head)]
        )
        self.output = nn.Linear(config.n_embed, config.n_embed)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Heads route independently; concatenation restores the model width.
        return self.output(torch.cat([head(x) for head in self.heads], dim=-1))


class FeedForward(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(config.n_embed, 4 * config.n_embed),
            nn.GELU(),
            nn.Linear(4 * config.n_embed, config.n_embed),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


class Block(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.attention = MultiHeadAttention()
        self.feed_forward = FeedForward()
        self.norm_attention = nn.LayerNorm(config.n_embed)
        self.norm_ffn = nn.LayerNorm(config.n_embed)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-norm residual updates: communicate first, transform second.
        x = x + self.attention(self.norm_attention(x))
        x = x + self.feed_forward(self.norm_ffn(x))
        return x


class MiniGPT(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embed)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embed)
        self.blocks = nn.Sequential(*[Block() for _ in range(config.n_layer)])
        self.final_norm = nn.LayerNorm(config.n_embed)
        self.language_head = nn.Linear(config.n_embed, config.vocab_size)

    def forward(
        self, indices: torch.Tensor, targets: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        _, time = indices.shape
        positions = torch.arange(time, device=indices.device)
        x = self.token_embedding(indices) + self.position_embedding(positions)
        x = self.blocks(x)
        logits = self.language_head(self.final_norm(x))  # (B,T,V)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, config.vocab_size), targets.reshape(-1)
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, indices: torch.Tensor, new_tokens: int) -> torch.Tensor:
        for _ in range(new_tokens):
            context = indices[:, -config.block_size :]
            logits, _ = self(context)
            probabilities = F.softmax(logits[:, -1, :], dim=-1)
            next_token = torch.multinomial(probabilities, num_samples=1)
            indices = torch.cat((indices, next_token), dim=1)
        return indices


model = MiniGPT().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)

for step in range(301):
    inputs, targets = get_batch(train_data)
    _, loss = model(inputs, targets)
    if loss is None:
        raise RuntimeError("training targets must produce a loss")
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    if step % 100 == 0:
        print(f"step {step}: train loss {loss.item():.4f}")

seed = torch.zeros((1, 1), dtype=torch.long, device=device)
generated = model.generate(seed, new_tokens=160)[0].tolist()
print("".join(itos[index] for index in generated))

Read it by following shapes. indices is $B\times T$. Token and position lookup produce $B\times T\times C$. Each head projects $C$ features to $H=C/h$, forms a $T\times T$ score matrix per batch item, applies the lower-triangular mask, and mixes values. Concatenating $h$ heads restores $C$. The FFN expands to $4C$ and returns to $C$. Every residual addition therefore has matching shape. The target is the input shifted by one position; cross-entropy compares vocabulary logits at each position with the actual next character; AdamW updates the parameters while the causal buffer stays fixed. Generation uses the same forward pass without targets, sampling from the final position’s logits and cropping to the context limit.

One incidental detail is worth flagging: language_head is a separate matrix from token_embedding — this toy does not tie its output head to its embedding table. Neither does Qwen3.8-27B, at vastly greater expense; hold that thought for the ladder below and for lesson 4.14, where the two untied matrices appear in the checkpoint as separate tensors.

Now climb the parameter ladder from this table to the plant floor, and let each rung be a number you can check. The toy: a character vocabulary of a few dozen entries, width 64, two layers, an ungated $4\times$ FFN — well under a million parameters, small enough that the printed loss moves within seconds. First rung, width: Qwen3.8-27B’s residual stream is 5120 wide, eighty times the toy’s 64, and parameter cost in the dense blocks grows roughly with the square of width. Second rung, depth: 64 layers against the toy’s two. Third rung, vocabulary: 248,320 rows against a few dozen, so the embedding table alone reaches about 1.271 billion parameters — and because the head is untied, the output matrix is another ~1.271 billion, roughly 2.54 billion before any block does any work. Fourth rung, the blocks: each layer’s gated FFN holds about 267 million parameters, and 64 of them total roughly 17 billion — the single largest line item, as lesson 4.9 established. The remainder of the ~27 billion budget covers the mixing layers — sixteen attention blocks with their 12288- and 1024-wide projections, forty-eight DeltaNet blocks — plus the vision tower that lets the model read images (lesson 4.17). Every rung is the same arithmetic you just ran on Config: rows times columns, summed over named modules. The plant is bigger than the tabletop press by a factor of hundreds of thousands, but it is inventoried with the multiplication you can now do by hand.