I Built the Smallest LLM I Could to Understand How They Work

I use LLMs every day and had no real idea what happens inside one. So I built the smallest one I could: a transformer with roughly 108,000 parameters that learns to write Shakespeare one character at a time, in under 200 lines of PyTorch. This post walks through every piece, plus a few experiments worth running.

What we’re building

A language model does one thing: given some text, predict what comes next. Chat, code, and reasoning are all that one trick at enormous scale. Ours is the same idea, shrunk down:

• Data: Tiny Shakespeare (from Andrej Karpathy’s char-rnn project), about 1.1 million characters

• Tokens: single characters (65 unique ones)

• Model: 2 transformer blocks, 4 attention heads, 64-dimensional embeddings

• Size: 107,968 parameters (the script prints this number when it starts). GPT-3 has 175 billion, over a million times more.

• Hardware: a laptop CPU is enough; the script uses a GPU automatically if one is available

Step 1: Text becomes numbers

Neural networks only do math, so the first job is mapping each character to an integer.

chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]

The word “hello” becomes something like [46, 43, 50, 50, 53]; the exact ids depend on the vocabulary. Real LLMs use subword tokenizers (BPE), where a token is a chunk like “ing” or “ the”. That gives shorter sequences but is the same idea. Characters keep ours dead simple.

Training examples come from sliding windows. Take 64 characters as input, and the same window shifted by one character as the target. At every position the model is asked the same question: what comes next? One window gives 64 training examples at once.

Step 2: Embeddings, or what and where

Each id looks up a row in a table of 64 learned numbers. That row is the character’s embedding. A second table encodes position, because attention on its own has no sense of order. We add the two together.

x = self.tok_emb(idx) + self.pos_emb(torch.arange(T))

Step 3: Self-attention, the core idea

Every character produces three vectors: a query (what am I looking for?), a key (what do I contain?), and a value (what do I hand over if someone attends to me?). Comparing a query with every key gives a score. A softmax turns the scores into weights that sum to 1, and the output is a weighted average of the values.

att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))
att = att.masked_fill(self.mask[:T, :T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
y = att @ v

Two details matter here. The causal mask stops position t from seeing anything after t; without it the model could simply read the answer. The division by the square root of the head size keeps the scores from growing so large that the softmax saturates. With 4 heads, the model runs 4 independent attention patterns in parallel and merges them.

Step 4: The transformer block

Attention gathers information from other positions. A small MLP then processes each position on its own. Two tricks make stacks of these trainable: residual connections (add the input back to the output) and LayerNorm.

x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))

Stack two blocks, apply a final LayerNorm, and use a linear layer to turn each position’s vector into 65 scores, one per possible next character. That is the entire model.

Step 5: Training

The loss is cross-entropy: the penalty is the negative log of the probability the model gave to the true next character. Confident and right costs almost nothing; confident and wrong costs a lot.

A sanity check worth knowing. Before training, the model should know nothing, meaning it spreads probability evenly over all 65 characters. That gives a loss of ln(65) ≈ 4.17. If your first loss is far from that, something is wrong. The script prints ln(65) at startup so you can compare. As a cautionary example, a NumPy simulation of an untrained model with PyTorch’s default embedding initialization (standard deviation 1) and a tied output layer gave a starting loss of about 19. Initializing the weights with a tiny standard deviation (0.02) brought it back to 4.18.

The training loop is the same one used for every neural network:

x, y = get_batch('train')
_, loss = model(x, y)
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()

The clipping line caps how large any single update can be, which protects a fairly high learning rate (0.003) from the occasional bad batch.

When you run the script, it prints the training and validation loss every 250 steps. The first reading should sit close to 4.17, then fall steadily as the model learns.

Watch the gap between train and validation loss. If validation stops improving while training loss keeps falling, the model has started memorizing instead of learning.

Step 6: Generating text

Generation is a loop: feed in the text so far, get a probability for every possible next character, sample one, append it, repeat. Temperature controls the sampling. Lower values give safer, more repetitive text; higher values give wilder text.

One easy mistake: the script calls model.eval() before sampling, which switches dropout off. Leave it on and the samples come out noisier than the model deserves.

The script prints one sample before training and one when training finishes. Here is what to expect from each:

• Before training: random characters with no structure at all.

• After training: you should typically see speaker names in capitals followed by a colon, then short lines that read like Shakespeare without meaning much.

For an in-between sample, run ITERS=300 python tinygpt.py. At that point you should see word-shaped clumps, with spaces and line breaks in plausible places.

A 108k-parameter model learns the shape of the language long before it learns sense, so don’t expect the lines to mean anything.

The complete code

Here is the whole script. Save it as tinygpt.py, run pip install torch, then python tinygpt.py. It downloads the dataset on the first run. For a quick smoke test that finishes in about a minute, run ITERS=200 python tinygpt.py.

"""
tinygpt.py - a character-level GPT in a single file.

Run:  pip install torch
      python tinygpt.py

Trains on Tiny Shakespeare (downloaded automatically) and prints samples.
Quick smoke test:  ITERS=200 python tinygpt.py
"""
import math
import os
import urllib.request

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

# ------------------------------------------------------------------
# 1. Config: tiny on purpose. Change these and watch what happens.
# ------------------------------------------------------------------
block_size = 64    # context window: how many characters the model can see
n_embd = 64        # size of each token's vector
n_head = 4         # attention heads (n_embd must divide evenly by this)
n_layer = 2        # number of transformer blocks
dropout = 0.1
batch_size = 32
lr = 3e-3
max_iters = int(os.environ.get("ITERS", 3000))
eval_every = 250

device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
torch.manual_seed(1337)

# ------------------------------------------------------------------
# 2. Data + tokenizer. Our "tokenizer" is just: one character = one id.
# ------------------------------------------------------------------
URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
if not os.path.exists("input.txt"):
    try:
        urllib.request.urlretrieve(URL, "input.txt")
    except Exception as e:
        if os.path.exists("input.txt"):
            os.remove("input.txt")  # don't leave a half-downloaded file behind
        raise SystemExit(
            f"Could not download the dataset ({e}).\n"
            f"Download it manually from:\n  {URL}\n"
            "and save it as input.txt next to this script."
        )
with open("input.txt", encoding="utf-8") as f:
    text = f.read()

chars = sorted(set(text))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}


def encode(s):
    return [stoi[c] for c in s]


def decode(ids):
    return "".join(itos[i] for i in ids)


data = torch.tensor(encode(text), dtype=torch.long)
split = int(0.9 * len(data))
train_data, val_data = data[:split], data[split:]


def get_batch(which):
    d = train_data if which == "train" else val_data
    ix = torch.randint(len(d) - block_size - 1, (batch_size,))
    x = torch.stack([d[i : i + block_size] for i in ix])
    y = torch.stack([d[i + 1 : i + block_size + 1] for i in ix])  # shifted by 1
    return x.to(device), y.to(device)


# ------------------------------------------------------------------
# 3. The model
# ------------------------------------------------------------------
class CausalSelfAttention(nn.Module):
    """Each position looks at earlier positions and mixes their information."""

    def __init__(self):
        super().__init__()
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd)
        self.drop = nn.Dropout(dropout)
        # lower-triangular mask: position t may only see positions <= t
        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)
        # split into heads: (B, n_head, T, head_dim)
        q, k, v = (t.view(B, T, n_head, C // n_head).transpose(1, 2) for t in (q, k, v))

        att = (q @ k.transpose(-2, -1)) / math.sqrt(k.size(-1))  # similarity scores
        att = att.masked_fill(self.mask[:T, :T] == 0, float("-inf"))  # hide the future
        att = F.softmax(att, dim=-1)  # scores -> weights that sum to 1
        y = self.drop(att) @ v  # weighted average of values

        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-merge heads
        return self.drop(self.proj(y))


class MLP(nn.Module):
    """Per-position 'thinking' step after attention has gathered context."""

    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP()

    def forward(self, x):
        x = x + self.attn(self.ln1(x))  # residual connection
        x = x + self.mlp(self.ln2(x))  # residual connection
        return x


class TinyGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, n_embd)  # what token is it?
        self.pos_emb = nn.Embedding(block_size, n_embd)  # where is it?
        self.blocks = nn.Sequential(*[Block() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight  # weight tying: fewer params
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m):
        # Small init so the untrained model predicts ~uniformly (loss ~ ln(vocab)).
        # Without this, PyTorch's default N(0,1) embeddings + weight tying give
        # huge initial logits and a starting loss around 19 instead of ~4.2.
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.zeros_(m.bias)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.tok_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device))
        x = self.blocks(x)
        logits = self.head(self.ln_f(x))  # (B, T, vocab_size)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0):
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -block_size:])  # crop to context window
            logits = logits[:, -1, :] / temperature  # only the last position
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)  # sample next char
            idx = torch.cat([idx, nxt], dim=1)
        return idx


# ------------------------------------------------------------------
# 4. Training
# ------------------------------------------------------------------
@torch.no_grad()
def estimate_loss(model, iters=50):
    model.eval()  # turn dropout off while measuring
    out = {}
    for which in ("train", "val"):
        losses = torch.zeros(iters)
        for i in range(iters):
            x, y = get_batch(which)
            losses[i] = model(x, y)[1].item()
        out[which] = losses.mean().item()
    model.train()
    return out


def sample(model, prompt="\n", n=300, temperature=0.8):
    model.eval()  # dropout off, otherwise samples are noisier than they should be
    idx = torch.tensor([encode(prompt)], dtype=torch.long, device=device)
    out = decode(model.generate(idx, n, temperature)[0].tolist())
    model.train()
    return out


if __name__ == "__main__":
    model = TinyGPT().to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"device={device}  vocab={vocab_size}  params={n_params:,}")
    print(f"a model that knows nothing should start near ln({vocab_size}) = {math.log(vocab_size):.3f}")

    print("\n--- before training (random noise) ---")
    print(sample(model, n=150))

    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    print()
    for step in range(max_iters):
        if step % eval_every == 0:
            l = estimate_loss(model)
            print(f"step {step:5d}  train {l['train']:.3f}  val {l['val']:.3f}")
        x, y = get_batch("train")
        _, loss = model(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # guard against spikes
        opt.step()
    l = estimate_loss(model)
    print(f"step {max_iters:5d}  train {l['train']:.3f}  val {l['val']:.3f}  (final)")

    print("\n--- after training ---")
    print(sample(model, n=500))

Experiments

Change one thing at a time, rerun, and compare the final validation loss with the baseline. These experiments are a good place to start, with what each is likely to show:

The last row is the instructive one. Without the mask the loss should look fantastic, because the model is reading the answer from the future. But that model can’t actually generate text, since at generation time there is no future to read. A good loss means nothing if the setup is cheating.

How this scales to a real LLM

The architecture is the same. What changes is mostly quantity: subword tokens instead of characters, dozens of layers instead of two, embeddings thousands of numbers wide instead of 64, context windows of many thousands of tokens, and training data measured in trillions of tokens. Production assistants then get extra training on top (instruction tuning and learning from human feedback) to turn a text predictor into something that follows requests.

Takeaways

• A language model is a next-token predictor. The surprising capabilities come from scale.

• Attention is a weighted average where the weights are learned and depend on content.

• Cheap sanity checks, like the starting loss equaling ln(vocab size), catch bugs before you waste hours of training.

• Changing one variable at a time teaches more than reading about it.

Credits and sources

• Dataset: Tiny Shakespeare, about 40,000 lines from a variety of Shakespeare’s plays, comes from Andrej Karpathy’s char-rnn repository as input.txt. The script downloads it from there. It became widely known through his post The Unreasonable Effectiveness of Recurrent Neural Networks.

• Suggested citation for the data: Karpathy, Andrej. char-rnn. 2015. https://github.com/karpathy/char-rnn

• Code approach: the structure follows the style of Karpathy’s nanoGPT and his video Let’s build GPT: from scratch, in code, spelled out. The script here is a smaller, heavily commented take on the same idea.

• Architecture: the transformer comes from Vaswani et al., Attention Is All You Need (2017).