Build a Mini GPT and Train It Three Ways

Introduction

30 Second Summary

Every time you chat with an AI assistant, a language model is predicting the next token based on patterns learned from vast amounts of text. But how does that model actually learn? The answer lives in three distinct training stages, each with its own loss function and its own effect on the model's behaviour.

In this project, you will build a complete mini GPT language model from scratch and train it through all three stages of the modern LLM pipeline: pre-training, supervised fine-tuning, and preference optimization.

What You'll Build

A from-scratch GPT implementation (~2M parameters) in a single Google Colab notebook, trained through pre-training on Shakespeare, supervised fine-tuning on instructions, and Direct Preference Optimization for alignment.

By the end of this project, you'll have:

• A working mini GPT model with every component (attention, feed-forward network, layer normalization, positional embeddings) implemented and understood mathematically.
• Trained weights from all three stages of the modern LLM pipeline: pre-training (next-token prediction), SFT (instruction following), and DPO (preference alignment).
• Loss curve visualizations and generated text comparisons showing how each training mode shapes model behaviour differently.
• Secret Mission: Implement GRPO (the algorithm behind DeepSeek-R1) with a verifiable arithmetic reward.

Are there any prerequisites?

You should be comfortable with basic Python and have a general understanding of neural network fundamentals (what a loss function is, what gradient descent does). Some familiarity with basic linear algebra (matrix multiplication) is helpful but not strictly required.

Before We Start

You are about to build a language model from raw math. Every component you will implement (attention, feed-forward layers, normalization, training loops) is the same machinery running inside GPT-4, Claude, and every other large language model.

In this step, you will open a Google Colab notebook, connect to a free GPU, and confirm that PyTorch can access it.

In this step, get ready to:

• Open a new Google Colab notebook.
• Set the runtime to a T4 GPU.
• Verify PyTorch and CUDA are working.

Open a new Google Colab notebook

• Go to colab.research.google.com and sign in with your Google account.
• Click New Notebook in the bottom-left corner of the dialog that appears.

What is Google Colab?

Google Colab is a free, browser-based notebook environment that gives you access to GPUs without installing anything on your computer. It comes with PyTorch pre-installed, so you can start writing deep learning code immediately.

Set the runtime to T4 GPU

Training a neural network on a CPU is painfully slow. A GPU processes thousands of matrix multiplications in parallel, cutting training time from hours to minutes.

• Click Runtime in the top menu bar.
• Click Change runtime type.
• In the Hardware accelerator dropdown, select T4 GPU.
• Click Save.

Why a T4 GPU?

The T4 is the standard free-tier GPU on Colab. It has 16 GB of VRAM, which is more than enough for our ~2M parameter model. Training will take about 3-5 minutes instead of 15+ minutes on CPU.

Import libraries and verify CUDA

Before building the model, you need to import the libraries you will use throughout this project and confirm the GPU is accessible.

• Paste this code into the first cell of your notebook and run it by pressing Shift+Enter:

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

What does this code do?

• torch is the core PyTorch library for tensor operations and neural networks.
• torch.nn provides building blocks like layers, loss functions, and embedding tables.
• torch.nn.functional gives you stateless functions like softmax and cross_entropy.
• matplotlib.pyplot is for plotting loss curves later.
• The device variable automatically picks the GPU if available, or falls back to CPU.

• Confirm you see output similar to:

PyTorch version: 2.x.x
CUDA available: True

If CUDA shows True, your GPU is ready.

Seeing CUDA available: False?

• Double-check that you selected T4 GPU (not None) in Runtime > Change runtime type.
• If the T4 option is greyed out, Colab's free GPU quota may be temporarily exhausted. Wait a few minutes and try again.
• If you cannot get a GPU at all, the project still works on CPU. Training will be slower (~15 minutes instead of ~3-5 minutes) but produces the same results.

Help me troubleshoot why CUDA is not available in my Google Colab notebook. I selected T4 GPU in the runtime settings.

✔️ Awesome, I've got everything!

Great. Double check you have saved the cell output showing CUDA available: True.

ⓧ I'd like to double check the full code

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

Your environment is ready. Next up, you will download the Tiny Shakespeare dataset and build a character-level tokenizer to feed text into your model.

Set Up Data and Hyperparameters

Before building the transformer architecture, you need two things: data to learn from and a clear set of hyperparameters that define the model's size and training behaviour.

In this step, you will download Shakespeare's works as training text, build a character-level tokenizer to convert text into numbers, and define the configuration that controls how large the model is and how it trains.

In this step, get ready to:

• Define model and training hyperparameters.
• Download and tokenize the Tiny Shakespeare dataset.
• Create a batch loading function for training.

Define model and training hyperparameters

Hyperparameters are the knobs you set before training begins. They control how large the model is, how long it trains, and how aggressively it learns.

• In a new code cell below your existing imports and device setup, add the hyperparameter configuration by pasting this code:

# Model hyperparameters
n_embd = 128
n_head = 4
n_layer = 4
block_size = 128
dropout = 0.1
batch_size = 32

# Training hyperparameters
max_iters = 3000
eval_interval = 100
learning_rate = 3e-4
warmup_steps = 100
min_lr = 3e-5
weight_decay = 0.01

What do these hyperparameters control?

The model hyperparameters define the architecture's shape:

• n_embd = 128: each token becomes a 128-dimensional vector.
• n_head = 4: the attention mechanism splits into 4 parallel heads.
• n_layer = 4: the model stacks 4 transformer blocks deep.
• block_size = 128: the model can look at 128 tokens of context at once.

The training hyperparameters control how the optimizer updates weights. A learning rate of 3e-4 with cosine decay and warmup is a proven recipe from GPT-2 and similar models.

• Run the cell.

You won't see any output yet. These values are stored in memory and used by the model and training loop you build next.

Download and tokenize the Tiny Shakespeare dataset

A language model learns by predicting the next character (or token) in a sequence. You need a dataset large enough to show patterns, but small enough to train in minutes on a free GPU.

• Download the Tiny Shakespeare dataset and load it into memory by adding this code in a new cell:

import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, 'input.txt')

with open('input.txt', 'r') as f:
    text = f.read()

print(f'Dataset size: {len(text):,} characters')

Why Tiny Shakespeare?

Tiny Shakespeare is approximately 1.1 million characters of Shakespeare's collected works. It is small enough to train a meaningful model in 3-5 minutes on a Colab T4 GPU, but large enough that real language patterns (grammar, word structure, character names) emerge during training.

• Run the cell. You should see output showing the dataset is approximately 1.1 million characters.

Seeing a download error?

Make sure your Colab runtime is connected (green checkmark in the top right). If the URL fails, check that you copied it exactly with no extra spaces.

Help me debug this download error

Now you need a way to convert characters into numbers that the model can process. This is called a tokenizer.

• Build the character-level tokenizer by adding this code in a new cell:

# Build character-level tokenizer
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# Test the tokenizer
print(f'Vocabulary size: {vocab_size} unique characters')
print(f'encode("hello") = {encode("hello")}')
print(f'decode(encode("hello")) = {decode(encode("hello"))}')

What does this tokenizer do?

The tokenizer maps each unique character to an integer and back. stoi (string-to-integer) is a dictionary mapping characters to numbers. itos (integer-to-string) is the reverse.

The encode function turns a string into a list of integers. The decode function turns integers back into readable text. The model only sees numbers during training, and you use decode to convert its output back to text.

• Run the cell. You should see that the vocabulary size is 65 unique characters, and that encoding then decoding "hello" returns "hello" unchanged.

Seeing a NameError for text?

Make sure you ran the previous cell (the one that downloads and loads the dataset) before running this one. In Colab, cells must be run in order.

Help me fix this NameError

Create the data loading function

The model learns by seeing batches of text sequences. For each input sequence, the target is the same sequence shifted one position to the right. This is the next-token prediction setup.

• Split the data into training and validation sets, then create the batch function by adding this code in a new cell:

# Convert text to tensor and split into train/val
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

def get_batch(split):
    """Return a random batch of (input, target) pairs."""
    d = train_data if split == 'train' else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([d[i:i+block_size] for i in ix])
    y = torch.stack([d[i+1:i+1+block_size] for i in ix])
    return x.to(device), y.to(device)

# Test the batch function
xb, yb = get_batch('train')
print(f'Input shape: {xb.shape}')
print(f'Target shape: {yb.shape}')
print(f'Train tokens: {len(train_data):,} | Val tokens: {len(val_data):,}')

How does next-token prediction work here?

Each batch contains batch_size sequences, each block_size tokens long. The target y is the input x shifted right by one position.

If x contains the tokens for "To be or not to b", then y contains "o be or not to be". The model's job is to predict each next character given everything before it.

• Run the cell. You should see that both input and target shapes are torch.Size([32, 128]) (32 sequences of 128 tokens each). The train/val split should show roughly 1 million train tokens and 100,000 val tokens.

Seeing a shape mismatch or NameError?

Make sure you ran the hyperparameter cell first. The get_batch function uses block_size and batch_size from that cell. Also confirm the tokenizer cell ran successfully before this one.

Help me fix my get_batch errors

✔️ Awesome, I've got everything!

Great. Make sure all four cells have been run in order (hyperparameters, download, tokenizer, batch function).

ⓧ I'd like to double check the full code

Here is the complete code you should have across all your cells at this point:

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Model hyperparameters
n_embd = 128
n_head = 4
n_layer = 4
block_size = 128
dropout = 0.1
batch_size = 32

# Training hyperparameters
max_iters = 3000
eval_interval = 100
learning_rate = 3e-4
warmup_steps = 100
min_lr = 3e-5
weight_decay = 0.01

import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, 'input.txt')

with open('input.txt', 'r') as f:
    text = f.read()

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

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

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

Your data pipeline is ready. You have a tokenizer that converts text to numbers and back, a batch function that serves training examples, and hyperparameters that will shape your model. Next up, you will build the attention mechanism that gives transformers their power.

Build the Attention Mechanism

Your data pipeline is ready. Now it is time to build the core mechanism that makes transformers work: self-attention.

Attention lets each token in a sequence look at every other token and decide which ones are relevant for predicting the next word. In this step, you will implement the full attention formula from the original "Attention Is All You Need" paper, then extend it to multi-head attention so the model can attend to different relationships simultaneously.

In this step, get ready to:

• Understand the scaled dot-product attention formula and why each piece exists.
• Implement a multi-head CausalSelfAttention class from scratch in PyTorch.
• Verify the mechanism produces correct output shapes and shows the expected causal pattern.

Understand scaled dot-product attention

Every transformer begins with one core operation: scaled dot-product attention. The formula is:

# Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
#
# Q = queries (what each token is looking for)
# K = keys (what each token advertises about itself)
# V = values (the actual content each token provides)
# d_k = dimension of each key vector

Why divide by sqrt(d_k)?

When the embedding dimension is large, the dot products between Q and K grow large in magnitude. Large values push softmax into regions where the gradients are extremely small (near 0 or 1). Dividing by sqrt(d_k) keeps the values in a range where softmax produces useful gradients during training.

Without this scaling, the model effectively learns nothing because gradients vanish through the softmax bottleneck.

For a language model, attention also needs a causal mask. Each token can only attend to tokens at earlier positions. This prevents the model from "cheating" by looking at future tokens it is supposed to predict.

What is a causal mask?

The causal mask is a lower-triangular matrix of ones. Positions above the diagonal get filled with negative infinity before softmax, which makes those attention weights become zero. Token 5 can see tokens 1-5, but not tokens 6 onward.

Implement the CausalSelfAttention class

Multi-head attention runs multiple attention computations in parallel, each learning different relationships. You split the embedding dimension across heads, compute attention independently in each head, then concatenate the results.

• Add a new code cell below your existing code and paste the full CausalSelfAttention class:

class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        )

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(n_embd, dim=2)
        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)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

What does this code do?

The __init__ method sets up the architecture:

• c_attn is a single linear layer that projects the input into Q, K, and V simultaneously (3x the embedding size). This is more efficient than three separate layers.
• c_proj projects the concatenated heads back to the original embedding dimension.
• The mask buffer is a pre-computed lower-triangular matrix. Using register_buffer ensures it moves to GPU with the model but is not a trainable parameter.

The forward method implements the attention formula step by step:

• Project input to Q, K, V, then reshape to (B, n_head, T, head_dim) so each head operates independently.
• Compute attention scores with the sqrt(head_dim) scaling factor.
• Apply the causal mask and softmax to get attention weights.
• Multiply weights by V, concatenate heads, and project the output.

• Run the cell to define the class.

Seeing an AssertionError?

The line assert n_embd % n_head == 0 checks that the embedding dimension divides evenly across heads. If you changed n_embd or n_head from the values set in Step 1, make sure n_embd (128) is divisible by n_head (4).

help me fix the AssertionError in CausalSelfAttention

Verify the attention mechanism

Before building the rest of the model, confirm that attention produces the correct output shape and shows the expected causal pattern.

• Add a new cell and run this verification code:

# Verify attention output shape
attn = CausalSelfAttention().to(device)
test_input = torch.randn(1, block_size, n_embd).to(device)
output = attn(test_input)
print(f"Input shape:  {test_input.shape}")
print(f"Output shape: {output.shape}")

# Visualize attention weights for one head
with torch.no_grad():
    qkv = attn.c_attn(test_input)
    q, k, v = qkv.split(n_embd, dim=2)
    q = q.view(1, block_size, n_head, n_embd // n_head).transpose(1, 2)
    k = k.view(1, block_size, n_head, n_embd // n_head).transpose(1, 2)
    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(n_embd // n_head))
    att = att.masked_fill(attn.mask[:, :, :block_size, :block_size] == 0, float('-inf'))
    att = F.softmax(att, dim=-1)

# Plot head 0 attention weights (first 16 positions for clarity)
plt.figure(figsize=(8, 6))
plt.imshow(att[0, 0, :16, :16].cpu(), cmap='Blues')
plt.title('Causal Attention Weights - Head 0')
plt.xlabel('Key Position')
plt.ylabel('Query Position')
plt.colorbar()
plt.show()

What should I see?

The output shape should be torch.Size([1, 128, 128]) matching the input. This confirms the attention mechanism preserves tensor dimensions.

The attention weight plot should show a lower-triangular pattern. Each row (query position) has non-zero values only at or before its column index. This is the causal mask at work. Token at position 5 attends to positions 0-5 but not beyond.

Plot not showing?

Make sure you ran the matplotlib import (import matplotlib.pyplot as plt) from your first cell in Step 1. If the plot is blank, check that your runtime has a GPU assigned by going to Runtime > Change runtime type.

help me debug why my attention visualization is not displaying

✔️ Awesome, I've got everything!

Great. Make sure your class definition cell has been executed successfully before moving on.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point:

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'

n_embd = 128
n_head = 4
n_layer = 4
block_size = 128
dropout = 0.1
batch_size = 32

max_iters = 3000
eval_interval = 100
learning_rate = 3e-4
warmup_steps = 100
min_lr = 3e-5
weight_decay = 0.01

import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, 'input.txt')

with open('input.txt', 'r') as f:
    text = f.read()

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

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

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

class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        )

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(n_embd, dim=2)
        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)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

Your attention mechanism is working. Next up, you will assemble the full GPT model by stacking transformer blocks, adding embeddings, and wiring up the generation loop.

Assemble the Full GPT Model

You have a working attention mechanism that lets tokens communicate with each other. Now you need the rest of the architecture that wraps around it to form a complete language model.

A transformer block combines attention with a feed-forward network, and stacking multiple blocks gives the model depth. In this step, you will build the full GPT model from these blocks, connect it to embeddings, and verify it can produce output (even if that output is gibberish for now).

In this step, get ready to:

• Build the transformer block (feed-forward network + residual connections).
• Assemble the full MiniGPT model with embeddings, stacked blocks, and weight tying.
• Implement autoregressive text generation and verify the model has ~2M parameters.

Build the transformer block

Every transformer block has two sub-layers: attention (which you already built) and a feed-forward network. The FFN expands the representation to 4x the embedding size, applies a non-linearity, then projects back down. This expansion gives the model more capacity to learn complex patterns.

• In a new cell below your CausalSelfAttention class, add the FeedForward class by pasting this code:

class FeedForward(nn.Module):
    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)

What does this code do?

• nn.Linear(n_embd, 4 * n_embd) expands the 128-dimensional representation to 512 dimensions. This larger space lets the network learn more complex transformations.
• nn.GELU() is the activation function used in GPT-2 and modern LLMs. It introduces non-linearity so the network can learn more than just linear relationships.
• The second nn.Linear(4 * n_embd, n_embd) projects back down to the original 128 dimensions so the output can feed into the next block.

• Run the cell.

Seeing a NameError?

Make sure you ran the configuration cell (with n_embd = 128 and dropout = 0.1) earlier in the notebook. If you restarted your runtime, re-run all cells from the top.

Help me debug a NameError in my FeedForward class.

Now you need the TransformerBlock that wires attention and the FFN together with residual connections and layer normalization.

• In a new cell below FeedForward, add the TransformerBlock class by pasting this code:

class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffn = FeedForward()

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

What does this code do?

This is the pre-norm architecture used by GPT-2 and all modern LLMs. Two key patterns make it work:

• Pre-norm: LayerNorm is applied BEFORE attention and FFN (not after). This makes training more stable because gradients flow more smoothly through the residual path.
• Residual connections: The x = x + ... pattern adds the sub-layer output back to its input. This creates a "skip highway" that prevents information from getting lost as data flows through many layers.

• Run the cell.

Seeing a NameError for CausalSelfAttention?

The TransformerBlock class references CausalSelfAttention and FeedForward. Make sure you ran both of those class cells first. In Colab, cells must be run in order from top to bottom.

Help me fix a NameError in my TransformerBlock.

Assemble the full GPT model

With the building blocks ready, you can now assemble the full model. The GPT architecture takes token indices as input, converts them to embeddings, passes them through a stack of transformer blocks, and projects back to vocabulary-sized logits that represent the probability of each next token.

• In a new cell, define the MiniGPT class with its constructor and forward method by pasting this code:

class MiniGPT(nn.Module):
    def __init__(self):
        super().__init__()
        # Token embedding: convert token indices to vectors
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        # Positional embedding: encode position in sequence
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        # Stack of transformer blocks
        self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(n_layer)])
        # Final layer norm
        self.ln_f = nn.LayerNorm(n_embd)
        # Output projection to vocabulary
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        # Weight tying: share embedding and output weights
        self.token_emb.weight = self.lm_head.weight
        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        B, T = idx.size()
        # Look up token and position embeddings
        tok_emb = self.token_emb(idx)
        pos_emb = self.pos_emb(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        # Pass through all transformer blocks
        x = self.blocks(x)
        x = self.ln_f(x)
        # Project to vocabulary logits
        logits = self.lm_head(x)
        # Compute loss if targets provided
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
        return logits, loss

What does this code do?

This is the complete GPT architecture minus generation. Here are the key design choices:

• Token + positional embeddings: The model needs to know both WHAT each token is and WHERE it sits in the sequence. These two embeddings are summed together.
• Weight tying: self.token_emb.weight = self.lm_head.weight forces the input embedding and output projection to share the same weight matrix. This saves ~30% of parameters and improves learning because the model uses the same representation space at both ends.
• _init_weights: Initializes all weights from a small normal distribution (std=0.02). Good initialization prevents exploding or vanishing gradients at the start of training.
• forward(): Takes token indices of shape (B, T), passes them through embeddings and blocks, and returns logits of shape (B, T, vocab_size). If targets are provided, it also computes cross-entropy loss.

Why weight tying?

The token embedding converts a token index to a vector (integer to meaning). The output projection converts a vector back to token probabilities (meaning to integer). These are inverse operations. Sharing their weights forces the model to use a consistent representation: if "king" maps to a certain direction in embedding space, the output head predicts "king" when it sees that same direction.

• Run the cell. You should see no errors (the class is defined but not yet instantiated).

Seeing a NameError for vocab_size or block_size?

The MiniGPT class uses vocab_size (from the tokenizer cell) and block_size (from the config cell). Re-run your data loading and configuration cells if you restarted the runtime.

Help me figure out which cells I need to re-run for my MiniGPT class.

Implement text generation and verify the model

The model can compute logits for a sequence, but it cannot generate new text yet. Autoregressive generation works by repeatedly predicting the next token, appending it to the sequence, and feeding the extended sequence back in.

• Add the generate method to your MiniGPT class. Paste it inside the class, after the forward method:

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=40):
        for _ in range(max_new_tokens):
            # Crop to the last block_size tokens (context window limit)
            idx_cond = idx[:, -block_size:]
            # Get predictions
            logits, _ = self(idx_cond)
            # Focus on the last time step
            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')
            # Convert to probabilities and sample
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            # Append to running sequence
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

What does this code do?

• @torch.no_grad() disables gradient tracking during generation. You only need gradients during training, not inference.
• temperature controls randomness. Lower values (e.g. 0.3) make the model more deterministic. Higher values (e.g. 1.5) make it more creative but potentially incoherent.
• top_k filtering zeroes out all tokens except the k most likely ones before sampling. This prevents the model from choosing extremely unlikely tokens.
• idx[:, -block_size:] crops the input to the context window. The model can only attend to 128 tokens at a time, so older tokens get dropped.

• Re-run the full MiniGPT class cell (with the new generate method included) so Python picks up the updated class definition.

Now instantiate the model and verify everything is connected correctly.

• In a new cell, instantiate the model and print the parameter count by running this code:

model = MiniGPT().to(device)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

What should I see?

You should see a parameter count of approximately 2 million (the exact number depends on your dataset's vocab_size). This is tiny by modern standards (GPT-3 has 175 billion), but large enough to learn real language patterns from Tiny Shakespeare.

• Test generation from the untrained model by running this in a new cell:

# Generate from the untrained model (expect random gibberish)
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print("Untrained model output:")
print(decode(model.generate(context, max_new_tokens=100)[0].tolist()))

• Run the cell. You should see random characters with no recognizable words or structure.

Why is the output random?

The model has never seen any training data. Its weights are random, so it assigns roughly equal probability to every character at every position. After pre-training in the next step, this same generate call will produce coherent Shakespeare-like text.

Getting a RuntimeError about tensor sizes?

Make sure the generate method is indented inside the MiniGPT class (not at module level). Also confirm you re-ran the full class cell after adding it.

If you see a CUDA device mismatch error, check that both your model and the context tensor use the same device variable.

Help me debug a RuntimeError in my MiniGPT generate method.

✔️ Awesome, I've got everything!

Great. Double check you have all four class cells and the instantiation cell saved and run in order.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point. Compare this against what you have in your cells:

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'

n_embd = 128
n_head = 4
n_layer = 4
block_size = 128
dropout = 0.1
batch_size = 32

max_iters = 3000
eval_interval = 100
learning_rate = 3e-4
warmup_steps = 100
min_lr = 3e-5
weight_decay = 0.01

import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, 'input.txt')

with open('input.txt', 'r') as f:
    text = f.read()

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

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

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

class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        )

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(n_embd, dim=2)
        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)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

class FeedForward(nn.Module):
    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 TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffn = FeedForward()

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

class MiniGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.token_emb.weight = self.lm_head.weight
        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        B, T = idx.size()
        tok_emb = self.token_emb(idx)
        pos_emb = self.pos_emb(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        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=0.8, top_k=40):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

model = MiniGPT().to(device)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

Your complete GPT architecture is assembled and verified. It can process sequences, compute loss, and generate text. The random output proves the model works end-to-end. In the next step, you will train it on Shakespeare and watch coherent language emerge from pure math.

Pre-train with Next-Token Prediction

Your model has all the right components, but right now it produces nothing but random noise. Every parameter is a random number. The model has no concept of English, no sense of grammar, and no understanding of Shakespeare.

Pre-training fixes this by showing the model millions of examples of "given these characters, predict the next one." Over thousands of iterations, the model's random parameters gradually shift toward values that make the correct next character more likely. This is the foundation of every LLM.

In this step, get ready to:

• Implement the optimizer and learning rate schedule.
• Train for 3000 iterations and plot the loss curves.
• Evaluate the pre-trained model by generating Shakespeare-like text.

Implement the optimizer and learning rate schedule

The model's forward() method already computes cross-entropy loss between its predictions and the actual next tokens. Cross-entropy is equivalent to maximizing the log-likelihood of the data: L = -1/T * sum(log P(x_t | x_<t)). Lower loss means the model assigns higher probability to the correct next token.

To minimize this loss, you need an optimizer and a learning rate schedule. The schedule starts with a gentle warmup (gradually increasing the learning rate) before decaying it with a cosine curve.

• Add the optimizer and learning rate schedule in a new cell below your model instantiation code:

# Set up the optimizer with weight decay for regularization
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

def get_lr(step):
    # Linear warmup for the first 100 steps
    if step < warmup_steps:
        return learning_rate * step / warmup_steps
    # Cosine decay after warmup
    decay_ratio = (step - warmup_steps) / (max_iters - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

What does this code do?

The AdamW optimizer updates model parameters using adaptive learning rates per-parameter, plus weight decay (a small penalty on large weights that prevents overfitting).

The get_lr() function implements a two-phase schedule. For the first 100 steps, the learning rate ramps up linearly from 0 to 3e-4 (warmup prevents early instability). After warmup, it smoothly decays following a cosine curve down to 3e-5 (10% of peak). This lets the model make big jumps early and fine-tune gently later.

Why warmup then decay?

At step 0, the model's gradients are noisy because every parameter is random. A large learning rate here can push the model into a bad region it never recovers from. Warmup lets the optimizer "feel out" the loss landscape before committing to big steps.

Cosine decay at the end ensures the model converges to a stable minimum instead of oscillating around it.

Run the pre-training loop

Now you will write the loop that actually trains your model. Each iteration samples a random batch, computes the loss, and updates the weights. Every 100 steps, it evaluates on validation data so you can monitor for overfitting.

• Add the training loop in the same cell, below the get_lr function:

train_losses = []
val_losses = []

for step in range(max_iters):
    # Update learning rate per step
    lr = get_lr(step)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Forward pass: get predictions and loss
    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)

    # Backward pass: compute gradients and update weights
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

    # Evaluate every 100 steps
    if step % eval_interval == 0:
        model.eval()
        with torch.no_grad():
            xv, yv = get_batch('val')
            _, val_loss = model(xv, yv)
        model.train()
        train_losses.append(loss.item())
        val_losses.append(val_loss.item())
        print(f"Step {step}: train loss {loss.item():.4f}, val loss {val_loss.item():.4f}")

What does this code do?

• get_lr(step) computes the current learning rate and manually sets it on the optimizer (this is the warmup + cosine schedule in action).
• get_batch('train') pulls a random batch of 32 sequences, each 128 characters long. The target is the same sequence shifted one position forward.
• model(xb, yb) runs the forward pass and computes cross-entropy loss against the true next tokens.
• loss.backward() computes gradients via backpropagation. optimizer.step() updates all parameters using those gradients.
• Every 100 steps, the loop switches to eval mode, checks validation loss, and logs both values. This tells you whether the model is learning (train loss drops) without memorizing (val loss stays close to train loss).

• Run this cell. Training takes approximately 3-5 minutes on a T4 GPU.

You should see the training loss drop from around 4.5 down to below 1.5 as the loop progresses. The validation loss should follow a similar curve.

Training is very slow or loss isn't dropping?

• Check that your runtime is set to GPU. Click **Runtime**, then **Change runtime type**, and confirm **T4 GPU** is selected. If you are on CPU, training will take around 15 minutes instead of 3-5.
• If loss stays flat above 4.0 after 500 steps, restart the runtime and re-run all cells from the top. Random initialization occasionally lands in a bad spot.

Still stuck? Help me debug why my training loss isn't decreasing during pre-training.

Once training completes, plot the loss curves to visualize the learning trajectory.

• Add the plotting code in a new cell and run it:

# Visualize the training and validation loss curves
plt.figure(figsize=(10, 4))
plt.plot(train_losses, label='Train')
plt.plot(val_losses, label='Val')
plt.xlabel('Evaluation Step')
plt.ylabel('Loss')
plt.title('Pre-training Loss (Next-Token Prediction)')
plt.legend()
plt.show()

What should I see?

You should see a smooth downward curve for both train and validation loss. Train loss starts around 4.5 (the model is nearly random) and drops to around 1.2-1.5 (the model predicts the next character correctly much of the time). If the validation curve stays close to the training curve, your model is generalizing well and not just memorizing.

Generate text from the pre-trained model

The true test of pre-training is whether the model generates coherent text. Before training, your model output random characters. Now it should produce something that looks like Shakespeare.

• Generate a 200-character sample by running this in a new cell:

# Generate text starting from a single zero token (beginning of sequence)
context = torch.zeros((1, 1), dtype=torch.long, device=device)
print("Generated text after pre-training:")
print(decode(model.generate(context, max_new_tokens=200)[0].tolist()))

What does this code do?

It creates a single-token context (the zero token, which is a newline in our character set) and uses the generate() method to autoregressively produce 200 new tokens. Each token is sampled from the model's predicted probability distribution, with temperature=0.8 adding slight randomness and top-k=40 filtering out unlikely characters.

Your output should contain real English words, proper line breaks, and character names that look like Shakespeare. It will not make perfect sense, but the structure is unmistakably language.

What just happened?

From purely random parameters, your model learned grammar, word structure, and even some thematic patterns. All from one simple objective: predict the next character. This is the exact same objective GPT-4 was trained with, just at a smaller scale.

The model never saw explicit rules about spelling, punctuation, or dialogue formatting. It discovered all of these patterns implicitly by minimizing the cross-entropy loss over millions of next-character predictions.

Output still looks like random characters?

• Confirm your final training loss was below 1.5. If it is still above 2.0, the model did not train enough. Re-run the training cell and let it complete all 3000 steps.
• Make sure you are generating from the trained model variable, not from a freshly instantiated MiniGPT().

Need more help? My model's generated text still looks random after training. How do I debug this?

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point:

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'

n_embd = 128
n_head = 4
n_layer = 4
block_size = 128
dropout = 0.1
batch_size = 32

max_iters = 3000
eval_interval = 100
learning_rate = 3e-4
warmup_steps = 100
min_lr = 3e-5
weight_decay = 0.01

import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, 'input.txt')

with open('input.txt', 'r') as f:
    text = f.read()

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

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

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

class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        )

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(n_embd, dim=2)
        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)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

class FeedForward(nn.Module):
    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 TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffn = FeedForward()

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

class MiniGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.token_emb.weight = self.lm_head.weight
        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        B, T = idx.size()
        tok_emb = self.token_emb(idx)
        pos_emb = self.pos_emb(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        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=0.8, top_k=40):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

model = MiniGPT().to(device)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

def get_lr(step):
    if step < warmup_steps:
        return learning_rate * step / warmup_steps
    decay_ratio = (step - warmup_steps) / (max_iters - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

train_losses = []
val_losses = []

for step in range(max_iters):
    lr = get_lr(step)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    if step % eval_interval == 0:
        model.eval()
        with torch.no_grad():
            xv, yv = get_batch('val')
            _, val_loss = model(xv, yv)
        model.train()
        train_losses.append(loss.item())
        val_losses.append(val_loss.item())
        print(f"Step {step}: train loss {loss.item():.4f}, val loss {val_loss.item():.4f}")

plt.figure(figsize=(10, 4))
plt.plot(train_losses, label='Train')
plt.plot(val_losses, label='Val')
plt.xlabel('Evaluation Step')
plt.ylabel('Loss')
plt.title('Pre-training Loss (Next-Token Prediction)')
plt.legend()
plt.show()

context = torch.zeros((1, 1), dtype=torch.long, device=device)
print("Generated text after pre-training:")
print(decode(model.generate(context, max_new_tokens=200)[0].tolist()))

Your model just learned language from scratch. Next up, you will teach it to follow instructions using supervised fine-tuning and then align its preferences with DPO.

Fine-tune with SFT and DPO

Your pre-trained model can generate Shakespeare-like text, but it has no concept of following instructions or producing answers to questions. It learned the structure of language, not how to be useful.

This is exactly the gap that real LLM companies bridge after pre-training. First, Supervised Fine-Tuning (SFT) teaches the model to follow a question-answer format. Then, Direct Preference Optimization (DPO) teaches it to prefer high-quality answers over low-quality ones. In this step, you will implement both.

In this step, get ready to:

• Fine-tune your pre-trained model on synthetic instruction data using masked cross-entropy loss.
• Implement the DPO loss function and train a preference-aligned model.
• Compare generated text and loss curves across all three training stages.

Create the SFT instruction dataset

SFT works by showing the model curated examples of the format you want it to produce. The model learns to replicate that format by predicting the answer tokens, while ignoring the prompt tokens during loss calculation.

• Create the synthetic instruction dataset by adding this code in a new cell below your pre-training output:

# Synthetic instruction dataset in Q/A format
sft_data = [
    "Q: What is the king's name?\nA: The king's name is Richard.\n",
    "Q: Who speaks first?\nA: The Duke speaks first.\n",
    "Q: Where does the scene take place?\nA: The scene takes place in the castle.\n",
    "Q: What season is it?\nA: It is winter in the kingdom.\n",
    "Q: Who is the villain?\nA: The villain is the false knight.\n",
    "Q: What does the lady want?\nA: The lady wants freedom from the tower.\n",
    "Q: How does the play end?\nA: The play ends with a royal wedding.\n",
    "Q: Who carries the sword?\nA: The brave soldier carries the sword.\n",
    "Q: What is the prophecy?\nA: The prophecy speaks of a new dawn.\n",
    "Q: Where is the treasure?\nA: The treasure is hidden beneath the throne.\n",
]

Why Shakespeare-themed Q/A pairs?

The model already knows Shakespeare vocabulary from pre-training. Using Shakespeare-themed questions lets SFT focus on teaching the format (question then answer) rather than fighting to introduce brand-new vocabulary. In production, SFT datasets contain thousands of examples. Ten is enough for our small model to demonstrate the principle.

• Run the cell. No output is expected since this just defines the data list.

Implement the SFT batch function with prompt masking

The key insight of SFT is that you only compute loss on the answer tokens. The prompt tokens are masked with -100 (PyTorch's ignore index), so the model learns to generate good answers without being penalized for "predicting" the question.

• Implement the batch function that handles this masking by adding this code in a new cell:

def get_sft_batch(batch_size=4):
    indices = torch.randint(len(sft_data), (batch_size,))
    max_len = block_size
    input_ids_list = []
    labels_list = []
    for idx in indices:
        text = sft_data[idx.item()]
        tokens = encode(text)[:max_len]
        input_ids = torch.tensor(tokens, dtype=torch.long)
        # Find where the answer starts (after "A: ")
        answer_start = text.find("\nA: ") + 4
        answer_start_token = len(encode(text[:answer_start]))
        # Labels: -100 for prompt tokens (masked), actual tokens for answer
        labels = torch.full_like(input_ids, -100)
        labels[answer_start_token:] = input_ids[answer_start_token:]
        # Pad to max_len
        pad_len = max_len - len(input_ids)
        if pad_len > 0:
            input_ids = F.pad(input_ids, (0, pad_len), value=0)
            labels = F.pad(labels, (0, pad_len), value=-100)
        input_ids_list.append(input_ids)
        labels_list.append(labels)
    return torch.stack(input_ids_list).to(device), torch.stack(labels_list).to(device)

What does this code do?

For each training example, this function splits the text into prompt and answer at the "\nA: " boundary. Every token position before the answer gets a label of -100, which tells PyTorch's cross-entropy to skip those positions entirely.

This is identical to how production SFT works. The model sees the full conversation, but only learns to predict the assistant's response. The prompt tokens provide context without contributing to the gradient.

• Run the cell. No output is expected yet.

Seeing a KeyError or encoding issue?

Make sure the sft_data cell ran successfully before this one. The apostrophes in strings like "king's" must use straight quotes, not curly quotes.

Help me debug my get_sft_batch function

Run the SFT training loop

SFT uses the same cross-entropy loss as pre-training, but only on answer tokens. You also use a lower learning rate to avoid destroying what the model already knows.

• Train the SFT model by adding and running this code in a new cell:

# Deep copy pre-trained model to preserve original weights
sft_model = copy.deepcopy(model)
sft_optimizer = torch.optim.AdamW(sft_model.parameters(), lr=1e-4, weight_decay=0.01)
sft_losses = []

for step in range(200):
    input_ids, labels = get_sft_batch()
    logits, _ = sft_model(input_ids)
    # Shift for next-token prediction alignment
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    # Cross-entropy ignoring -100 positions (prompt tokens)
    loss = F.cross_entropy(shift_logits.view(-1, vocab_size), shift_labels.view(-1), ignore_index=-100)
    sft_optimizer.zero_grad(set_to_none=True)
    loss.backward()
    sft_optimizer.step()
    if step % 50 == 0:
        sft_losses.append(loss.item())
        print(f"SFT Step {step}: loss {loss.item():.4f}")

What does this code do?

The loop copies your pre-trained model with copy.deepcopy(model) so you keep the original intact for comparison later. The learning rate is 1e-4 (three times lower than pre-training) to make gentle updates that preserve existing knowledge.

The shift operation aligns predictions with targets. Position i of logits predicts position i+1 of the input, which is the standard next-token prediction setup. The ignore_index=-100 argument makes cross-entropy skip all masked prompt positions.

• Run the cell. You should see loss values printed every 50 steps, starting around 3-4 and dropping below 1.0 by step 150.

Loss not decreasing?

If the loss stays flat or oscillates above 3.0, confirm that your pre-training ran successfully in the previous step. The SFT model is initialized from the pre-trained weights, so a bad starting point means SFT has too much to learn in 200 steps.

Help me debug my SFT training loss

Why DPO instead of RLHF?

Traditional RLHF (Reinforcement Learning from Human Feedback) requires training a separate reward model and running a complex PPO loop with a value network. DPO achieves the same goal (learning human preferences) by deriving a closed-form loss that works directly on preference pairs. No reward model, no RL loop, and the same mathematical objective.

Define the preference data and DPO helpers

DPO needs pairs of completions for the same prompt: one "chosen" (preferred) and one "rejected" (dispreferred). The model learns to increase the probability of the chosen completion relative to the rejected one.

• Define your preference pairs by adding this code in a new cell:

# Preference pairs: (prompt, chosen completion, rejected completion)
preference_data = [
    ("Q: Who is king?\nA: ", "The rightful king is Henry.", "king king king king"),
    ("Q: What happened?\nA: ", "A great battle was fought at dawn.", "nothing nothing nothing"),
    ("Q: Where are we?\nA: ", "We stand in the royal court of England.", "asdlfkj random words"),
    ("Q: Who speaks?\nA: ", "The noble Lord Gloucester speaks.", "speak speak speak"),
    ("Q: What time?\nA: ", "It is the hour before midnight.", "time time time time"),
]

What makes a good preference pair?

Each tuple contains the same prompt with two possible completions. The chosen response is coherent and informative. The rejected response is repetitive nonsense. In production, these pairs come from human annotators ranking model outputs. Here, the contrast is exaggerated so our tiny model can learn the signal in just 100 steps.

• Run the cell. No output is expected.

• Now implement the helper function that computes log-probabilities over response tokens, in a new cell:

def compute_sequence_logprobs(model_to_eval, token_ids, response_start):
    """Compute sum of log-probabilities for the response tokens."""
    with torch.set_grad_enabled(model_to_eval.training):
        logits, _ = model_to_eval(token_ids)
        log_probs = F.log_softmax(logits, dim=-1)
        # Gather log-probs for actual next tokens
        target_ids = token_ids[:, 1:]
        gathered = log_probs[:, :-1, :].gather(2, target_ids.unsqueeze(-1)).squeeze(-1)
        # Sum only over response tokens
        mask = torch.zeros_like(gathered)
        mask[:, response_start:] = 1.0
        return (gathered * mask).sum(dim=-1)

What does this function compute?

For a given sequence of tokens, this computes the total log-probability that the model assigns to the response portion only (everything after response_start). DPO needs this value for both the "chosen" and "rejected" completions to compare them.

The torch.set_grad_enabled(model_to_eval.training) trick enables gradients only when the model is in training mode. This lets you reuse the same function for both the policy model (needs gradients) and the frozen reference model (no gradients).

• Run the cell. No output is expected.

• Implement the DPO loss function in a new cell:

def dpo_loss(policy_chosen_logps, policy_rejected_logps, ref_chosen_logps, ref_rejected_logps, beta=0.1):
    """Standard DPO loss."""
    # How much more the policy prefers chosen over rejected
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    # How much more the reference prefers chosen over rejected
    ref_logratios = ref_chosen_logps - ref_rejected_logps
    # DPO logits: policy preference minus reference preference
    logits = pi_logratios - ref_logratios
    # Negative log-sigmoid: pushes logits positive (prefer chosen)
    loss = -F.logsigmoid(beta * logits).mean()
    # Implicit rewards for monitoring
    chosen_reward = beta * (policy_chosen_logps - ref_chosen_logps).detach()
    rejected_reward = beta * (policy_rejected_logps - ref_rejected_logps).detach()
    return loss, chosen_reward, rejected_reward

What is the DPO loss doing mathematically?

The loss computes how much the policy model prefers the chosen completion over the rejected one, relative to a frozen reference model. The logsigmoid turns this into a binary classification-style loss: push the policy to prefer the chosen completion more than the reference does.

The beta=0.1 parameter controls how far the policy can deviate from the reference. Lower beta means the model is allowed to shift its preferences more aggressively. The reference model prevents the policy from drifting too far and forgetting how to produce coherent language.

• Run the cell. No output is expected.

NameError on F.logsigmoid?

Make sure you ran the very first cell that imports torch.nn.functional as F. If your Colab runtime restarted, you need to re-run all cells from the top.

Help me fix my DPO loss function

Train the DPO model

DPO requires a frozen reference model (the SFT model) and a trainable policy model (a copy of the SFT model that learns preferences). The reference acts as an anchor so the policy cannot drift arbitrarily far.

• Run the DPO training loop by adding and running this code in a new cell:

# Freeze the SFT model as our reference (anchor)
ref_model = copy.deepcopy(sft_model).eval()
for p in ref_model.parameters():
    p.requires_grad = False

# Policy model starts as a copy of SFT
dpo_model = copy.deepcopy(sft_model)
dpo_optimizer = torch.optim.AdamW(dpo_model.parameters(), lr=5e-5, weight_decay=0.01)
dpo_losses = []
reward_margins = []

for step in range(100):
    # Cycle through preference pairs
    idx = step % len(preference_data)
    prompt, chosen, rejected = preference_data[idx]
    chosen_text = prompt + chosen
    rejected_text = prompt + rejected
    response_start = len(encode(prompt)) - 1
    chosen_ids = torch.tensor([encode(chosen_text)[:block_size]], dtype=torch.long, device=device)
    rejected_ids = torch.tensor([encode(rejected_text)[:block_size]], dtype=torch.long, device=device)
    # Policy log-probs (gradients enabled)
    dpo_model.train()
    policy_chosen_logps = compute_sequence_logprobs(dpo_model, chosen_ids, response_start)
    policy_rejected_logps = compute_sequence_logprobs(dpo_model, rejected_ids, response_start)
    # Reference log-probs (no gradients)
    with torch.no_grad():
        ref_chosen_logps = compute_sequence_logprobs(ref_model, chosen_ids, response_start)
        ref_rejected_logps = compute_sequence_logprobs(ref_model, rejected_ids, response_start)
    loss, chosen_reward, rejected_reward = dpo_loss(
        policy_chosen_logps, policy_rejected_logps,
        ref_chosen_logps, ref_rejected_logps, beta=0.1
    )
    dpo_optimizer.zero_grad(set_to_none=True)
    loss.backward()
    dpo_optimizer.step()
    if step % 20 == 0:
        dpo_losses.append(loss.item())
        margin = (chosen_reward - rejected_reward).mean().item()
        reward_margins.append(margin)
        print(f"DPO Step {step}: loss {loss.item():.4f}, reward margin {margin:.4f}")

What should I see?

The first line should print a loss near 0.693. This is -log(0.5), meaning the model has no preference between chosen and rejected. As training progresses, the loss should decrease and the reward margin should become more positive, showing the model learning to prefer the chosen completions.

• Run the cell. Confirm the DPO loss starts near 0.693 at step 0 and decreases over 100 steps.

DPO loss stays exactly at 0.693?

If the loss never moves, the policy and reference are producing identical log-probabilities. Verify that ref_model has requires_grad = False on all parameters and that dpo_model is a separate copy.deepcopy (not the same object as ref_model).

Help me debug my DPO training loop

Compare all three training stages

Now for the payoff. You will feed the same prompt to all three models and see how each training stage shaped the model's behavior differently.

• Generate the three-stage comparison and plot all loss curves by adding and running this code in a new cell:

print("\n" + "="*60)
print("COMPARISON: Same prompt, three training stages")
print("="*60)

prompt_text = "Q: Who is the king?\nA: "
prompt_ids = torch.tensor([encode(prompt_text)], dtype=torch.long, device=device)

# Generate from each model with the same prompt
for name, m in [("Pre-trained", model), ("SFT", sft_model), ("DPO", dpo_model)]:
    m.eval()
    output = m.generate(prompt_ids, max_new_tokens=50, temperature=0.7)
    generated = decode(output[0].tolist())[len(prompt_text):]
    print(f"\n[{name}]: {generated[:100]}")

# Plot all three loss curves side by side
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
axes[0].plot(train_losses)
axes[0].set_title('Pre-training Loss')
axes[0].set_xlabel('Step (x100)')
axes[1].plot(sft_losses)
axes[1].set_title('SFT Loss')
axes[1].set_xlabel('Step (x50)')
axes[2].plot(dpo_losses)
axes[2].set_title('DPO Loss')
axes[2].set_xlabel('Step (x20)')
plt.tight_layout()
plt.show()

What should the comparison show?

The pre-trained model will produce Shakespeare-style continuation that ignores the Q/A format entirely. The SFT model will respond in the expected Q/A format with a coherent answer. The DPO model should produce a similar format but with slightly more articulate or preferred-style responses.

The loss curves tell the story of what each stage learned. Pre-training loss drops from ~4.5 to below 1.5 as the model learns language. SFT loss drops quickly as the model memorizes the format. DPO loss starts at 0.693 (random preference) and decreases as preferences are learned.

• Run the cell. You should see three different text outputs showing how each model responds to the same prompt, followed by a three-panel loss curve plot.

All three models producing identical output?

This usually means the SFT and DPO training did not actually modify their respective model copies. Check that you used copy.deepcopy(model) for sft_model and copy.deepcopy(sft_model) for dpo_model. If you used simple assignment (sft_model = model), both names point to the same object and training modified the original.

Help me debug my model comparison output

✔️ Awesome, I've got everything!

Great. Double check you have run all cells in order and your three-panel plot is visible.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point. Compare it against your cells to find any differences.

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import copy
import math

device = 'cuda' if torch.cuda.is_available() else 'cpu'

n_embd = 128
n_head = 4
n_layer = 4
block_size = 128
dropout = 0.1
batch_size = 32

max_iters = 3000
eval_interval = 100
learning_rate = 3e-4
warmup_steps = 100
min_lr = 3e-5
weight_decay = 0.01

import urllib.request
url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
urllib.request.urlretrieve(url, 'input.txt')

with open('input.txt', 'r') as f:
    text = f.read()

chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

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

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

class CausalSelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)
        )

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(n_embd, dim=2)
        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)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

class FeedForward(nn.Module):
    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 TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention()
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffn = FeedForward()

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

class MiniGPT(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.token_emb.weight = self.lm_head.weight
        self.apply(self._init_weights)

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

    def forward(self, idx, targets=None):
        B, T = idx.size()
        tok_emb = self.token_emb(idx)
        pos_emb = self.pos_emb(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        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=0.8, top_k=40):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

model = MiniGPT().to(device)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

def get_lr(step):
    if step < warmup_steps:
        return learning_rate * step / warmup_steps
    decay_ratio = (step - warmup_steps) / (max_iters - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)

train_losses = []
val_losses = []

for step in range(max_iters):
    lr = get_lr(step)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    xb, yb = get_batch('train')
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    if step % eval_interval == 0:
        model.eval()
        with torch.no_grad():
            xv, yv = get_batch('val')
            _, val_loss = model(xv, yv)
        model.train()
        train_losses.append(loss.item())
        val_losses.append(val_loss.item())
        print(f"Step {step}: train loss {loss.item():.4f}, val loss {val_loss.item():.4f}")

plt.figure(figsize=(10, 4))
plt.plot(train_losses, label='Train')
plt.plot(val_losses, label='Val')
plt.xlabel('Evaluation Step')
plt.ylabel('Loss')
plt.title('Pre-training Loss (Next-Token Prediction)')
plt.legend()
plt.show()

context = torch.zeros((1, 1), dtype=torch.long, device=device)
print("Generated text after pre-training:")
print(decode(model.generate(context, max_new_tokens=200)[0].tolist()))

sft_data = [
    "Q: What is the king's name?\nA: The king's name is Richard.\n",
    "Q: Who speaks first?\nA: The Duke speaks first.\n",
    "Q: Where does the scene take place?\nA: The scene takes place in the castle.\n",
    "Q: What season is it?\nA: It is winter in the kingdom.\n",
    "Q: Who is the villain?\nA: The villain is the false knight.\n",
    "Q: What does the lady want?\nA: The lady wants freedom from the tower.\n",
    "Q: How does the play end?\nA: The play ends with a royal wedding.\n",
    "Q: Who carries the sword?\nA: The brave soldier carries the sword.\n",
    "Q: What is the prophecy?\nA: The prophecy speaks of a new dawn.\n",
    "Q: Where is the treasure?\nA: The treasure is hidden beneath the throne.\n",
]

def get_sft_batch(batch_size=4):
    indices = torch.randint(len(sft_data), (batch_size,))
    max_len = block_size
    input_ids_list = []
    labels_list = []
    for idx in indices:
        text = sft_data[idx.item()]
        tokens = encode(text)[:max_len]
        input_ids = torch.tensor(tokens, dtype=torch.long)
        answer_start = text.find("\nA: ") + 4
        answer_start_token = len(encode(text[:answer_start]))
        labels = torch.full_like(input_ids, -100)
        labels[answer_start_token:] = input_ids[answer_start_token:]
        pad_len = max_len - len(input_ids)
        if pad_len > 0:
            input_ids = F.pad(input_ids, (0, pad_len), value=0)
            labels = F.pad(labels, (0, pad_len), value=-100)
        input_ids_list.append(input_ids)
        labels_list.append(labels)
    return torch.stack(input_ids_list).to(device), torch.stack(labels_list).to(device)

sft_model = copy.deepcopy(model)
sft_optimizer = torch.optim.AdamW(sft_model.parameters(), lr=1e-4, weight_decay=0.01)
sft_losses = []

for step in range(200):
    input_ids, labels = get_sft_batch()
    logits, _ = sft_model(input_ids)
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    loss = F.cross_entropy(shift_logits.view(-1, vocab_size), shift_labels.view(-1), ignore_index=-100)
    sft_optimizer.zero_grad(set_to_none=True)
    loss.backward()
    sft_optimizer.step()
    if step % 50 == 0:
        sft_losses.append(loss.item())
        print(f"SFT Step {step}: loss {loss.item():.4f}")

preference_data = [
    ("Q: Who is king?\nA: ", "The rightful king is Henry.", "king king king king"),
    ("Q: What happened?\nA: ", "A great battle was fought at dawn.", "nothing nothing nothing"),
    ("Q: Where are we?\nA: ", "We stand in the royal court of England.", "asdlfkj random words"),
    ("Q: Who speaks?\nA: ", "The noble Lord Gloucester speaks.", "speak speak speak"),
    ("Q: What time?\nA: ", "It is the hour before midnight.", "time time time time"),
]

def compute_sequence_logprobs(model_to_eval, token_ids, response_start):
    with torch.set_grad_enabled(model_to_eval.training):
        logits, _ = model_to_eval(token_ids)
        log_probs = F.log_softmax(logits, dim=-1)
        target_ids = token_ids[:, 1:]
        gathered = log_probs[:, :-1, :].gather(2, target_ids.unsqueeze(-1)).squeeze(-1)
        mask = torch.zeros_like(gathered)
        mask[:, response_start:] = 1.0
        return (gathered * mask).sum(dim=-1)

def dpo_loss(policy_chosen_logps, policy_rejected_logps, ref_chosen_logps, ref_rejected_logps, beta=0.1):
    pi_logratios = policy_chosen_logps - policy_rejected_logps
    ref_logratios = ref_chosen_logps - ref_rejected_logps
    logits = pi_logratios - ref_logratios
    loss = -F.logsigmoid(beta * logits).mean()
    chosen_reward = beta * (policy_chosen_logps - ref_chosen_logps).detach()
    rejected_reward = beta * (policy_rejected_logps - ref_rejected_logps).detach()
    return loss, chosen_reward, rejected_reward

ref_model = copy.deepcopy(sft_model).eval()
for p in ref_model.parameters():
    p.requires_grad = False

dpo_model = copy.deepcopy(sft_model)
dpo_optimizer = torch.optim.AdamW(dpo_model.parameters(), lr=5e-5, weight_decay=0.01)
dpo_losses = []
reward_margins = []

for step in range(100):
    idx = step % len(preference_data)
    prompt, chosen, rejected = preference_data[idx]
    chosen_text = prompt + chosen
    rejected_text = prompt + rejected
    response_start = len(encode(prompt)) - 1
    chosen_ids = torch.tensor([encode(chosen_text)[:block_size]], dtype=torch.long, device=device)
    rejected_ids = torch.tensor([encode(rejected_text)[:block_size]], dtype=torch.long, device=device)
    dpo_model.train()
    policy_chosen_logps = compute_sequence_logprobs(dpo_model, chosen_ids, response_start)
    policy_rejected_logps = compute_sequence_logprobs(dpo_model, rejected_ids, response_start)
    with torch.no_grad():
        ref_chosen_logps = compute_sequence_logprobs(ref_model, chosen_ids, response_start)
        ref_rejected_logps = compute_sequence_logprobs(ref_model, rejected_ids, response_start)
    loss, chosen_reward, rejected_reward = dpo_loss(
        policy_chosen_logps, policy_rejected_logps,
        ref_chosen_logps, ref_rejected_logps, beta=0.1
    )
    dpo_optimizer.zero_grad(set_to_none=True)
    loss.backward()
    dpo_optimizer.step()
    if step % 20 == 0:
        dpo_losses.append(loss.item())
        margin = (chosen_reward - rejected_reward).mean().item()
        reward_margins.append(margin)
        print(f"DPO Step {step}: loss {loss.item():.4f}, reward margin {margin:.4f}")

print("\n" + "="*60)
print("COMPARISON: Same prompt, three training stages")
print("="*60)

prompt_text = "Q: Who is the king?\nA: "
prompt_ids = torch.tensor([encode(prompt_text)], dtype=torch.long, device=device)

for name, m in [("Pre-trained", model), ("SFT", sft_model), ("DPO", dpo_model)]:
    m.eval()
    output = m.generate(prompt_ids, max_new_tokens=50, temperature=0.7)
    generated = decode(output[0].tolist())[len(prompt_text):]
    print(f"\n[{name}]: {generated[:100]}")

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
axes[0].plot(train_losses)
axes[0].set_title('Pre-training Loss')
axes[0].set_xlabel('Step (x100)')
axes[1].plot(sft_losses)
axes[1].set_title('SFT Loss')
axes[1].set_xlabel('Step (x50)')
axes[2].plot(dpo_losses)
axes[2].set_title('DPO Loss')
axes[2].set_xlabel('Step (x20)')
plt.tight_layout()
plt.show()

Secret mission

Implement GRPO with Arithmetic Rewards

You just trained a model with DPO using offline preference pairs. But what if you could check whether the model's answer is objectively correct? GRPO (Group Relative Policy Optimization) is the algorithm behind DeepSeek-R1's reasoning capabilities. It replaces the learned reward model with a verifiable reward and eliminates the value network entirely. Your mission: implement GRPO with a simple arithmetic verifier and watch your model learn to do math.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on Google Colab's free tier, so there are no ongoing costs. Decide whether to keep your notebook available, disconnect the runtime to free GPU quota, or delete everything entirely.

Resources you used:

• Google Colab notebook (stored in your Google Drive).
• Downloaded input.txt file (Tiny Shakespeare dataset, stored in the Colab ephemeral filesystem).
• GPU runtime session (T4 GPU allocated from Colab's free quota).

Keep everything running

No action needed. Choose this if you want to revisit the notebook and experiment with different hyperparameters or training data.

• Your notebook stays in Google Drive and you can reopen it anytime.
• The Colab runtime will disconnect automatically after a period of inactivity, freeing the GPU.
• The downloaded input.txt file lives in the ephemeral filesystem and disappears when the runtime disconnects. It will be re-downloaded next time you run the cell.

Pause - I'll come back to this later

Disconnect your GPU runtime to free up quota for other Colab users. Your notebook stays saved in Google Drive.

• In Colab, click Runtime in the top menu bar.
• Click Disconnect and delete runtime.
• This frees the GPU and clears all variables from memory. Your notebook code and outputs remain saved in Google Drive.

Delete - I don't want to use this again

Remove the notebook from Google Drive entirely. This cannot be undone.

• First, disconnect the runtime: click Runtime in the top menu, then Disconnect and delete runtime.
• Open Google Drive at drive.google.com.
• Navigate to the Colab Notebooks folder (or wherever you saved your notebook).
• Right-click the notebook file and select Move to trash.
• Empty your Google Drive trash to permanently delete the file.

Nice Work!

Nice Work!

You just built a language model from raw math and trained it through every stage of the modern LLM pipeline. That is no small feat.

What you learned:

• Built a complete GPT model from scratch, implementing every component (self-attention, feed-forward networks, LayerNorm, positional embeddings, weight tying) with the underlying math explained.
• Trained the model through the full modern LLM pipeline: pre-training on raw text, supervised fine-tuning on instructions, and preference optimization with DPO.
• Implemented and understood three distinct loss functions and how each one shapes model behaviour differently: cross-entropy for language modeling, masked cross-entropy for instruction following, and the DPO objective for alignment.
• (Secret Mission) Implemented GRPO, the online RL algorithm behind DeepSeek-R1, with a verifiable arithmetic reward.

Ready to quiz yourself?