Build a Tiny LLM From Scratch

The plan: every piece, in one file

Fourteen articles is a lot of surface area, and it's tempting to assume a real GPT implementation must match it: thousands of lines spread across a dozen files, a research codebase you'd need a week to onboard onto. It isn't. The entire architecture, tokenizer, embeddings, positional encoding, a stack of transformer blocks, a loss function, an optimizer, and a sampling loop, fits in well under 300 lines, small enough to read start to finish in one sitting and hold the whole thing in your head at once.

A Cessna and a 747 are both, mechanically, wings, an engine, and control surfaces. Nobody confuses the two, but if you understand how the Cessna flies, you understand the shape of how the 747 flies too, the differences are in scale and refinement, not in kind. That's the relationship between what gets built in this article and a production model. Same operations, same tensor shapes, same training loop, same sampling logic. The only thing that changes going from here to a frontier model is size: more parameters, more data, more compute, applied to the identical mechanism.

The target for this build: a character-level tokenizer, a d_model of 128, 4 attention heads, 6 stacked transformer blocks, and a context window of 128 tokens. That works out to roughly 1.2 million tunable parameters, small enough to train on a plain text corpus a few hundred kilobytes in size, on a laptop CPU, in minutes rather than days. Obviously nowhere near a real model, a frontier LLM runs somewhere between four and five orders of magnitude more parameters and roughly nine orders of magnitude more training tokens. But every operation this tiny model performs is the same operation a 400-billion parameter model performs. Nothing about the mechanism knows or cares how big the numbers get.

Here's how the last fourteen articles map onto the four sections that follow. Tokenization and Embeddings and Positional Encoding become the input pipeline, next. Neural Network Foundations and Attention and The Transformer Block become the model itself, the section after that. Loss and Backpropagation, and the toy-scale version of Pretraining at Scale, become the training loop. And Sampling and Generation, closing the loop first opened by What Is a Language Model, becomes the function that actually produces text. From Base Model to Assistant and Scaling Laws don't have a toy-scale equivalent worth building, instruction tuning needs a preference dataset this corpus doesn't have, and scaling laws are an empirical claim about trends across many models, not something one model demonstrates on its own, so both are set aside here. Context Windows and the KV Cache gets one direct mention, in the generation function where the tradeoff it describes actually shows up in code. Everything else gets used, by name, exactly where it belongs.

Tokenizer and embeddings, wired together

It's tempting to reach straight for a real byte-pair encoder, the Tokenization article covered exactly how one is built, learned merges and all. But a full BPE vocabulary is built to handle open-domain text efficiently, and that's not the constraint here. The toy corpus for this build is a plain text file a few hundred kilobytes in size, doesn't matter what it contains, and a tokenizer's job, however it's implemented, is unchanged from that article: turn text into integers and back. So the simplest tokenizer that does that job honestly is a character-level one, every unique character in the corpus gets one integer ID. Same interface as a BPE tokenizer, encode and decode, just a coarser vocabulary.

1class CharTokenizer:
2    """Simplest possible stand-in for the BPE tokenizer from the
3    Tokenization article. Same job, text to integers and back, a coarser
4    vocabulary: one entry per unique character instead of learned merges."""
5
6    def __init__(self, text: str):
7        chars = sorted(set(text))
8        self.vocab_size = len(chars)
9        self.stoi = {ch: i for i, ch in enumerate(chars)}
10        self.itos = {i: ch for i, ch in enumerate(chars)}
11
12    def encode(self, text: str) -> list[int]:
13        return [self.stoi[ch] for ch in text]
14
15    def decode(self, ids: list[int]) -> str:
16        return "".join(self.itos[i] for i in ids)
17
18
19text = open("corpus.txt").read()
20tokenizer = CharTokenizer(text)
21print(tokenizer.vocab_size)
22print(tokenizer.encode("hello"))

vocab_sizehere comes out to 65 for a typical English corpus, uppercase and lowercase letters, punctuation, whitespace. That number matters, it's the input dimension of the embedding table built next, and it's referenced again unchanged all the way through the output projection at the far end of the model.

The Embeddings article established that an embedding table is, mechanically, a lookup table: one row of learned numbers per vocabulary entry, and nn.Embeddingis exactly that, nothing more exotic than an indexable matrix that gradients flow into during training. Token ID 46 doesn't mean anything on its own, it's an address. The row at that address is what carries meaning, and that row gets updated every time backpropagation runs, the same way it would for a 100,000-entry BPE vocabulary.

On its own, a token embedding can't tell “cat sat” from “sat cat,” the vectors get looked up identically regardless of position. That's the gap the Positional Encoding article closed. This build uses the simpler of the two options that article covered, a learned position embedding table instead of a fixed sinusoidal one, one row per position from 0 up to block_size. Simpler to reason about, same purpose: give the model a way to tell first token from fifth from fiftieth.

1import torch
2import torch.nn as nn
3
4class TokenAndPositionEmbedding(nn.Module):
5    def __init__(self, vocab_size: int, d_model: int, block_size: int):
6        super().__init__()
7        self.token_emb = nn.Embedding(vocab_size, d_model)
8        self.pos_emb = nn.Embedding(block_size, d_model)
9
10    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
11        B, T = token_ids.shape
12        positions = torch.arange(T, device=token_ids.device)
13        # broadcast: (T, d_model) added onto every row in the batch
14        return self.token_emb(token_ids) + self.pos_emb(positions)

Wiring the two pieces together end to end, a batch of raw text turns into a batch of position-aware vectors in two calls:

1d_model = 128
2block_size = 128
3
4embed = TokenAndPositionEmbedding(tokenizer.vocab_size, d_model, block_size)
5
6batch_text = ["to be or not to", "the trophy did"]
7token_ids = torch.tensor([tokenizer.encode(t) for t in batch_text])  # (B, T)
8x = embed(token_ids)                                                # (B, T, d_model)
9print(x.shape)

x is now a batch of shape (B, T, d_model), one 128-dimensional vector per token, position baked in by addition rather than left implicit. Every downstream piece in this article, every transformer block, the loss function, the sampling loop, only ever operates on tensors of this shape. Get this part right and the rest is composition.

The transformer block stack

It's easy to look at a transformer architecture diagram and assume the block itself is some new invention, a fifth fundamental operation sitting alongside attention and feedforward layers. It isn't. Everything in this section is the weighted sum from the Attention article and the stacked linear-plus-nonlinearity layers from Neural Network Foundations, wrapped in exactly two addition operations and two calls to LayerNorm. The Transformer Block article called that wrapping load-bearing, and it is, but it isn't new arithmetic. It's assembly.

Start with multi-head attention itself, the query-key-value mechanism from the Attention article, run several times in parallel with independently learned projections, then concatenated back together. One addition here that article's toy example skipped for clarity: a causal mask. Token 5 can attend to tokens 0 through 5, never to token 6 or beyond, otherwise the model would be predicting the next token by looking at it directly.

1import math
2import torch
3import torch.nn as nn
4import torch.nn.functional as F
5
6class MultiHeadAttention(nn.Module):
7    def __init__(self, d_model: int, n_heads: int, block_size: int, dropout: float = 0.1):
8        super().__init__()
9        assert d_model % n_heads == 0
10        self.n_heads = n_heads
11        self.head_dim = d_model // n_heads
12
13        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
14        self.out_proj = nn.Linear(d_model, d_model)
15        self.dropout = nn.Dropout(dropout)
16
17        # token t can only attend to tokens <= t
18        mask = torch.tril(torch.ones(block_size, block_size))
19        self.register_buffer("mask", mask)
20
21    def forward(self, x: torch.Tensor) -> torch.Tensor:
22        B, T, C = x.shape
23        q, k, v = self.qkv_proj(x).split(C, dim=-1)
24
25        # (B, T, C) -> (B, n_heads, T, head_dim), one slice per head
26        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
27        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
28        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
29
30        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
31        scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
32        weights = self.dropout(F.softmax(scores, dim=-1))
33
34        out = weights @ v
35        out = out.transpose(1, 2).contiguous().view(B, T, C)
36        return self.out_proj(out)

Attention on its own is a weighted average, a linear combination of value vectors. Neural Network Foundations covered why stacking linear operations without a nonlinearity between them collapses into a single linear operation, no additional expressive power gained. That's exactly why a feedforward layer follows every attention layer, a widen-then-project stack with a nonlinearity in between: the same neuron-and-activation-function machinery from that article, applied identically to every token position.

1class FeedForward(nn.Module):
2    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
3        super().__init__()
4        self.net = nn.Sequential(
5            nn.Linear(d_model, d_ff),
6            nn.GELU(),
7            nn.Linear(d_ff, d_model),
8            nn.Dropout(dropout),
9        )
10
11    def forward(self, x: torch.Tensor) -> torch.Tensor:
12        return self.net(x)

Now the assembly The Transformer Block article walked through: attention and feedforward, each wrapped in a residual connection, addition rather than replacement, plus a LayerNorm before each one to keep activations from drifting as they pass through six stacked copies of this block.

1class TransformerBlock(nn.Module):
2    def __init__(self, d_model: int, n_heads: int, d_ff: int, block_size: int, dropout: float = 0.1):
3        super().__init__()
4        self.ln1 = nn.LayerNorm(d_model)
5        self.attn = MultiHeadAttention(d_model, n_heads, block_size, dropout)
6        self.ln2 = nn.LayerNorm(d_model)
7        self.ff = FeedForward(d_model, d_ff, dropout)
8
9    def forward(self, x: torch.Tensor) -> torch.Tensor:
10        x = x + self.attn(self.ln1(x))
11        x = x + self.ff(self.ln2(x))
12        return x

Stack six of these, add a final LayerNorm, and project back from d_model down to vocab_size, and the model is complete: input token IDs in, one logit per vocabulary entry out, for every position in the sequence.

1class TinyGPT(nn.Module):
2    def __init__(
3        self,
4        vocab_size: int,
5        d_model: int = 128,
6        n_heads: int = 4,
7        n_layers: int = 6,
8        block_size: int = 128,
9        dropout: float = 0.1,
10    ):
11        super().__init__()
12        self.block_size = block_size
13        self.embed = TokenAndPositionEmbedding(vocab_size, d_model, block_size)
14        self.blocks = nn.ModuleList([
15            TransformerBlock(d_model, n_heads, 4 * d_model, block_size, dropout)
16            for _ in range(n_layers)
17        ])
18        self.ln_final = nn.LayerNorm(d_model)
19        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
20        self.lm_head.weight = self.embed.token_emb.weight  # weight tying
21
22    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
23        x = self.embed(token_ids)          # (B, T, d_model)
24        for block in self.blocks:
25            x = block(x)                   # (B, T, d_model)
26        x = self.ln_final(x)
27        logits = self.lm_head(x)           # (B, T, vocab_size)
28        return logits
29
30
31model = TinyGPT(vocab_size=tokenizer.vocab_size)
32n_params = sum(p.numel() for p in model.parameters())
33print(f"{n_params:,} parameters")

One detail worth pausing on: lm_head.weightis set equal to the token embedding's weight, not copied, the same tensor, shared. The row of numbers the model uses to represent a token going in is literally reused as the row of numbers it uses to score that token coming out. Fewer parameters to learn, and a representation that has to pull double duty, which in practice makes it a better one.

The training loop

A freshly initialized TinyGPTis a random function. Feed it a prompt and it produces logits, but the logits are noise, the model has never seen a correct answer. It's tempting to think turning that noise into something coherent takes some separate mysterious optimization machinery. It doesn't. It takes exactly the four steps the Loss and Backpropagation article named: forward pass, compute loss, backward pass, optimizer step, repeated over batches until the loss stops going down.

First, batches. Every training step needs input sequences and, for each one, the correct next token at every position. Because this is next-token prediction, the target sequence is just the input sequence shifted one position to the right, the same objective the very first article in this series defined and every article since has assumed.

1def get_batch(data: torch.Tensor, block_size: int, batch_size: int):
2    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
3    x = torch.stack([data[i : i + block_size] for i in ix])
4    y = torch.stack([data[i + 1 : i + block_size + 1] for i in ix])
5    return x, y
6
7
8all_ids = torch.tensor(tokenizer.encode(text), dtype=torch.long)
9n = int(0.9 * len(all_ids))
10train_data, val_data = all_ids[:n], all_ids[n:]
11
12xb, yb = get_batch(train_data, block_size=128, batch_size=32)
13print(xb.shape, yb.shape)

xb and yb are identical in shape and offset by exactly one position. Position t of yb is the correct answer for what should come after position t of xb. Every position in every sequence in the batch produces its own training signal at once, not just the last one, which is why a batch of 32 sequences of length 128 yields 4,096 individual next-token predictions per step.

Now the loop itself. Forward pass through TinyGPT to get logits, cross-entropy between those logits and yb, loss.backward()to run backpropagation and populate every parameter's .grad, and optimizer.step()to nudge every parameter downhill against its gradient. Adam, the adaptive variant of gradient descent covered in that article, tracks a running estimate of each parameter's gradient variance and scales its step size accordingly, which is why it's the default choice here rather than plain gradient descent.

1model = TinyGPT(vocab_size=tokenizer.vocab_size)
2optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
3
4block_size = 128
5batch_size = 32
6
7for step in range(3000):
8    xb, yb = get_batch(train_data, block_size, batch_size)
9
10    logits = model(xb)                              # (B, T, vocab_size)
11    B, T, V = logits.shape
12    loss = F.cross_entropy(logits.view(B * T, V), yb.view(B * T))
13
14    optimizer.zero_grad()
15    loss.backward()
16    optimizer.step()
17
18    if step % 500 == 0:
19        print(f"step {step:>5} | train loss {loss.item():.4f}")

logits.view(B * T, V)is the one reshape worth pausing on: cross-entropy in PyTorch expects one row of logits per prediction, not per sequence, so the batch and time dimensions get flattened together before comparing against the flattened targets. Every one of those 4,096 flattened rows contributes its own gradient, and it's the sum of all of them, run backward through every transformer block, every attention projection, every feedforward layer, all the way to the embedding table, that loss.backward()computes in that single line.

Zoom out and this is the Pretraining at Scale article, at toy scale instead of production scale. Same objective, same loop shape, same optimizer. The only differences are the ones that article predicted: a corpus measured in hundreds of kilobytes instead of trillions of tokens, and a training run measured in a few thousand steps on a laptop instead of hundreds of thousands of steps across a cluster of GPUs running for weeks.

Generation: sampling text from the trained model

The obvious way to turn logits into text is to always take the highest-scoring token, argmax, and repeat. Try it on a trained model and the output degrades fast: loops, “the the the the,” the same safe phrase recurring every few words. Argmax is deterministic and greedy, and greedy decoding walks straight into the most probable immediate next word at every step without ever considering whether a slightly less probable word now leads somewhere better later. The Sampling and Generation article covered the fix: don't always take the top token, sample from the distribution, shaped by temperature, top-k, and top-p so the sampling stays sensible instead of picking uniformly at random.

Generation itself is the loop that What Is a Language Model first described at the very start of this series: feed in a prompt, get a probability distribution over the next token, pick one, append it, and feed the whole thing back in for the next step. TinyGPT's forward pass already returns exactly that distribution, as logits, one per position. The generatefunction below only needs the last position's logits, reshapes them with temperature, filters them with top-k and top-p, and samples.

1@torch.no_grad()
2def generate(
3    model: TinyGPT,
4    tokenizer: CharTokenizer,
5    prompt: str,
6    max_new_tokens: int = 200,
7    temperature: float = 1.0,
8    top_k: int | None = None,
9    top_p: float | None = None,
10) -> str:
11    model.eval()
12    token_ids = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long)
13
14    for _ in range(max_new_tokens):
15        context = token_ids[:, -model.block_size :]
16        logits = model(context)
17        next_logits = logits[:, -1, :] / temperature   # (1, vocab_size)
18
19        if top_k is not None:
20            values, _ = torch.topk(next_logits, top_k)
21            threshold = values[:, -1, None]
22            next_logits = next_logits.masked_fill(next_logits < threshold, float("-inf"))
23
24        if top_p is not None:
25            sorted_logits, sorted_idx = torch.sort(next_logits, descending=True)
26            probs = F.softmax(sorted_logits, dim=-1)
27            cumulative = torch.cumsum(probs, dim=-1)
28            drop = cumulative > top_p
29            drop[:, 1:] = drop[:, :-1].clone()
30            drop[:, 0] = False
31            sorted_logits[drop] = float("-inf")
32            next_logits = torch.full_like(next_logits, float("-inf"))
33            next_logits.scatter_(1, sorted_idx, sorted_logits)
34
35        probs = F.softmax(next_logits, dim=-1)
36        next_id = torch.multinomial(probs, num_samples=1)
37        token_ids = torch.cat([token_ids, next_id], dim=1)
38
39    return tokenizer.decode(token_ids[0].tolist())

Temperature divides the logits before softmax, below 1.0 it sharpens the distribution toward the model's favorite tokens, above 1.0 it flattens it toward uniform, more surprises, more incoherence past a point. Top-k keeps only the k highest-scoring tokens and zeroes out the rest before sampling, a hard cutoff on how many candidates get considered. Top-p, nucleus sampling, is the adaptive version: keep the smallest set of top tokens whose cumulative probability crosses p, so a confident distribution keeps only one or two candidates and an uncertain one keeps many, the cutoff width adjusts itself instead of staying fixed. Run with a prompt and reasonable settings for both:

1print(generate(
2    model,
3    tokenizer,
4    prompt="the meaning of ",
5    max_new_tokens=120,
6    temperature=0.8,
7    top_k=40,
8    top_p=0.9,
9))

One honest limitation worth naming: generaterecomputes the entire forward pass, through all six blocks, on every single new token, feeding the whole growing sequence back in each time. Production models don't do this, they cache the key and value tensors from every previous step instead of recomputing them, exactly the KV cache the Context Windows article covered. At block_size= 128 the cost of not caching is trivial. At the 32,000-token context windows real deployments run, it's the difference between a response that streams back in seconds and one that doesn't. And if you wanted to know precisely how well this model is doing, not just eyeballing its output, the Evaluating LLMs article's answer is sitting right there in the training loop already: perplexity is just e raised to the average cross-entropy loss, computed on held-out text the model never trained on.

Read back over what just ran. A tokenizer turned raw text into integers. An embedding table turned those integers into vectors, and a positional table told the model where each one sat in the sequence. Multi-head attention let every token look at every earlier token and decide what mattered. Feedforward layers gave the model somewhere to actually transform what attention gathered. Residual connections let all of that stack six layers deep without the signal degrading, and layer norm kept the numbers flowing through it well-behaved. Cross-entropy loss measured exactly how wrong each prediction was, backpropagation traced that error back through every one of those layers to every parameter responsible for it, and gradient descent, by way of Adam, nudged each one to be less wrong next time. And sampling, temperature and top-k and top-p together, turned the trained model's output distribution back into the one thing this entire series has been building toward from its very first sentence: the next word.