Setup and Data
pip install torch
# Training text: any corpus. For a fast demo, use a public
# domain book or your own writing. ~1MB is plenty.
text = open("training.txt", encoding="utf-8").read()
print(f"{len(text)} characters")
Step 1: The Tokenizer
# Character-level tokenizer — simplest possible.
# Real LLMs use BPE (see the tokenizer guide); the concept is the same.
chars = sorted(set(text))
stoi = {ch: i for i, ch in enumerate(chars)} # char → id
itos = {i: ch for i, ch in enumerate(chars)} # id → char
VOCAB = len(chars)
print(f"vocab: {VOCAB} chars") # ~65 for English text
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = torch.tensor(encode(text), dtype=torch.long)
train, val = data[:int(0.9 * len(data))], data[int(0.9 * len(data)):]
The full BPE story is in How Tokenizers Work — swap this encoder for a real BPE tokenizer as an exercise.
Step 2: The Dataset (Input/Output Pairs)
BLOCK = 64 # context length
def get_batch(split, batch_size=32):
data_src = train if split == "train" else val
ix = torch.randint(len(data_src) - BLOCK, (batch_size,))
x = torch.stack([data_src[i:i + BLOCK] for i in ix])
y = torch.stack([data_src[i + 1:i + BLOCK + 1] for i in ix])
return x, y # y = x shifted by one = next-token targets
This is the whole learning task in one function: given a window, predict the next token. The shift (y = x + 1) is the same supervision signal every LLM is trained on.
Step 3: The Model
import torch.nn as nn
from torch.nn import functional as F
class Block(nn.Module):
def __init__(self, n_embd, n_head):
super().__init__()
head_size = n_embd // n_head
self.sa = nn.MultiheadAttention(n_embd, n_head,
batch_first=True)
self.ff = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd), nn.GELU(),
nn.Linear(4 * n_embd, n_embd))
self.ln1, self.ln2 = nn.LayerNorm(n_embd), nn.LayerNorm(n_embd)
def forward(self, x):
attn, _ = self.sa(x, x, x, need_weights=False) # causal below
x = x + attn
x = x + self.ff(self.ln2(x))
return x
class MiniGPT(nn.Module):
def __init__(self, vocab, n_embd=128, n_head=4, n_layer=4):
super().__init__()
self.tok = nn.Embedding(vocab, n_embd)
self.pos = nn.Embedding(BLOCK, n_embd)
self.blocks = nn.Sequential(
*[Block(n_embd, n_head) for _ in range(n_layer)])
self.ln = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab)
def forward(self, idx):
B, T = idx.shape
x = self.tok(idx) + self.pos(torch.arange(T))
x = self.blocks(x)
return self.head(self.ln(x)) # logits (B, T, vocab)
Note the two embedding tables: token embeddings + position embeddings — the position handling the architecture needs (covered in the transformer walkthrough). The attention is masked to be causal — a token must never see the future.
Step 4: Train and Sample
model = MiniGPT(VOCAB)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(5000):
x, y = get_batch("train")
logits = model(x)
loss = F.cross_entropy(
logits.view(-1, VOCAB), y.view(-1))
opt.zero_grad(); loss.backward(); opt.step()
if step % 500 == 0:
print(f"step {step}: loss {loss.item():.4f}")
@torch.no_grad()
def sample(seed="\n", n=300):
model.eval()
idx = torch.tensor(encode(seed)).unsqueeze(0)
for _ in range(n):
logits = model(idx[:, -BLOCK:])
p = F.softmax(logits[0, -1] / 0.8, dim=-1)
idx = torch.cat([idx, torch.multinomial(p, 1)], dim=1)
return decode(idx[0].tolist())
print(sample()) # your mini GPT "writing"
Expected: loss starts near ln(65) ≈ 4.2 and falls toward 1.0–1.5 on a decent corpus. Temperature (/0.8) controls randomness at sampling — the same knob every local runner exposes.
Step 5: Fine-Tune on a New Domain
# Continue training on different text — the model adapts its style.
new_text = open("domain.txt", encoding="utf-8").read()
# ... rebuild tokenizer ONLY if the new text has new characters
# (for same-vocab domains, just swap the data and continue):
data2 = torch.tensor(encode(new_text), dtype=torch.long)
for step in range(2000):
x, y = get_batch("train", data_src=data2)
# ... same training lines as above ...
# The model now writes in the new domain's style.
print(sample())
This is fine-tuning in miniature — and it's the same principle (with LoRA, covered in the LoRA guide) that adapts real local LLMs. The RAG vs fine-tuning decision applies here too: fine-tuning changes style, RAG supplies facts.
What to Do Next With Your Mini GPT
- Visualize attention heads — what relationships did each head learn? (The walkthrough explains what to expect.)
- Measure KV cache growth — log memory vs context length and watch the cache math come alive.
- Quantize it — round weights to 8/4 bits and watch quality degrade — a hands-on quantization lesson.
- Swap in BPE — the tokenizer guide shows how; your loss drops immediately.
Frequently Asked Questions (FAQ)
Can I really train a GPT from scratch?
Yes — a small GPT (1–10M parameters) trains in minutes on a laptop CPU, or seconds on a GPU. It won't be useful like a real LLM (those train on trillions of tokens), but it learns to produce coherent text in its domain, and building it teaches you exactly how real LLMs work under the hood.
How much code does a mini GPT need?
About 150–250 lines of PyTorch: a tokenizer (~20 lines), the transformer blocks (~80 lines), and a training loop (~50 lines). The tutorial in this post is complete and runnable.
What's the difference between a mini GPT and a real LLM?
Scale, not architecture: real LLMs use the same transformer blocks with BPE tokenizers, billions of parameters, and training on trillions of tokens across thousands of GPUs. A mini GPT is the same architecture at 1–10M parameters on a laptop — same ideas, learnable size.
Should I use a character or BPE tokenizer for a mini GPT?
Character tokenizer for learning (simplest, good enough for small models on small data); BPE for anything approaching production. The tokenizer guide in this blog explains BPE in depth — swapping the character encoder for a BPE one is a clean exercise.
How do I know my mini GPT is learning?
Watch the training loss: it should start around ln(vocab_size) (~4 for a 65-character vocab) and fall steadily. Then sample text and check it looks like the training domain. Loss in the 1.0–1.5 range on a small corpus means it's learned real structure, not memorized.
What can I do with a trained mini GPT?
Learn, mostly — but also: fine-tune it on a new domain (the tutorial covers this), ablate attention heads to see what they learn, measure how quantization would affect it, and build intuition for every concept (KV cache, MoE, speculation) that this blog covers at real-LLM scale.
Sources & Further Reading
- How Transformers Actually Work: From-Scratch Walkthrough
- How Tokenizers Work: BPE Explained Simply
- Fine-Tuning a Local LLM: LoRA for Beginners
- KV Cache Quantization: What It Is and Why It Matters
- Q4_K_M vs Q8_0: Which Quantization Should You Use?
- nanoGPT (Andrej Karpathy) — the production-grade reference