Build a Transformer from Scratch

I built a Transformer from scratch and it predicted "the the the the." Honestly, some meetings feel the same way.

Norm, Repetitively Decoded AI Agent

1. What We Are Building

We will implement a decoder-only Transformer (the GPT architecture) that performs character-level language modeling. Given a sequence of characters, the model predicts the next character at every position. We choose character-level modeling because it eliminates the need for a tokenizer (Chapter 2), letting us focus entirely on the architecture.

Our model will have these hyperparameters:

This is a small model (~1.6M parameters) that trains in a few minutes on a single GPU (or even on CPU for a few epochs). The architecture is identical to GPT; only the scale differs.

Figure 4.2.2: Architecture of our decoder-only Transformer. N blocks of self-attention + FFN with Pre-LN ordering, followed by a final normalization and linear projection.

2. The Complete Implementation

Below is the full model in a single file. We break it into logical pieces and explain each one. The complete code (all pieces assembled) is approximately 300 lines including comments.

2.1 Imports and Configuration

We begin by importing PyTorch and defining a configuration dataclass that holds all model hyperparameters.

"""
mini_transformer.py
A minimal decoder-only Transformer for character-level language modeling.
~300 lines of annotated PyTorch.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class TransformerConfig:
 """All hyperparameters in one place."""
 vocab_size: int = 65 # number of unique characters
 block_size: int = 128 # maximum context length
 n_layers: int = 4 # number of Transformer blocks
 n_heads: int = 4 # number of attention heads
 d_model: int = 128 # embedding / residual stream dimension
 d_ff: int = 512 # feed-forward inner dimension
 dropout: float = 0.1 # dropout probability
 bias: bool = False # use bias in Linear layers?

# Training loop for the mini-transformer: character-level batches,
# cross-entropy loss, AdamW optimizer with gradient clipping.

import time
from torch.utils.data import DataLoader

def train(config=None, text_path='input.txt', max_steps=5000,
 batch_size=64, learning_rate=3e-4, eval_interval=500):
 """Complete training procedure."""

 # ---- Setup ----
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Using device: {device}")

 # Load text data
 with open(text_path, 'r', encoding='utf-8') as f:
 text = f.read()

 # Create dataset
 dataset = CharDataset(text, block_size=128)
 print(f"Vocabulary size: {dataset.vocab_size}")
 print(f"Dataset size: {len(dataset):,} examples")

 # Create config with correct vocab size
 if config is None:
 config = TransformerConfig(vocab_size=dataset.vocab_size)
 else:
 config.vocab_size = dataset.vocab_size

 # Create model
 model = MiniTransformer(config).to(device)
 n_params = sum(p.numel() for p in model.parameters())
 print(f"Model parameters: {n_params:,}")

 # Optimizer
 optimizer = torch.optim.AdamW(
 model.parameters(),
 lr=learning_rate,
 betas=(0.9, 0.95),
 weight_decay=0.1
 )

 # Data loader
 loader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
 num_workers=0, pin_memory=True)
 data_iter = iter(loader)

 # ---- Training ----
 model.train()
 t0 = time.time()

 for step in range(max_steps):
 # Get batch (cycle through data)
 try:
 xb, yb = next(data_iter)
 except StopIteration:
 data_iter = iter(loader)
 xb, yb = next(data_iter)

 xb, yb = xb.to(device), yb.to(device)

 # Forward pass
 logits, loss = model(xb, yb)

 # Backward pass
 optimizer.zero_grad(set_to_none=True)
 loss.backward()

 # Gradient clipping (standard practice)
 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

 optimizer.step()

 # Logging
 if step % eval_interval == 0 or step == max_steps - 1:
 dt = time.time() - t0
 print(f"step {step:5d} | loss {loss.item():.4f} | "
 f"time {dt:.1f}s")

 # ---- Generation ----
 model.eval()
 prompt = "\n"
 context = torch.tensor(
 [dataset.encode(prompt)], dtype=torch.long, device=device
 )
 generated = model.generate(context, max_new_tokens=500, temperature=0.8)
 print("\n" + "=" * 50)
 print("Generated text:")
 print("=" * 50)
 print(dataset.decode(generated[0].tolist()))

 return model, dataset

if __name__ == '__main__':
 train()

We use a dataclass so that every hyperparameter is explicit, documented, and easy to modify. Setting bias=False follows the LLaMA convention and marginally reduces parameter count. Code Fragment 4.2.2 below puts this into practice.

2.2 Causal Self-Attention

This module implements multi-head causal self-attention with a triangular mask that prevents positions from attending to future tokens.

# Causal self-attention with a triangular mask: each token can only
# attend to itself and earlier positions, enforcing left-to-right generation.

class CausalSelfAttention(nn.Module):
 """Multi-head causal (masked) self-attention."""

 def __init__(self, config: TransformerConfig):
 super().__init__()
 assert config.d_model % config.n_heads == 0

 # Key, Query, Value projections combined into one matrix
 self.qkv_proj = nn.Linear(config.d_model, 3 * config.d_model, bias=config.bias)
 # Output projection
 self.out_proj = nn.Linear(config.d_model, config.d_model, bias=config.bias)

 self.attn_dropout = nn.Dropout(config.dropout)
 self.resid_dropout = nn.Dropout(config.dropout)

 self.n_heads = config.n_heads
 self.d_model = config.d_model
 self.d_k = config.d_model // config.n_heads

 # Causal mask: lower-triangular boolean matrix
 # Register as buffer so it moves to GPU with the model
 mask = torch.tril(torch.ones(config.block_size, config.block_size))
 self.register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size))

 def forward(self, x):
 B, T, C = x.shape # batch, sequence length, d_model

 # Compute Q, K, V in one matrix multiply, then split
 qkv = self.qkv_proj(x)
 q, k, v = qkv.split(self.d_model, dim=2)

 # Reshape for multi-head: (B, T, C) -> (B, n_heads, T, d_k)
 q = q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
 k = k.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
 v = v.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

 # Scaled dot-product attention
 # (B, n_heads, T, d_k) @ (B, n_heads, d_k, T) -> (B, n_heads, T, T)
 scores = (q @ k.transpose(-2, -1)) * (self.d_k ** -0.5)

 # Apply causal mask: positions beyond current token get -inf
 scores = scores.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))

 attn_weights = F.softmax(scores, dim=-1)
 attn_weights = self.attn_dropout(attn_weights)

 # Weighted sum of values
 # (B, n_heads, T, T) @ (B, n_heads, T, d_k) -> (B, n_heads, T, d_k)
 out = attn_weights @ v

 # Concatenate heads: (B, n_heads, T, d_k) -> (B, T, C)
 out = out.transpose(1, 2).contiguous().view(B, T, C)

 # Final linear projection + dropout
 return self.resid_dropout(self.out_proj(out))

2.3 Feed-Forward Network

The position-wise feed-forward network applies two linear transformations with a nonlinearity in between.

# Position-wise FFN: expand to 4*d_model with ReLU, then project back.
# Applied independently at every token position in the sequence.

class FeedForward(nn.Module):
 """Position-wise feed-forward network with ReLU activation."""

 def __init__(self, config: TransformerConfig):
 super().__init__()
 self.fc1 = nn.Linear(config.d_model, config.d_ff, bias=config.bias)
 self.fc2 = nn.Linear(config.d_ff, config.d_model, bias=config.bias)
 self.dropout = nn.Dropout(config.dropout)

 # Forward pass: define computation graph
 def forward(self, x):
 x = F.relu(self.fc1(x))
 x = self.dropout(x)
 x = self.fc2(x)
 x = self.dropout(x)
 return x

This is the simplest version. For a more advanced variant, you can swap in SwiGLU. Code Fragment 4.2.4 below puts this into practice.

# SwiGLU FFN variant (LLaMA, PaLM): replace ReLU with a gated SiLU
# activation, using 2/3 of the hidden dimension for the gate projection.

class SwiGLUFeedForward(nn.Module):
 """SwiGLU feed-forward (used in LLaMA, PaLM)."""

 def __init__(self, config: TransformerConfig):
 super().__init__()
 # SwiGLU uses 3 weight matrices instead of 2
 # To keep param count comparable, the hidden dim is often 2/3 of d_ff
 hidden = int(2 * config.d_ff / 3)
 self.w1 = nn.Linear(config.d_model, hidden, bias=config.bias)
 self.w2 = nn.Linear(hidden, config.d_model, bias=config.bias)
 self.w3 = nn.Linear(config.d_model, hidden, bias=config.bias) # gate
 self.dropout = nn.Dropout(config.dropout)

 def forward(self, x):
 # SiLU(x * W1) * (x * W3) then project back
 return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))

2.4 Transformer Block

Each block combines causal self-attention and a feed-forward network with residual connections and layer normalization.

# Single transformer block: Pre-LN ordering with residual connections.
# Input flows through norm, attention, add residual, norm, FFN, add residual.

class TransformerBlock(nn.Module):
 """A single Transformer block with Pre-LN ordering."""

 def __init__(self, config: TransformerConfig):
 super().__init__()
 self.ln1 = nn.LayerNorm(config.d_model)
 self.attn = CausalSelfAttention(config)
 self.ln2 = nn.LayerNorm(config.d_model)
 self.ffn = FeedForward(config)

 # Forward pass: define computation graph
 def forward(self, x):
 # Pre-LN: normalize before each sub-layer
 x = x + self.attn(self.ln1(x)) # residual + attention
 x = x + self.ffn(self.ln2(x)) # residual + FFN
 return x

This is remarkably simple. Two lines of actual computation, each following the pattern: x = x + SubLayer(LayerNorm(x)). The residual connection is the x + at the beginning; the Pre-LN ordering means we normalize the input to each sub-layer, not the output.

2.5 The Complete Model

The full model stacks an embedding layer, multiple Transformer blocks, and a final linear head for next-token prediction.

# Full decoder-only transformer: stack N blocks, add token + position
# embeddings, project to vocabulary logits, and implement autoregressive generate().

class MiniTransformer(nn.Module):
 """Decoder-only Transformer for character-level language modeling."""

 def __init__(self, config: TransformerConfig):
 super().__init__()
 self.config = config

 # Token and position embeddings
 self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
 self.pos_emb = nn.Embedding(config.block_size, config.d_model)
 self.drop = nn.Dropout(config.dropout)

 # Stack of Transformer blocks
 self.blocks = nn.ModuleList([
 TransformerBlock(config) for _ in range(config.n_layers)
 ])

 # Final layer norm (needed with Pre-LN)
 self.ln_final = nn.LayerNorm(config.d_model)

 # Output head: project from d_model to vocab_size
 self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

 # Weight tying: share embedding and output weights
 self.token_emb.weight = self.lm_head.weight

 # Initialize weights
 self.apply(self._init_weights)
 # Scale residual projections
 for block in self.blocks:
 nn.init.normal_(
 block.attn.out_proj.weight,
 mean=0.0,
 std=0.02 / math.sqrt(2 * config.n_layers)
 )
 nn.init.normal_(
 block.ffn.fc2.weight,
 mean=0.0,
 std=0.02 / math.sqrt(2 * config.n_layers)
 )

 def _init_weights(self, module):
 if isinstance(module, nn.Linear):
 nn.init.normal_(module.weight, mean=0.0, std=0.02)
 if module.bias is not None:
 nn.init.zeros_(module.bias)
 elif isinstance(module, nn.Embedding):
 nn.init.normal_(module.weight, mean=0.0, std=0.02)
 elif isinstance(module, nn.LayerNorm):
 nn.init.ones_(module.weight)
 nn.init.zeros_(module.bias)

 def forward(self, idx, targets=None):
 """
 Args:
 idx: (B, T) tensor of token indices
 targets: (B, T) tensor of target token indices (optional)
 Returns:
 logits: (B, T, vocab_size)
 loss: scalar cross-entropy loss (only if targets provided)
 """
 B, T = idx.shape
 assert T <= self.config.block_size, \
 f"Sequence length {T} exceeds block_size {self.config.block_size}"

 # Token embeddings + positional embeddings
 positions = torch.arange(0, T, device=idx.device) # (T,)
 x = self.token_emb(idx) + self.pos_emb(positions) # (B, T, d_model)
 x = self.drop(x)

 # Pass through all Transformer blocks
 for block in self.blocks:
 x = block(x)

 # Final normalization
 x = self.ln_final(x)

 # Project to vocabulary
 logits = self.lm_head(x) # (B, T, vocab_size)

 # Compute loss if targets are provided
 loss = None
 if targets is not None:
 loss = F.cross_entropy(
 logits.view(-1, logits.size(-1)),
 targets.view(-1)
 )

 return logits, loss

 @torch.no_grad()
 def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
 """
 Auto-regressive generation.
 Args:
 idx: (B, T) conditioning sequence
 max_new_tokens: number of tokens to generate
 temperature: softmax temperature (lower = more deterministic)
 top_k: if set, only sample from top-k most likely tokens
 """
 for _ in range(max_new_tokens):
 # Crop context to block_size if needed
 idx_cond = idx[:, -self.config.block_size:]

 # Forward pass
 logits, _ = self(idx_cond)

 # Take logits at the last position and apply temperature
 logits = logits[:, -1, :] / temperature

 # Optional top-k filtering
 if top_k is not None:
 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
 logits[logits < v[:, [-1]]] = float('-inf')

 # Sample from the distribution
 probs = F.softmax(logits, dim=-1)
 next_token = torch.multinomial(probs, num_samples=1)

 # Append to sequence
 idx = torch.cat([idx, next_token], dim=1)

 return idx

# Production equivalent: load a pre-trained causal LM
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
inputs = tokenizer("The future of AI", return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=50, temperature=0.8)
print(tokenizer.decode(output[0]))

3. Data Preparation

For training, we use a simple character-level dataset. Any plain text file will work. We will use a small text corpus (a few hundred KB) for quick experimentation. Code Fragment 4.2.7 below puts this into practice.

# Define CharDataset; implement __len__, __getitem__, decode
# See inline comments for step-by-step details.

class CharDataset:
 """Character-level dataset that produces (input, target) pairs."""

 def __init__(self, text, block_size):
 self.block_size = block_size
 # Build character vocabulary
 chars = sorted(set(text))
 self.vocab_size = len(chars)
 self.stoi = {ch: i for i, ch in enumerate(chars)}
 self.itos = {i: ch for ch, i in self.stoi.items()}
 # Encode entire text as integers
 self.data = torch.tensor([self.stoi[c] for c in text], dtype=torch.long)

 def __len__(self):
 return len(self.data) - self.block_size

 def __getitem__(self, idx):
 chunk = self.data[idx : idx + self.block_size + 1]
 x = chunk[:-1] # input: characters 0..block_size-1
 y = chunk[1:] # target: characters 1..block_size
 return x, y

 def decode(self, indices):
 """Convert list of integer indices back to string."""
 return ''.join(self.itos[i] for i in indices)

 def encode(self, text):
 """Convert string to list of integer indices."""
 return [self.stoi[c] for c in text]

4. The Training Loop

5. Understanding the Shapes

Tracking tensor shapes is one of the most valuable debugging skills when working with Transformers. Here is a shape trace through the forward pass:

6. Running the Lab

6.1 Getting Data

Download a small text file for training. Shakespeare's collected works (~1.1 MB) is the classic choice. Code Fragment 4.2.9 below puts this into practice.

# Download the tiny Shakespeare dataset
import urllib.request
url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
urllib.request.urlretrieve(url, "input.txt")

6.2 Training

With the dataset prepared, we can now train the model and observe the loss decreasing over several thousand steps.

# Train with default settings
model, dataset = train(max_steps=5000)

# Expected output after 5000 steps (loss around 1.4-1.5):
# step 0 | loss 4.1742 | time 0.0s
# step 500 | loss 1.9831 | time 12.3s
# step 1000 | loss 1.6524 | time 24.7s
# ...
# step 5000 | loss 1.4208 | time 123.5s

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

input_ids = tokenizer("To be, or not to be", return_tensors="pt").input_ids
output = model.generate(input_ids, max_new_tokens=50, temperature=0.8, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=True))

6.3 Evaluating the Output

After training, the model will generate text that resembles the style of the training data. At ~5000 steps with our small model, you should see recognizable words, approximate sentence structure, and character-level patterns that match the training corpus. The text will not be coherent, but it should clearly be "trying" to produce English in the style of the training data.

6.4 Experiments to Try

• Increase n_layers from 4 to 6 or 8. Does the loss improve? How much slower is training?
• Increase d_model from 128 to 256. Compare the parameter count and training speed.
• Remove positional embeddings entirely. What happens to the generated text?
• Switch to SwiGLU (replace FeedForward with SwiGLUFeedForward). Does the loss curve change?
• Remove the causal mask. The model can now "cheat" by looking at future tokens. What happens to the training loss? What happens to generation quality?
• Try temperature values of 0.5, 1.0, and 1.5 during generation. Observe the diversity/quality tradeoff.

7. Common Bugs and Debugging

When implementing Transformers from scratch, certain bugs appear repeatedly. Here are the most common ones and how to detect them: