PyTorch Zero to Hero 10: Build a Small GPT-Style Language Model From Scratch
Build a Small GPT-Style Language Model From Scratch in PyTorch
This is the final post in the PyTorch: Zero to Hero series.
We have spent the previous posts learning the machinery underneath PyTorch:
- tensors and shapes;
- autograd;
- manual neural networks;
nn.Moduleand parameter registration;- DataLoader performance;
- convolutional networks;
- attention and masks;
- training failures;
- CUDA performance and
torch.compile.
Now we put it together.
The goal is not to download a pretrained model.
The goal is to build a small autoregressive language model ourselves so that every major tensor transformation is visible.
By the end, we will have this pipeline:
flowchart TD
A[Text] --> B[Token IDs]
B --> C[Token Embeddings + Positional Embeddings]
C --> D[Transformer Blocks]
D --> E[Final LayerNorm]
E --> F[Vocabulary Logits]
F --> G[Cross Entropy Loss]
G --> H[Backpropagation]
H --> I[AdamW Update]
I --> J[Trained Model]
J --> K[Autoregressive Generation]
text
↓
token ids
↓
token embeddings + positional embeddings
↓
transformer blocks
↓
final layer norm
↓
vocabulary logits
↓
cross entropy
↓
backpropagation
↓
AdamW
↓
trained language model
↓
autoregressive generation
This is not intended to compete with frontier models.
It is intended to remove the magic.
1. What are we actually building?
We are building an autoregressive language model.
Given tokens:
The cat sat on the
we want the model to predict a probability distribution for the next token.
Then we append one sampled token:
The cat sat on the mat
and ask again.
Training therefore looks like next-token prediction.
If our token sequence is:
[10, 42, 7, 13, 99]
then a training pair can be:
input: [10, 42, 7, 13]
target: [42, 7, 13, 99]
Every position predicts the next token.
That simple shift is the entire training objective.
2. A deliberately small model
We will keep the default model small enough to experiment with locally.
from dataclasses import dataclass
@dataclass
class GPTConfig:
vocab_size: int
block_size: int = 128
n_layer: int = 4
n_head: int = 4
n_embd: int = 256
dropout: float = 0.1
The important relationship is:
n_embd % n_head == 0
because each attention head receives:
head_dim = n_embd // n_head
With the defaults:
n_embd = 256
n_head = 4
head_dim = 64
The tensor shape that will dominate the attention implementation is:
(B, T, C)
where:
B = batch size
T = sequence length
C = embedding dimension
Inside attention it becomes:
(B, H, T, D)
where:
H = number of heads
D = head dimension
C = H * D
If you understand those two shapes, much of transformer code becomes straightforward.
Part I — Data
3. Start with text
For a real project you would usually use a proper tokenizer.
For this article, we will first build a character-level tokenizer because it exposes the mechanism clearly and requires no additional dependency.
Suppose we have:
text = open("input.txt", "r", encoding="utf-8").read()
Inspect it immediately:
print("characters:", len(text))
print("preview:", repr(text[:200]))
Do not skip this.
A training pipeline that accidentally reads an empty file, the wrong encoding, duplicated data, or HTML instead of text can happily train for hours.
Runtime evidence first.
4. Build a character tokenizer
chars = sorted(set(text))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}
def encode(s: str) -> list[int]:
return [stoi[ch] for ch in s]
def decode(ids: list[int]) -> str:
return "".join(itos[i] for i in ids)
Test round-trip correctness:
sample = text[:100]
ids = encode(sample)
restored = decode(ids)
assert restored == sample
If your tokenizer cannot round-trip the input, stop there.
Do not train the model and hope it sorts itself out.
5. Convert text into one long tensor
import torch
data = torch.tensor(encode(text), dtype=torch.long)
print(data.shape)
print(data[:20])
Token IDs must be integer tensors because we will use them as indices into an embedding table.
A common failure is accidentally converting them to floating-point tensors.
Check explicitly:
assert data.dtype == torch.long
6. Train/validation split
split = int(0.9 * len(data))
train_data = data[:split]
val_data = data[split:]
print("train tokens:", len(train_data))
print("val tokens:", len(val_data))
This gives us a simple holdout set.
Do not report training loss alone.
A model can memorize increasingly well while getting worse on unseen text.
7. Build batches
We need random windows of length block_size.
def get_batch(source: torch.Tensor, batch_size: int, block_size: int, device: str):
max_start = len(source) - block_size - 1
starts = torch.randint(
low=0,
high=max_start,
size=(batch_size,),
)
x = torch.stack([
source[i : i + block_size]
for i in starts
])
y = torch.stack([
source[i + 1 : i + block_size + 1]
for i in starts
])
return x.to(device), y.to(device)
Inspect one batch:
x, y = get_batch(
train_data,
batch_size=4,
block_size=16,
device="cpu",
)
print("x:", x.shape) # (B, T)
print("y:", y.shape) # (B, T)
print(x[0])
print(y[0])
The contract is:
x: (B, T)
y: (B, T)
And this relationship must hold:
assert torch.equal(x[:, 1:], y[:, :-1])
That is one of the most useful tests in the entire post.
It proves that the target really is the next-token shift.
Part II — Embeddings
8. Token embeddings
A token ID is just an integer.
The model needs a vector representation.
import torch.nn as nn
token_embedding = nn.Embedding(vocab_size, 256)
If:
input ids: (B, T)
then:
x = token_embedding(x)
gives:
(B, T, C)
where C = 256 here.
Inspect it:
print(x.shape)
The embedding table itself has shape:
(vocab_size, n_embd)
Each token ID selects one row.
9. Positional embeddings
Self-attention by itself does not know whether a token is first, fifth or fiftieth.
For this small model, we will use learned positional embeddings.
position_embedding = nn.Embedding(128, 256)
For a sequence of length T:
positions = torch.arange(T)
pos = position_embedding(positions)
pos has shape:
(T, C)
while token embeddings have shape:
(B, T, C)
Broadcasting lets us add them:
x = token_embeddings + position_embeddings
The position tensor broadcasts across the batch dimension.
Result:
(B, T, C)
This is exactly the kind of broadcasting rule from Step 01 that becomes useful later.
Part III — Attention
10. A causal self-attention module
We already built attention carefully in Step 07.
Now we package it into the model.
import torch
import torch.nn as nn
import torch.nn.functional as F
class CausalSelfAttention(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
assert config.n_embd % config.n_head == 0
self.n_head = config.n_head
self.n_embd = config.n_embd
self.head_dim = config.n_embd // config.n_head
self.dropout = config.dropout
self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
self.resid_dropout = nn.Dropout(config.dropout)
def forward(self, x):
B, T, C = x.shape
qkv = self.qkv(x)
q, k, v = qkv.chunk(3, dim=-1)
q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
y = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=None,
dropout_p=self.dropout if self.training else 0.0,
is_causal=True,
)
y = y.transpose(1, 2).contiguous().view(B, T, C)
y = self.proj(y)
y = self.resid_dropout(y)
return y
Shape trace:
flowchart LR
A["x: (B,T,C)"] --> B["qkv: (B,T,3C)"]
B --> C["q/k/v: (B,T,C)"]
C --> D["reshape: (B,T,H,D)"]
D --> E["transpose: (B,H,T,D)"]
E --> F["attention output: (B,H,T,D)"]
F --> G["transpose: (B,T,H,D)"]
G --> H["merge heads: (B,T,C)"]
H --> I["projection: (B,T,C)"]
x (B, T, C)
qkv (B, T, 3C)
q/k/v (B, T, C)
reshape (B, T, H, D)
transpose (B, H, T, D)
attention output (B, H, T, D)
transpose (B, T, H, D)
merge heads (B, T, C)
projection (B, T, C)
Keep that beside you when debugging transformer code.
11. Why causal attention matters
During training, the target at position t is the token at t + 1.
If token t could attend to future tokens, training would leak the answer.
Causal attention ensures position t can only attend to positions:
0 ... t
not:
t + 1 ... T - 1
Using:
is_causal=True
makes this intent explicit.
12. Test attention before building the full model
config = GPTConfig(vocab_size=vocab_size)
attn = CausalSelfAttention(config)
x = torch.randn(2, 16, config.n_embd)
y = attn(x)
print(y.shape) # (2, 16, 256)
assert y.shape == x.shape
Backward smoke test:
loss = y.square().mean()
loss.backward()
for name, p in attn.named_parameters():
assert p.grad is not None, name
If this fails, do not proceed to the full transformer.
Debug the smallest failing component.
Part IV — The feed-forward network
13. MLP block
A transformer block is not only attention.
It also contains a per-token feed-forward network.
class MLP(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
hidden = 4 * config.n_embd
self.net = nn.Sequential(
nn.Linear(config.n_embd, hidden),
nn.GELU(),
nn.Linear(hidden, config.n_embd),
nn.Dropout(config.dropout),
)
def forward(self, x):
return self.net(x)
Input:
(B, T, C)
Output:
(B, T, C)
The MLP operates independently at every sequence position.
The sequence dimension remains intact.
Part V — Transformer block
14. Residual connections and LayerNorm
We will use a pre-normalization block:
class Block(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.ln1 = nn.LayerNorm(config.n_embd)
self.attn = CausalSelfAttention(config)
self.ln2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
There are two residual paths:
flowchart TD
X[x] --> ADD1((+))
X --> LN1[LayerNorm]
LN1 --> ATTN[Attention]
ATTN --> ADD1
ADD1 --> Y[y]
Y --> ADD2((+))
Y --> LN2[LayerNorm]
LN2 --> MLP[MLP]
MLP --> ADD2
ADD2 --> OUT[output]
Every block preserves:
(B, T, C)
That invariant makes stacking blocks easy.
15. Unit-test the block
block = Block(config)
x = torch.randn(2, 16, config.n_embd)
y = block(x)
assert y.shape == x.shape
assert torch.isfinite(y).all()
Gradient smoke test:
y.mean().backward()
missing = [
name
for name, p in block.named_parameters()
if p.requires_grad and p.grad is None
]
assert not missing, missing
Again: fail early, locally and loudly.
Part VI — The complete GPT-style model
16. Assemble the model
class TinyGPT(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.config = config
self.token_embedding = nn.Embedding(
config.vocab_size,
config.n_embd,
)
self.position_embedding = nn.Embedding(
config.block_size,
config.n_embd,
)
self.dropout = nn.Dropout(config.dropout)
self.blocks = nn.ModuleList([
Block(config)
for _ in range(config.n_layer)
])
self.ln_f = nn.LayerNorm(config.n_embd)
self.lm_head = nn.Linear(
config.n_embd,
config.vocab_size,
bias=False,
)
def forward(self, idx, targets=None):
B, T = idx.shape
if T > self.config.block_size:
raise ValueError(
f"sequence length {T} exceeds block_size "
f"{self.config.block_size}"
)
positions = torch.arange(T, device=idx.device)
tok = self.token_embedding(idx)
pos = self.position_embedding(positions)
x = self.dropout(tok + pos)
for block in self.blocks:
x = block(x)
x = self.ln_f(x)
logits = self.lm_head(x)
loss = None
if targets is not None:
loss = F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
targets.reshape(-1),
)
return logits, loss
This is the central object of the article.
Everything we learned in earlier posts appears here.
17. Shape contract for the full model
Input:
idx: (B, T)
After token embedding:
(B, T, C)
After every transformer block:
(B, T, C)
Vocabulary logits:
(B, T, V)
where V = vocab_size.
Cross entropy receives flattened tensors:
logits: (B*T, V)
targets: (B*T)
That is the loss contract.
18. Test the full forward pass
model = TinyGPT(config)
x, y = get_batch(
train_data,
batch_size=4,
block_size=config.block_size,
device="cpu",
)
logits, loss = model(x, y)
print("logits:", logits.shape)
print("loss:", loss.item())
assert logits.shape == (
x.size(0),
x.size(1),
config.vocab_size,
)
assert torch.isfinite(loss)
At initialization, the model is effectively guessing.
A useful rough sanity check is:
import math
expected_random_loss = math.log(config.vocab_size)
print(expected_random_loss)
The initial loss should usually be in the same broad region.
If it is wildly different, investigate before training.
Part VII — Parameter accounting
19. Count parameters
def count_parameters(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
total, trainable = count_parameters(model)
print(f"total parameters: {total:,}")
print(f"trainable parameters: {trainable:,}")
Parameter count should be treated as basic model metadata.
If an LLM generated your architecture, this is one of the first things worth checking.
A missing ModuleList, accidental duplicated layer or unexpected projection can change the count dramatically.
20. Inspect the largest tensors
def largest_parameters(model, n=10):
rows = []
for name, p in model.named_parameters():
rows.append((p.numel(), name, tuple(p.shape)))
rows.sort(reverse=True)
for numel, name, shape in rows[:n]:
print(f"{numel:>12,} {name:<40} {shape}")
largest_parameters(model)
This quickly shows where the model capacity lives.
Part VIII — Training
21. Choose the device
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
print("device:", device)
Move the model:
model = model.to(device)
And remember:
model parameters and input tensors must be on the same device
22. AdamW optimizer
optimizer = torch.optim.AdamW(
model.parameters(),
lr=3e-4,
weight_decay=0.1,
)
Before training, verify optimizer membership if you have dynamically replaced parameters or submodules.
def optimizer_parameter_ids(optimizer):
return {
id(p)
for group in optimizer.param_groups
for p in group["params"]
}
opt_ids = optimizer_parameter_ids(optimizer)
missing = [
name
for name, p in model.named_parameters()
if p.requires_grad and id(p) not in opt_ids
]
assert not missing, missing
That is a direct carry-over from Step 04 and Step 08.
23. Estimate validation loss
@torch.no_grad()
def estimate_loss(
model,
train_data,
val_data,
batch_size,
block_size,
device,
eval_batches=50,
):
model.eval()
out = {}
for split_name, source in [
("train", train_data),
("val", val_data),
]:
losses = []
for _ in range(eval_batches):
x, y = get_batch(
source,
batch_size,
block_size,
device,
)
_, loss = model(x, y)
losses.append(loss.item())
out[split_name] = sum(losses) / len(losses)
model.train()
return out
Do not confuse:
model.eval()
with:
torch.no_grad()
They solve different problems.
24. A minimal training loop
batch_size = 32
max_steps = 5000
eval_interval = 250
model.train()
for step in range(max_steps):
if step % eval_interval == 0:
losses = estimate_loss(
model,
train_data,
val_data,
batch_size,
config.block_size,
device,
)
print(
f"step={step:05d} "
f"train={losses['train']:.4f} "
f"val={losses['val']:.4f}"
)
x, y = get_batch(
train_data,
batch_size,
config.block_size,
device,
)
optimizer.zero_grad(set_to_none=True)
_, loss = model(x, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
That is enough to train the model.
But for a serious workflow we should instrument more than the loss.
We can quickly visualize the training progress by recording the loss at every step:
import matplotlib.pyplot as plt
losses_log = []
for step in range(max_steps):
# ... training steps as above ...
losses_log.append(loss.item())
plt.plot(losses_log)
plt.xlabel('Step')
plt.ylabel('Loss')
plt.title('Training loss over steps')
plt.grid(True)
plt.show()
Part IX — Prove that training is actually happening
25. Check parameter updates
Take a snapshot before a step:
before = {
name: p.detach().clone()
for name, p in model.named_parameters()
}
Run one optimization step.
Then:
with torch.no_grad():
changed = []
for name, p in model.named_parameters():
delta = (p - before[name]).abs().max().item()
if delta > 0:
changed.append(name)
print("changed parameters:", len(changed))
If the loss exists but zero parameters change, the training loop is not training.
That fact is more useful than twenty guesses about hyperparameters.
26. Gradient health
def gradient_report(model):
rows = []
for name, p in model.named_parameters():
if p.grad is None:
rows.append((name, None, None))
continue
grad = p.grad.detach()
rows.append((
name,
grad.norm().item(),
torch.isfinite(grad).all().item(),
))
return rows
Print suspicious entries:
for name, norm, finite in gradient_report(model):
if norm is None:
print("NO GRAD", name)
elif not finite:
print("NON-FINITE", name, norm)
This is why the series spent so much time on debugging.
Once the model becomes a few thousand lines of generated and library code, observable runtime state is what lets you cut through the abstraction.
Part X — Mixed precision
27. Optional CUDA AMP
For CUDA training, mixed precision can improve throughput and reduce memory use depending on hardware and workload.
The important point is not to assume that it helps.
Measure it.
A compact pattern:
use_amp = device == "cuda"
scaler = torch.amp.GradScaler(
"cuda",
enabled=use_amp,
)
for step in range(max_steps):
x, y = get_batch(
train_data,
batch_size,
config.block_size,
device,
)
optimizer.zero_grad(set_to_none=True)
with torch.autocast(
device_type="cuda",
dtype=torch.bfloat16,
enabled=use_amp,
):
_, loss = model(x, y)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
scaler.step(optimizer)
scaler.update()
If your GPU does not support the dtype/workload well, benchmark another configuration.
Do not cargo-cult AMP settings.
Part XI — torch.compile
28. Optional compilation
Once the eager version is correct:
compiled_model = torch.compile(model)
Then benchmark steady-state throughput.
Do not include the first compilation step in your steady-state measurement.
Do not assume compilation helped simply because the line executed successfully.
The performance post in Step 09 already established the rule:
Optimization is a hypothesis until the benchmark confirms it.
Part XII — Checkpointing
29. Save everything required to resume
def save_checkpoint(
path,
model,
optimizer,
step,
config,
stoi,
itos,
):
torch.save(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step,
"config": config.__dict__,
"stoi": stoi,
"itos": itos,
},
path,
)
A useful checkpoint contains more than weights.
You want enough information to reconstruct:
architecture
vocabulary
optimizer state
training step
Otherwise you have a weight file, not necessarily a resumable training state.
30. Load a checkpoint
def load_checkpoint(path, device):
checkpoint = torch.load(
path,
map_location=device,
weights_only=False,
)
config = GPTConfig(**checkpoint["config"])
model = TinyGPT(config).to(device)
model.load_state_dict(checkpoint["model"])
optimizer = torch.optim.AdamW(
model.parameters(),
lr=3e-4,
weight_decay=0.1,
)
optimizer.load_state_dict(checkpoint["optimizer"])
return (
model,
optimizer,
checkpoint["step"],
checkpoint["stoi"],
checkpoint["itos"],
)
Production code should treat checkpoint formats as schemas.
Version them if you expect the architecture to evolve.
Part XIII — Generation
31. Greedy generation first
Before temperature and top-k, make the simplest generator work.
@torch.no_grad()
def generate_greedy(model, idx, max_new_tokens):
model.eval()
for _ in range(max_new_tokens):
idx_cond = idx[:, -model.config.block_size :]
logits, _ = model(idx_cond)
next_logits = logits[:, -1, :]
next_token = torch.argmax(
next_logits,
dim=-1,
keepdim=True,
)
idx = torch.cat((idx, next_token), dim=1)
return idx
Why greedy first?
Because it removes sampling as a source of bugs.
If generation is structurally broken, random sampling can make diagnosis harder.
32. Temperature
Temperature rescales logits before softmax:
next_logits = next_logits / temperature
Lower temperature sharpens the distribution.
Higher temperature flattens it.
A safe guard:
if temperature <= 0:
raise ValueError("temperature must be > 0")
33. Top-k sampling
Top-k keeps only the k largest logits.
def apply_top_k(logits, k):
if k is None:
return logits
k = min(k, logits.size(-1))
values, _ = torch.topk(logits, k)
threshold = values[:, [-1]]
return logits.masked_fill(
logits < threshold,
float("-inf"),
)
Then sample from the remaining distribution.
34. Full sampling function
@torch.no_grad()
def generate(
model,
idx,
max_new_tokens,
temperature=1.0,
top_k=None,
):
if temperature <= 0:
raise ValueError("temperature must be > 0")
model.eval()
for _ in range(max_new_tokens):
idx_cond = idx[:, -model.config.block_size :]
logits, _ = model(idx_cond)
logits = logits[:, -1, :]
logits = logits / temperature
logits = apply_top_k(logits, top_k)
probs = F.softmax(logits, dim=-1)
if not torch.isfinite(probs).all():
raise RuntimeError("non-finite sampling probabilities")
next_token = torch.multinomial(
probs,
num_samples=1,
)
idx = torch.cat(
(idx, next_token),
dim=1,
)
return idx
35. Generate text
Start with a prompt:
prompt = "The "
prompt_ids = encode(prompt)
idx = torch.tensor(
[prompt_ids],
dtype=torch.long,
device=device,
)
Generate:
out = generate(
model,
idx,
max_new_tokens=300,
temperature=0.8,
top_k=40,
)
Decode:
generated = decode(out[0].tolist())
print(generated)
And there it is.
A language model you built yourself.
Part XIV — Why generation gets slower
36. The naive implementation recomputes the prefix
Notice this line:
idx_cond = idx[:, -model.config.block_size :]
At every generation step we run the entire visible context through the model again.
That is simple and correct.
It is not efficient.
Production autoregressive inference commonly caches attention keys and values so previous positions do not need to be recomputed in full.
That optimization is deliberately outside the minimum model because correctness comes first.
The important lesson is to understand the baseline before optimizing it.
Part XV — A complete runnable implementation
37. Full model code
The following version puts the main pieces together in one file.
from __future__ import annotations
from dataclasses import dataclass
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
@dataclass
class GPTConfig:
vocab_size: int
block_size: int = 128
n_layer: int = 4
n_head: int = 4
n_embd: int = 256
dropout: float = 0.1
class CausalSelfAttention(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
if config.n_embd % config.n_head != 0:
raise ValueError(
"n_embd must be divisible by n_head"
)
self.n_head = config.n_head
self.n_embd = config.n_embd
self.head_dim = config.n_embd // config.n_head
self.dropout = config.dropout
self.qkv = nn.Linear(
config.n_embd,
3 * config.n_embd,
bias=False,
)
self.proj = nn.Linear(
config.n_embd,
config.n_embd,
bias=False,
)
self.resid_dropout = nn.Dropout(
config.dropout
)
def forward(self, x):
B, T, C = x.shape
qkv = self.qkv(x)
q, k, v = qkv.chunk(3, dim=-1)
q = (
q.view(B, T, self.n_head, self.head_dim)
.transpose(1, 2)
)
k = (
k.view(B, T, self.n_head, self.head_dim)
.transpose(1, 2)
)
v = (
v.view(B, T, self.n_head, self.head_dim)
.transpose(1, 2)
)
y = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=None,
dropout_p=(
self.dropout
if self.training
else 0.0
),
is_causal=True,
)
y = (
y.transpose(1, 2)
.contiguous()
.view(B, T, C)
)
y = self.proj(y)
y = self.resid_dropout(y)
return y
class MLP(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
hidden = 4 * config.n_embd
self.net = nn.Sequential(
nn.Linear(config.n_embd, hidden),
nn.GELU(),
nn.Linear(hidden, config.n_embd),
nn.Dropout(config.dropout),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.ln1 = nn.LayerNorm(config.n_embd)
self.attn = CausalSelfAttention(config)
self.ln2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
class TinyGPT(nn.Module):
def __init__(self, config: GPTConfig):
super().__init__()
self.config = config
self.token_embedding = nn.Embedding(
config.vocab_size,
config.n_embd,
)
self.position_embedding = nn.Embedding(
config.block_size,
config.n_embd,
)
self.dropout = nn.Dropout(config.dropout)
self.blocks = nn.ModuleList([
Block(config)
for _ in range(config.n_layer)
])
self.ln_f = nn.LayerNorm(config.n_embd)
self.lm_head = nn.Linear(
config.n_embd,
config.vocab_size,
bias=False,
)
def forward(self, idx, targets=None):
B, T = idx.shape
if T > self.config.block_size:
raise ValueError(
f"sequence length {T} > "
f"block_size {self.config.block_size}"
)
positions = torch.arange(
T,
device=idx.device,
)
tok = self.token_embedding(idx)
pos = self.position_embedding(positions)
x = self.dropout(tok + pos)
for block in self.blocks:
x = block(x)
x = self.ln_f(x)
logits = self.lm_head(x)
loss = None
if targets is not None:
loss = F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
targets.reshape(-1),
)
return logits, loss
@torch.no_grad()
def generate(
self,
idx,
max_new_tokens,
temperature=1.0,
top_k=None,
):
if temperature <= 0:
raise ValueError(
"temperature must be > 0"
)
was_training = self.training
self.eval()
try:
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.config.block_size :]
logits, _ = self(idx_cond)
logits = logits[:, -1, :]
logits = logits / temperature
if top_k is not None:
k = min(top_k, logits.size(-1))
values, _ = torch.topk(logits, k)
cutoff = values[:, [-1]]
logits = logits.masked_fill(
logits < cutoff,
float("-inf"),
)
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(
probs,
num_samples=1,
)
idx = torch.cat(
(idx, next_token),
dim=1,
)
return idx
finally:
self.train(was_training)
This is already enough to train and generate.
Part XVI — Training script
38. End-to-end training example
import torch
# --------------------------------------------------
# data
# --------------------------------------------------
text = open(
"input.txt",
"r",
encoding="utf-8",
).read()
chars = sorted(set(text))
stoi = {
ch: i
for i, ch in enumerate(chars)
}
itos = {
i: ch
for ch, i in stoi.items()
}
def encode(s):
return [stoi[ch] for ch 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 = data[:split]
val_data = data[split:]
# --------------------------------------------------
# configuration
# --------------------------------------------------
device = (
"cuda"
if torch.cuda.is_available()
else "cpu"
)
config = GPTConfig(
vocab_size=len(chars),
block_size=128,
n_layer=4,
n_head=4,
n_embd=256,
dropout=0.1,
)
batch_size = 32
learning_rate = 3e-4
max_steps = 5000
eval_interval = 250
# --------------------------------------------------
# batches
# --------------------------------------------------
def get_batch(source):
max_start = (
len(source)
- config.block_size
- 1
)
starts = torch.randint(
0,
max_start,
(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)
# --------------------------------------------------
# model
# --------------------------------------------------
model = TinyGPT(config).to(device)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate,
weight_decay=0.1,
)
# --------------------------------------------------
# evaluation
# --------------------------------------------------
@torch.no_grad()
def estimate_loss(eval_batches=50):
model.eval()
result = {}
for name, source in [
("train", train_data),
("val", val_data),
]:
losses = []
for _ in range(eval_batches):
x, y = get_batch(source)
_, loss = model(x, y)
losses.append(loss.item())
result[name] = (
sum(losses)
/ len(losses)
)
model.train()
return result
# --------------------------------------------------
# training
# --------------------------------------------------
model.train()
for step in range(max_steps):
if step % eval_interval == 0:
losses = estimate_loss()
print(
f"step={step:05d} "
f"train={losses['train']:.4f} "
f"val={losses['val']:.4f}"
)
x, y = get_batch(train_data)
optimizer.zero_grad(set_to_none=True)
_, loss = model(x, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
1.0,
)
optimizer.step()
# --------------------------------------------------
# generation
# --------------------------------------------------
prompt = "The "
idx = torch.tensor(
[encode(prompt)],
dtype=torch.long,
device=device,
)
out = model.generate(
idx,
max_new_tokens=500,
temperature=0.8,
top_k=40,
)
print(decode(out[0].tolist()))
That is a complete miniature language-model training program.
Part XVII — Debugging the model
39. The model runs but the loss does not decrease
Use the same forensic order from Step 08.
Do not immediately change ten hyperparameters.
Check:
1. data
2. targets
3. baseline
4. forward pass
5. loss
6. gradients
7. optimizer membership
8. parameter updates
9. tiny-batch overfit
10. only then tune
40. Verify next-token alignment
x, y = get_batch(train_data)
assert torch.equal(
x[:, 1:].cpu(),
y[:, :-1].cpu(),
)
If this assertion fails, your model is learning the wrong target relationship.
Fix that before anything else.
41. Verify token ranges
assert x.min() >= 0
assert y.min() >= 0
assert x.max() < config.vocab_size
assert y.max() < config.vocab_size
Embedding lookup errors often come from corrupted or mismatched tokenizer IDs.
42. Verify logits shape
logits, loss = model(x, y)
assert logits.shape == (
x.size(0),
x.size(1),
config.vocab_size,
)
Do not rely on visual inspection.
Make the contract executable.
43. Verify finite values
assert torch.isfinite(logits).all()
assert torch.isfinite(loss)
Then after backward:
loss.backward()
for name, p in model.named_parameters():
if p.grad is None:
continue
assert torch.isfinite(p.grad).all(), name
The first non-finite tensor is much more useful than the final NaN loss.
44. Overfit one tiny batch
This remains one of the strongest debugging tests for a training system.
x, y = get_batch(train_data)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
for step in range(500):
optimizer.zero_grad(set_to_none=True)
_, loss = model(x, y)
loss.backward()
optimizer.step()
if step % 50 == 0:
print(step, loss.item())
If a sufficiently expressive model cannot drive loss down on one fixed batch, suspect the training machinery before blaming dataset scale.
45. Inspect activation statistics
handles = []
def hook(name):
def fn(module, inputs, output):
if isinstance(output, torch.Tensor):
x = output.detach()
print(
name,
tuple(x.shape),
"mean=", x.mean().item(),
"std=", x.std().item(),
"min=", x.min().item(),
"max=", x.max().item(),
)
return fn
for name, module in model.named_modules():
if isinstance(module, (nn.LayerNorm, nn.Linear)):
handles.append(
module.register_forward_hook(
hook(name)
)
)
x, y = get_batch(train_data)
model(x, y)
for handle in handles:
handle.remove()
This is noisy.
Use it when you need evidence.
Part XVIII — Common transformer mistakes
46. Wrong head reshape
Bad:
q = q.view(B, self.n_head, T, self.head_dim)
This may create a tensor with the expected final shape while scrambling the semantic grouping of sequence and head dimensions.
Correct pattern:
q = (
q.view(B, T, self.n_head, self.head_dim)
.transpose(1, 2)
)
Shape alone is not always enough.
Dimension meaning matters.
47. Forgetting contiguity when merging heads
After:
y = y.transpose(1, 2)
the tensor may be non-contiguous.
This can break:
y.view(B, T, C)
Safer:
y = (
y.transpose(1, 2)
.contiguous()
.view(B, T, C)
)
Or use a reshape where appropriate and still understand what memory movement may occur.
48. Softmax before cross entropy
Bad:
probs = F.softmax(logits, dim=-1)
loss = F.cross_entropy(probs, targets)
Cross entropy expects logits.
Use:
loss = F.cross_entropy(logits, targets)
The softmax is already part of the computation implied by the loss.
49. Position table too short
If:
T > block_size
then learned positional embeddings cannot index those positions.
The explicit check:
if T > self.config.block_size:
raise ValueError(...)
is much better than allowing an obscure embedding-index failure later.
50. Generating without truncating context
Bad:
logits, _ = model(idx)
once idx has grown past block_size.
Correct baseline:
idx_cond = idx[:, -model.config.block_size :]
logits, _ = model(idx_cond)
The model can only consume the context length it was designed for.
51. Forgetting dropout behavior during generation
Generation should normally use evaluation behavior:
model.eval()
Otherwise dropout remains active and generation includes training-time stochasticity from dropout in addition to sampling stochasticity.
Part XIX — Search-oriented failure guide
52. “index out of range in self” from nn.Embedding
Check token IDs:
print(x.min().item())
print(x.max().item())
print(config.vocab_size)
Required:
0 <= token_id < vocab_size
Likely causes:
- tokenizer/model vocabulary mismatch;
- stale checkpoint with different vocabulary;
- corrupted tokenization;
- incorrect special-token IDs.
53. “mat1 and mat2 shapes cannot be multiplied”
Trace the last dimension.
Transformer linear layers expect:
(..., in_features)
If n_embd=256, most residual-stream linear layers expect their input last dimension to be 256.
Print:
print(x.shape)
print(layer.weight.shape)
Do not guess.
54. CUDA out of memory during attention
Attention cost grows quickly with sequence length because the attention score structure relates positions to positions.
Immediate experiments:
reduce batch size
reduce block_size
reduce n_layer
reduce n_embd
use AMP where appropriate
use activation checkpointing
profile allocated/peak memory
Change one variable at a time and measure.
55. Loss decreases but generated text is poor
Possible reasons include:
- insufficient data;
- model too small;
- training too short;
- validation loss still high;
- character tokenization creates long dependencies;
- sampling settings are poor;
- overfitting;
- prompt distribution differs from training data.
Do not evaluate the system from one amusing or terrible sample.
Track validation loss and generate multiple fixed prompts periodically.
Part XX — Better tokenizer, same model
56. Moving beyond characters
Character tokenization is useful pedagogically because it makes the entire pipeline transparent.
For better efficiency, real language-model pipelines usually use subword or byte-based tokenization schemes.
The important architectural point is that the transformer does not fundamentally care how the IDs were produced.
It receives:
integer token IDs
with a vocabulary size.
Change the tokenizer and update:
config.vocab_size
while preserving the same model contract.
Part XXI — Weight tying
57. Optional token/output weight tying
The input embedding and output vocabulary projection both involve the vocabulary.
A common design choice is to share those weights.
With matching shapes:
self.lm_head.weight = self.token_embedding.weight
If you experiment with this, verify registration and parameter counts afterward.
for name, p in model.named_parameters():
print(name, tuple(p.shape))
Do not assume sharing happened just because the assignment executed.
Inspect the model state.
Part XXII — Reproducibility
58. Seed the experiment
import random
import torch
seed = 1337
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
Reproducibility can still depend on hardware, kernels and backend behavior, but fixing obvious random seeds makes debugging much easier.
Log the seed with your experiment.
59. Save configuration alongside metrics
run_config = {
"seed": seed,
"batch_size": batch_size,
"learning_rate": learning_rate,
"model": config.__dict__,
}
print(run_config)
A loss curve without the configuration that produced it is difficult to reproduce.
Part XXIII — A useful training log
60. Log more than loss
At minimum, periodically record:
step
train loss
validation loss
learning rate
gradient norm
tokens/second
peak CUDA memory
Example:
lr = optimizer.param_groups[0]["lr"]
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(),
1.0,
)
print(
f"step={step} "
f"loss={loss.item():.4f} "
f"lr={lr:.3e} "
f"grad_norm={float(grad_norm):.4f}"
)
This gives future debugging sessions evidence rather than memories.
Part XXIV — Why this final model matters
61. We have now removed most of the magic
Look back at the model.
There is no mysterious pipeline() call.
There is no pretrained checkpoint.
There is no hidden service.
The language model is:
embedding lookup
+ positional representation
+ repeated attention/MLP residual blocks
+ layer normalization
+ vocabulary projection
+ cross entropy
+ gradient descent
That does not mean modern frontier systems are simple.
Scale, data, distributed training, kernels, optimization, architecture, inference systems and evaluation make them enormously more complex.
But the core computational path is no longer opaque.
Part XXV — The LLM-era programmer
62. Why understanding the mechanism still matters
An LLM can generate nearly every class in this post.
It can generate the tokenizer.
It can generate the training loop.
It can generate the optimizer configuration.
It can generate the checkpoint code.
It can even generate plausible explanations for why the loss is not decreasing.
That makes typing the code less valuable.
It does not make understanding the runtime less valuable.
When the generated system breaks, your useful questions become:
Are the targets aligned?
Are the token IDs valid?
Are the tensor dimensions semantically correct?
Are all parameters registered?
Do gradients exist?
Are gradients finite?
Does the optimizer own the parameters?
Did the parameters actually move?
Is validation loss improving?
Where is the memory going?
Did torch.compile actually improve throughput?
Those questions cannot be answered by looking at source code alone.
They are answered by interrogating the running system.
That is why debugging has occupied so much of this series.
In an LLM-heavy programming workflow, generated code is cheap.
Evidence about what that code actually did is valuable.
Part XXVI — Final checklist
63. Before calling the model “working”
Data
[ ] tokenizer round-trips known text
[ ] token IDs are in range
[ ] training/validation split is correct
[ ] x/y next-token shift is verified
[ ] batch shapes are asserted
Model
[ ] n_embd is divisible by n_head
[ ] every block preserves (B,T,C)
[ ] logits are (B,T,V)
[ ] loss is finite
[ ] parameter count is plausible
[ ] all expected parameters are registered
Training
[ ] gradients exist
[ ] gradients are finite
[ ] optimizer contains trainable parameters
[ ] parameters actually change
[ ] tiny batch can be overfit
[ ] training loss falls
[ ] validation loss is tracked
Performance
[ ] throughput is measured
[ ] CUDA timing is synchronized correctly
[ ] peak memory is measured
[ ] AMP is benchmarked rather than assumed
[ ] torch.compile is benchmarked after warm-up
Generation
[ ] greedy generation works first
[ ] context is cropped to block_size
[ ] generation uses eval mode
[ ] temperature is validated
[ ] top-k is bounded by vocab size
[ ] sampling probabilities are finite
Reproducibility
[ ] seed is recorded
[ ] model config is saved
[ ] vocabulary/tokenizer state is saved
[ ] optimizer state is checkpointed
[ ] training step is checkpointed
64. The complete Zero to Hero path
We started with one scalar parameter.
Then:
Step 00 — What Are We Actually Doing?
Step 01 — Tensor Shapes and Broadcasting
Step 02 — Autograd Debugging
Step 03 — Neural Network From Scratch
Step 04 — nn.Module, Parameters and state_dict
Step 05 — DataLoader Performance
Step 06 — CNN Shape Debugging
Step 07 — Attention Shapes and Masks
Step 08 — Model Not Learning Debugging
Step 09 — CUDA Performance and torch.compile
Step 10 — Small GPT-Style Language Model From Scratch
The final model is not the end because there is always another layer:
distributed training
better tokenization
KV caching
Flash Attention
quantization
LoRA
mixture of experts
long-context attention
evaluation
serving
compiler optimization
But those topics now have somewhere to attach.
You are no longer approaching them as isolated API calls.
You have the underlying model.
Final thought
The most useful skill in PyTorch is not remembering every function name.
It is being able to look at a running model and ask:
What tensor is this?
What shape should it have?
What does each dimension mean?
Where did this value come from?
Does the gradient reach this parameter?
Did the optimizer update it?
What did this optimization actually improve?
If you can answer those questions, the abstractions stop being magic.
And when an LLM writes the next thousand lines for you, you still know how to find out whether they work.
