Fine-Tune a Language Model: SFT and DPO
Build a from-scratch SFT and DPO pipeline to fine-tune and align GPT-2.
Introduction
30 Second Summary
Every language model starts as a general text predictor. The difference between a raw model and a helpful assistant comes down to two post-training stages: teaching it your domain knowledge, then aligning its responses to human preferences.
In this project, you will build a complete fine-tuning pipeline from scratch that takes a pretrained GPT-2 model through Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) to produce an aligned model. You will write every training loop by hand to understand the mechanics behind modern LLM post-training.
What You'll Build
A from-scratch fine-tuning pipeline that takes your Project 3 JSONL training data and produces an aligned language model through three stages: SFT, preference optimization with DPO, and evaluation.
By the end of this project, you'll have:
- A working SFT training loop that fine-tunes GPT-2 (124M) on your Project 3 JSONL dataset, written from scratch without high-level trainer abstractions.
- A DPO alignment step implemented from scratch that adjusts the fine-tuned model using synthetic preference pairs.
- A quantitative evaluation comparing base GPT-2, SFT-tuned, and DPO-aligned outputs on the same prompts using perplexity and generation comparison.
- Secret Mission: Add a cosine decay learning rate schedule and reward margin tracking to visualize how DPO shifts model preferences over training steps.
Are there any prerequisites?
This project assumes you have completed Project 3 and have a JSONL dataset ready for training. You should have a basic understanding of Python and PyTorch (tensors, forward/backward passes).
Before We Start
Fine-tuning a language model involves multiple stages, and each stage requires specific tools. You need PyTorch for writing training loops from scratch, Hugging Face Transformers for loading pretrained model weights, and tiktoken for tokenization.
In this step, you'll prepare your development environment with all the dependencies needed to run the SFT and DPO training pipelines you'll build in later steps.
In this step, get ready to:
- Create a project folder and virtual environment.
- Install PyTorch, Hugging Face Transformers, and tiktoken.
- Download GPT-2 (124M) model weights and verify they load.
Create a virtual environment and install PyTorch and Transformers
A virtual environment keeps your project's packages isolated from the rest of your system. This prevents version conflicts between projects.
- Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type VS Code and press Enter to open it.
- Open a terminal in VS Code by pressing Ctrl+` (backtick key, same on both macOS and Windows).
- Move to your Desktop by running this command:
cd ~/Desktop- Create your project folder and move into it by running these commands:
mkdir fine-tune-llm
cd fine-tune-llm- Create a virtual environment by running:
python3 -m venv venv- Activate the virtual environment by running:
macOS
source venv/bin/activateWindows
venv\Scripts\Activate.ps1You should see (venv) appear at the beginning of your terminal prompt.
Don't see (venv) in your prompt?
On Windows, you may need to allow script execution first. Run Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned then try the activate command again.
Make sure you're inside the fine-tune-llm folder where you created the virtual environment.
Help me activate my Python virtual environment
- Install PyTorch (CPU version) by running:
pip install torchWhy CPU-only PyTorch?
The CPU version of PyTorch is smaller to download and works on any machine. You'll only run 50-100 training steps in this project, which takes a few minutes on CPU.
- Install Hugging Face Transformers by running:
pip install "transformers[torch]"The install may take a minute or two depending on your internet connection.
Seeing install errors?
Make sure your virtual environment is active (you should see (venv) in your prompt). If not, re-run the activate command above.
Help me fix my pip install error for PyTorch or Transformers
Install tiktoken and verify all packages load
GPT-2 uses byte-level BPE tokenization. The tiktoken library provides the canonical tokenizer for GPT-2, so your token IDs will match the model's vocabulary exactly.
- Install tiktoken by running:
pip install tiktoken- Verify all three packages load correctly by running this command:
python3 -c "import torch; import transformers; import tiktoken; print(f'PyTorch: {torch.__version__}'); print(f'Transformers: {transformers.__version__}'); print('tiktoken: OK')"You should see version numbers for PyTorch and Transformers, followed by tiktoken: OK.
Seeing ModuleNotFoundError?
This usually means the package didn't install into your virtual environment. Check that (venv) is visible in your prompt, then re-run the failing pip install command.
Help me fix a ModuleNotFoundError for my Python packages
Download GPT-2 model weights and verify they load
GPT-2 (124M parameters) is the model you'll fine-tune throughout this project. Downloading it now means you won't have to wait during the training steps later.
- Download and verify the GPT-2 model by running this command:
python3 -c "from transformers import GPT2LMHeadModel; model = GPT2LMHeadModel.from_pretrained('openai-community/gpt2'); print(f'Model loaded: {model.config.n_layer} layers, {model.config.n_embd} embedding dim'); print(f'Parameters: {sum(p.numel() for p in model.parameters()):,}')"What's happening here?
GPT2LMHeadModel.from_pretrained('openai-community/gpt2') downloads the pretrained GPT-2 weights from Hugging Face Hub (about 500MB) and loads them into memory.
The model has 12 layers, 12 attention heads, and 768-dimensional embeddings. You'll use this as the starting point for both SFT and DPO training.
You should see output showing 12 layers, 768 embedding dim, and approximately 124 million parameters.
Download taking a long time or failing?
The first download is about 500MB. If your connection is slow, give it a few minutes. If it fails, check your internet connection and try again. The model is cached locally after the first successful download.
Help me download the GPT-2 model from Hugging Face
- Copy your Project 3 JSONL dataset into the project folder. Your training_data.jsonl file should be inside the fine-tune-llm folder.
- Confirm the file is in place by running:
lsYou should see training_data.jsonl and venv listed in the output.
Your environment is ready. Next up, you'll build the data loader that converts your JSONL dataset into tokenized training batches.
Build the Data Loader
Your environment is ready with PyTorch, Transformers, and tiktoken installed. Now it's time to build the bridge between your raw JSONL training data and the model.
A language model learns by predicting the next token in a sequence. Your data loader needs to read each JSONL example, tokenize it, and shift the tokens by one position so the model always has an input and a corresponding target to learn from.
In this step, get ready to:
- Write a dataset class that reads JSONL and tokenizes examples for causal language modeling.
- Implement a collation function that pads variable-length sequences into uniform batches.
- Verify the data loader outputs correctly shaped tensors from your training data.
Write the JSONLDataset class
Causal language modeling works by training a model to predict token N+1 given tokens 1 through N. To set this up, you split every tokenized sequence into an input (all tokens except the last) and a label (all tokens except the first). This one-position shift creates the prediction targets.
- In VS Code, create a new file called data_loader.py in your project folder.
- Add the following code to define the JSONLDataset class:
import json
import torch
from torch.utils.data import Dataset, DataLoader
class JSONLDataset(Dataset):
"""Loads JSONL file and tokenizes for causal language modeling."""
def __init__(self, jsonl_path, tokenizer, max_length=256):
self.examples = []
self.tokenizer = tokenizer
self.max_length = max_length
with open(jsonl_path, "r", encoding="utf-8") as f:
for line in f:
data = json.loads(line.strip())
# Combine instruction and response into a single training text
if "instruction" in data and "response" in data:
text = f"### Instruction:\n{data['instruction']}\n\n### Response:\n{data['response']}"
elif "text" in data:
text = data["text"]
elif "prompt" in data and "completion" in data:
text = f"{data['prompt']}{data['completion']}"
else:
continue
self.examples.append(text)
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
text = self.examples[idx]
token_ids = self.tokenizer.encode(text)
# Truncate to max_length
token_ids = token_ids[: self.max_length]
# Input is all tokens except last, target is all tokens except first
input_ids = torch.tensor(token_ids[:-1], dtype=torch.long)
labels = torch.tensor(token_ids[1:], dtype=torch.long)
return input_ids, labelsWhat does this code do?
The JSONLDataset class handles three common JSONL formats from Project 3: instruction/response pairs, plain text, and prompt/completion pairs. It combines each example into one training string.
- __init__ reads every line in your JSONL file and stores each example as a formatted text string.
- __getitem__ tokenizes the text, truncates to max_length, then shifts by one position: input_ids is all tokens except the last, labels is all tokens except the first. This creates the next-token prediction target.
- Save data_loader.py.
- Confirm the file exists by running:
lsYou should see data_loader.py alongside your training_data.jsonl file in the output.
Don't see data_loader.py?
Make sure you saved the file in the same directory where your terminal session is running. Check your current directory with pwd (macOS/Linux) or cd (Windows) and confirm it matches the folder shown in VS Code's Explorer sidebar.
Help me find where my file was saved
Add batching with padding and attention masks
Different examples in your JSONL file tokenize to different lengths. To train in batches, every sequence in a batch must be the same length. A collation function pads shorter sequences and creates an attention mask so the model knows which tokens are real and which are padding.
- In data_loader.py, add the collate_fn function below the JSONLDataset class:
def collate_fn(batch, pad_token_id=50256):
"""Pads sequences in a batch to the same length."""
input_ids_list, labels_list = zip(*batch)
# Find max length in batch
max_len = max(ids.size(0) for ids in input_ids_list)
padded_inputs = []
padded_labels = []
attention_masks = []
for input_ids, labels in zip(input_ids_list, labels_list):
pad_len = max_len - input_ids.size(0)
# Pad on the right
padded_inputs.append(
torch.cat([input_ids, torch.full((pad_len,), pad_token_id, dtype=torch.long)])
)
padded_labels.append(
torch.cat([labels, torch.full((pad_len,), -100, dtype=torch.long)])
)
attention_masks.append(
torch.cat([torch.ones(input_ids.size(0), dtype=torch.long), torch.zeros(pad_len, dtype=torch.long)])
)
return {
"input_ids": torch.stack(padded_inputs),
"labels": torch.stack(padded_labels),
"attention_mask": torch.stack(attention_masks),
}What does this code do?
- pad_token_id=50256 uses GPT-2's end-of-text token as the padding token. This is standard for GPT-2 since it has no dedicated pad token.
- Labels are padded with -100, which tells PyTorch's cross_entropy loss to ignore those positions. The model won't be penalized for padding predictions.
- The attention mask is 1 for real tokens and 0 for padding. This tells the model's attention mechanism which positions to attend to.
- Now add the create_dataloader function below collate_fn:
def create_dataloader(jsonl_path, tokenizer, batch_size=4, max_length=256, shuffle=True):
"""Creates a DataLoader from a JSONL file."""
dataset = JSONLDataset(jsonl_path, tokenizer, max_length=max_length)
loader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=shuffle,
collate_fn=lambda batch: collate_fn(batch, pad_token_id=50256),
)
return loaderWhat does this code do?
This is a convenience wrapper that ties everything together. It creates the dataset, then wraps it in PyTorch's DataLoader which handles batching and shuffling. The lambda passes the correct pad token ID to your collate_fn.
- Save data_loader.py.
You won't see a visible change yet. The next substep will run the data loader and confirm everything works.
Verify the data loader output
Before moving on to training, you need to confirm the data loader produces tensors with the right shapes. A quick verification script will load a sample batch and print its dimensions.
- At the bottom of data_loader.py, add this verification block:
if __name__ == "__main__":
import tiktoken
tokenizer = tiktoken.get_encoding("gpt2")
# Replace with your actual JSONL file path from Project 3
loader = create_dataloader("training_data.jsonl", tokenizer, batch_size=2, max_length=128)
# Verify shapes
batch = next(iter(loader))
print(f"input_ids shape: {batch['input_ids'].shape}")
print(f"labels shape: {batch['labels'].shape}")
print(f"attention_mask shape: {batch['attention_mask'].shape}")
print(f"Sample tokens: {batch['input_ids'][0][:20]}")What does this code do?
- tiktoken.get_encoding("gpt2") loads the byte-level BPE tokenizer that matches GPT-2's vocabulary.
- The script creates a loader with batch_size=2 and max_length=128 (smaller values for a quick test), pulls one batch, and prints the tensor shapes to confirm everything is wired correctly.
- Save data_loader.py.
- Run the script by executing this command in your terminal:
python data_loader.pyYou should see output showing tensor shapes like input_ids shape: torch.Size([2, N]) where 2 is the batch size and N is the sequence length. The labels and attention_mask tensors should have the same shape.
Seeing a FileNotFoundError?
This means your training_data.jsonl file is not in the same directory as data_loader.py. Either move your JSONL file to the project folder or update the filename string in the __main__ block to match the actual path to your Project 3 JSONL file.
If your JSONL file has a different name, update "training_data.jsonl" in the code to match your file's actual name.
Help me fix my data loader file path
Getting an empty dataset or zero-length error?
Your JSONL file might use a format the loader doesn't recognize. The loader handles three formats: instruction/response pairs, a text field, or prompt/completion pairs. Open your JSONL file and check that each line matches one of these structures.
Help me debug my JSONL format
✔️ Awesome, I've got everything!
Great. Double check you've saved data_loader.py and that your terminal output shows matching shapes for all three tensors.
ⓧ I'd like to double check the full code
Here is the complete data_loader.py file for comparison:
import json
import torch
from torch.utils.data import Dataset, DataLoader
class JSONLDataset(Dataset):
"""Loads JSONL file and tokenizes for causal language modeling."""
def __init__(self, jsonl_path, tokenizer, max_length=256):
self.examples = []
self.tokenizer = tokenizer
self.max_length = max_length
with open(jsonl_path, "r", encoding="utf-8") as f:
for line in f:
data = json.loads(line.strip())
# Combine instruction and response into a single training text
if "instruction" in data and "response" in data:
text = f"### Instruction:\n{data['instruction']}\n\n### Response:\n{data['response']}"
elif "text" in data:
text = data["text"]
elif "prompt" in data and "completion" in data:
text = f"{data['prompt']}{data['completion']}"
else:
continue
self.examples.append(text)
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
text = self.examples[idx]
token_ids = self.tokenizer.encode(text)
# Truncate to max_length
token_ids = token_ids[: self.max_length]
# Input is all tokens except last, target is all tokens except first
input_ids = torch.tensor(token_ids[:-1], dtype=torch.long)
labels = torch.tensor(token_ids[1:], dtype=torch.long)
return input_ids, labels
def collate_fn(batch, pad_token_id=50256):
"""Pads sequences in a batch to the same length."""
input_ids_list, labels_list = zip(*batch)
# Find max length in batch
max_len = max(ids.size(0) for ids in input_ids_list)
padded_inputs = []
padded_labels = []
attention_masks = []
for input_ids, labels in zip(input_ids_list, labels_list):
pad_len = max_len - input_ids.size(0)
# Pad on the right
padded_inputs.append(
torch.cat([input_ids, torch.full((pad_len,), pad_token_id, dtype=torch.long)])
)
padded_labels.append(
torch.cat([labels, torch.full((pad_len,), -100, dtype=torch.long)])
)
attention_masks.append(
torch.cat([torch.ones(input_ids.size(0), dtype=torch.long), torch.zeros(pad_len, dtype=torch.long)])
)
return {
"input_ids": torch.stack(padded_inputs),
"labels": torch.stack(padded_labels),
"attention_mask": torch.stack(attention_masks),
}
def create_dataloader(jsonl_path, tokenizer, batch_size=4, max_length=256, shuffle=True):
"""Creates a DataLoader from a JSONL file."""
dataset = JSONLDataset(jsonl_path, tokenizer, max_length=max_length)
loader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=shuffle,
collate_fn=lambda batch: collate_fn(batch, pad_token_id=50256),
)
return loader
if __name__ == "__main__":
import tiktoken
tokenizer = tiktoken.get_encoding("gpt2")
# Replace with your actual JSONL file path from Project 3
loader = create_dataloader("training_data.jsonl", tokenizer, batch_size=2, max_length=128)
# Verify shapes
batch = next(iter(loader))
print(f"input_ids shape: {batch['input_ids'].shape}")
print(f"labels shape: {batch['labels'].shape}")
print(f"attention_mask shape: {batch['attention_mask'].shape}")
print(f"Sample tokens: {batch['input_ids'][0][:20]}")Your data pipeline is ready. Next up, you'll write the SFT training loop that feeds these batches through GPT-2 and teaches it to predict your training data.
Build Shared Utilities
Your data loader is ready to serve tokenized batches. Before writing the training loops, you need two helper functions that both SFT and DPO will depend on.
The first helper computes log-probabilities for token sequences. The second generates text from a prompt using top-k sampling. Both live in a shared utils.py file so your training and evaluation scripts can import them without duplicating code.
In this step, get ready to:
- Write a log-probability computation function for scoring model outputs.
- Write a text generation function with top-k sampling.
- Verify both utilities are importable from your training scripts.
Write the log-probability computation function
DPO needs to compare how likely a response is under two different models. That comparison happens through log-probabilities. This function takes a model's output logits, applies log-softmax, and gathers the probabilities for the actual tokens in the sequence.
- Create a new file called utils.py in your project folder.
- Add the imports and the compute_log_probs function by pasting the following code:
import torch
import torch.nn.functional as F
def compute_log_probs(model, input_ids, attention_mask, labels):
"""
Compute per-token log probabilities for a sequence.
Returns the mean log probability over non-padding tokens.
"""
with torch.no_grad():
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Shift: logits[..., :-1, :] predict labels[..., 1:]
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
# Compute log probabilities
log_probs = F.log_softmax(shift_logits, dim=-1)
# Gather log probs for actual tokens
token_log_probs = log_probs.gather(2, shift_labels.unsqueeze(-1)).squeeze(-1)
# Mask out padding (-100 in labels)
mask = (shift_labels != -100).float()
masked_log_probs = token_log_probs * mask
# Mean log prob per sequence
seq_log_probs = masked_log_probs.sum(dim=-1) / mask.sum(dim=-1).clamp(min=1)
return seq_log_probsWhat does this code do?
- torch.no_grad() disables gradient tracking because this function only measures probabilities. It never updates the model.
- The shift aligns logits with labels. In causal language modeling, the logit at position i predicts the token at position i+1.
- F.log_softmax converts raw logits into log-probabilities, then .gather() picks out the log-prob for each actual token in the sequence.
- The mask ignores padding tokens (marked with -100) so they don't skew the average. The final output is the mean log-probability per sequence in the batch.
- Save utils.py.
Why do both SFT and DPO need log-probabilities?
SFT uses cross-entropy loss, which is closely related to log-probabilities. DPO explicitly compares log-probabilities of chosen vs. rejected responses under two models (the policy and the reference).
By factoring this computation into a shared utility, you avoid duplicating the tricky shift-and-gather logic across multiple scripts.
Write the text generation function
You need a way to actually produce text from your models so you can compare outputs after SFT and DPO training. This function implements autoregressive generation with top-k sampling.
- In utils.py, add the generate_text function below the compute_log_probs function:
def generate_text(model, tokenizer, prompt, max_new_tokens=50, temperature=1.0, top_k=50):
"""Generate text from a prompt using top-k sampling."""
model.eval()
token_ids = tokenizer.encode(prompt)
input_ids = torch.tensor([token_ids], dtype=torch.long)
with torch.no_grad():
for _ in range(max_new_tokens):
# Only use last 1024 tokens (GPT-2 context window)
context = input_ids[:, -1024:]
outputs = model(input_ids=context)
logits = outputs.logits[:, -1, :] / temperature
# Top-k filtering
if top_k > 0:
top_k_logits, top_k_indices = torch.topk(logits, top_k, dim=-1)
logits = torch.full_like(logits, float("-inf"))
logits.scatter_(1, top_k_indices, top_k_logits)
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
input_ids = torch.cat([input_ids, next_token], dim=1)
# Stop at end-of-text token
if next_token.item() == 50256:
break
generated_ids = input_ids[0].tolist()
return tokenizer.decode(generated_ids)What does this code do?
- The function generates text one token at a time (autoregressive). At each step it feeds the current sequence into the model and samples the next token.
- temperature controls randomness. Lower values (like 0.7) make output more focused. Higher values (like 1.5) make it more random.
- Top-k filtering zeroes out all logits except the k highest-probability tokens. This prevents the model from sampling extremely unlikely tokens.
- The context window is clamped to 1024 tokens because that is GPT-2's maximum sequence length.
- Generation stops when the model emits token 50256 (GPT-2's end-of-text token) or reaches max_new_tokens.
- Save utils.py.
Seeing a syntax error when you save?
Check that the generate_text function starts at the same indentation level as compute_log_probs (no extra indentation). Make sure there is a blank line between the two functions.
Help me debug a syntax error in utils.py
Verify both functions import correctly
Before moving on, confirm that Python can import your utilities without errors. This catches any typos or indentation issues now rather than midway through training.
- In your terminal, run this command to test the imports:
python -c "from utils import compute_log_probs, generate_text; print('Both functions imported successfully')"- You should see Both functions imported successfully printed to your terminal.
Seeing an ImportError or ModuleNotFoundError?
- Make sure you are running the command from the same folder where utils.py lives.
- Check that your virtual environment is still activated (you should see the environment name in your terminal prompt).
- Confirm the file is named exactly utils.py (not util.py or Utils.py).
Help me fix an import error with my utils.py file
✔️ Awesome, I've got everything!
Great. Make sure you have saved utils.py.
ⓧ I'd like to double check the full code
Here is the complete utils.py file for comparison:
import torch
import torch.nn.functional as F
def compute_log_probs(model, input_ids, attention_mask, labels):
"""
Compute per-token log probabilities for a sequence.
Returns the mean log probability over non-padding tokens.
"""
with torch.no_grad():
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Shift: logits[..., :-1, :] predict labels[..., 1:]
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
# Compute log probabilities
log_probs = F.log_softmax(shift_logits, dim=-1)
# Gather log probs for actual tokens
token_log_probs = log_probs.gather(2, shift_labels.unsqueeze(-1)).squeeze(-1)
# Mask out padding (-100 in labels)
mask = (shift_labels != -100).float()
masked_log_probs = token_log_probs * mask
# Mean log prob per sequence
seq_log_probs = masked_log_probs.sum(dim=-1) / mask.sum(dim=-1).clamp(min=1)
return seq_log_probs
def generate_text(model, tokenizer, prompt, max_new_tokens=50, temperature=1.0, top_k=50):
"""Generate text from a prompt using top-k sampling."""
model.eval()
token_ids = tokenizer.encode(prompt)
input_ids = torch.tensor([token_ids], dtype=torch.long)
with torch.no_grad():
for _ in range(max_new_tokens):
# Only use last 1024 tokens (GPT-2 context window)
context = input_ids[:, -1024:]
outputs = model(input_ids=context)
logits = outputs.logits[:, -1, :] / temperature
# Top-k filtering
if top_k > 0:
top_k_logits, top_k_indices = torch.topk(logits, top_k, dim=-1)
logits = torch.full_like(logits, float("-inf"))
logits.scatter_(1, top_k_indices, top_k_logits)
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
input_ids = torch.cat([input_ids, next_token], dim=1)
# Stop at end-of-text token
if next_token.item() == 50256:
break
generated_ids = input_ids[0].tolist()
return tokenizer.decode(generated_ids)Your shared utilities are ready. Next up, you will use these helpers inside a full SFT training loop that fine-tunes GPT-2 on your dataset.
Implement the SFT Training Loop
In the previous step, you built the data pipeline and utility functions that handle tokenization and batching. Now it is time to write the actual training loop that fine-tunes GPT-2 on your dataset.
Supervised Fine-Tuning (SFT) teaches a pretrained model to follow your data's patterns by computing how wrong the model's next-token predictions are (cross-entropy loss), then adjusting the model's weights to reduce that error. You will write every piece of this loop by hand so you understand exactly what high-level frameworks like TRL abstract away.
In this step, get ready to:
- Write the forward pass, cross-entropy loss computation, backward pass, and optimizer step with gradient clipping.
- Add learning rate warmup and periodic loss logging.
- Run SFT training for 100 steps and save the fine-tuned checkpoint.
Write the training infrastructure in sft_train.py
The SFT script needs to load the pretrained GPT-2 model, feed it batches from your data loader, and set up the optimizer that will update the model's weights. This infrastructure wraps around the training loop itself.
- Create a new file called sft_train.py in your project folder.
- Add the imports, function signature, and model/data loading setup by pasting this code:
import time
import torch
import torch.nn.functional as F
import tiktoken
from transformers import GPT2LMHeadModel
from data_loader import create_dataloader
def train_sft(
jsonl_path,
output_dir="checkpoints/sft",
batch_size=4,
max_length=256,
learning_rate=5e-5,
max_steps=100,
warmup_steps=10,
grad_clip=1.0,
log_every=10,
):
"""Supervised Fine-Tuning training loop from scratch."""
# Load tokenizer and model
print("Loading GPT-2 model...")
tokenizer = tiktoken.get_encoding("gpt2")
model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2")
model.train()
# Create data loader
print(f"Loading data from {jsonl_path}...")
dataloader = create_dataloader(jsonl_path, tokenizer, batch_size=batch_size, max_length=max_length)
# Set up optimizer (AdamW with weight decay)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)What does this code do?
GPT2LMHeadModel.from_pretrained("openai-community/gpt2") downloads the pretrained GPT-2 124M weights from Hugging Face. Calling model.train() puts the model in training mode so dropout layers are active.
AdamW is the optimizer used in most LLM training runs. It decouples weight decay from the gradient update, which prevents the regularization from interfering with the adaptive learning rate. The weight_decay=0.01 value is standard for language model fine-tuning.
Write the core training loop
The training loop iterates over batches, computes how wrong the model's predictions are, and adjusts the weights to reduce that error. Each iteration follows the standard deep learning cycle: forward pass, loss computation, backward pass, gradient clipping, optimizer step.
- Add the following training loop below the optimizer setup in sft_train.py:
# Training loop
print(f"Starting SFT training for {max_steps} steps...")
step = 0
total_loss = 0.0
start_time = time.time()
while step < max_steps:
for batch in dataloader:
if step >= max_steps:
break
input_ids = batch["input_ids"]
labels = batch["labels"]
attention_mask = batch["attention_mask"]
# Learning rate warmup (linear)
if step < warmup_steps:
lr = learning_rate * (step + 1) / warmup_steps
else:
lr = learning_rate
for param_group in optimizer.param_groups:
param_group["lr"] = lr
# Forward pass
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Compute cross-entropy loss (shift logits and labels)
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
# Flatten for cross_entropy
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100,
)
# Backward pass
optimizer.zero_grad()
loss.backward()
# Gradient clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
# Optimizer step
optimizer.step()
total_loss += loss.item()
step += 1
# Logging
if step % log_every == 0:
avg_loss = total_loss / log_every
elapsed = time.time() - start_time
print(
f"Step {step}/{max_steps} | Loss: {avg_loss:.4f} | "
f"LR: {lr:.2e} | Time: {elapsed:.1f}s"
)
total_loss = 0.0What does this training loop do?
The loop has five key stages that repeat every step:
- Forward pass: the model predicts the next token at every position in the input sequence.
- Loss computation: shift_logits and shift_labels align predictions with their targets (logit at position i predicts the token at position i+1). ignore_index=-100 tells PyTorch to skip padded positions.
- Backward pass: loss.backward() computes gradients for every parameter, telling each weight which direction to adjust.
- Gradient clipping: clip_grad_norm_ caps the total gradient magnitude at 1.0 to prevent unstable updates (exploding gradients).
- Optimizer step: optimizer.step() applies the clipped gradients to update model weights using the AdamW algorithm.
Why shift the logits and labels?
In causal language modeling, the model predicts the next token given all previous tokens. The logit output at position 0 is the model's prediction for what comes at position 1. Shifting aligns each prediction with its ground truth target so the loss function can compare them directly.
Understand the warmup and logging
Two important additions make the training loop production-ready: learning rate warmup and periodic loss logging.
Look at the warmup section you just added:
# Learning rate warmup (linear)
if step < warmup_steps:
lr = learning_rate * (step + 1) / warmup_steps
else:
lr = learning_rate
for param_group in optimizer.param_groups:
param_group["lr"] = lrWhy warm up the learning rate?
At the very start of training, the model's gradients can be noisy and large because the optimizer's momentum estimates have not stabilized yet. Starting with a small learning rate and linearly increasing it over the first 10 steps prevents the optimizer from making destructive updates early on.
After warmup completes, the learning rate stays at its full value (5e-5) for the remaining steps.
The logging section prints the average loss every 10 steps:
# Logging
if step % log_every == 0:
avg_loss = total_loss / log_every
elapsed = time.time() - start_time
print(
f"Step {step}/{max_steps} | Loss: {avg_loss:.4f} | "
f"LR: {lr:.2e} | Time: {elapsed:.1f}s"
)
total_loss = 0.0What does the logging tell you?
You should see the loss decrease over steps, which means the model is learning your dataset's patterns. If loss stays flat or increases, something is wrong with the data or hyperparameters. The elapsed time helps you estimate how long the full run takes on CPU.
Run SFT training and save the checkpoint
The final piece of the script saves the fine-tuned model weights and provides an entry point to run training from the terminal.
- Add the checkpoint saving and main block at the bottom of sft_train.py, below the training loop:
# Save checkpoint
import os
os.makedirs(output_dir, exist_ok=True)
model.save_pretrained(output_dir)
print(f"SFT model saved to {output_dir}")
return model
if __name__ == "__main__":
# Replace with your Project 3 JSONL file path
train_sft("training_data.jsonl", max_steps=100)What does save_pretrained do?
model.save_pretrained() writes the model weights and configuration to the checkpoints/sft/ directory in Hugging Face format. This means you can reload it later with GPT2LMHeadModel.from_pretrained("checkpoints/sft") without retraining.
- Save sft_train.py.
- Run the SFT training by executing this command in your terminal:
python sft_train.pyTraining will take approximately 5-10 minutes on CPU. You should see output showing the loss decreasing every 10 steps, similar to:
Loading GPT-2 model...
Loading data from training_data.jsonl...
Starting SFT training for 100 steps...
Step 10/100 | Loss: 4.2315 | LR: 5.00e-05 | Time: 32.1s
Step 20/100 | Loss: 3.8742 | LR: 5.00e-05 | Time: 64.5s
...Your exact loss values will differ depending on your dataset, but the trend should be downward.
Training is too slow or runs out of memory?
If training takes longer than 15 minutes for 100 steps, reduce max_length to 128 in the train_sft() call. Shorter sequences process faster.
If you see a memory error, reduce batch_size to 2. Smaller batches use less RAM at the cost of noisier gradient estimates.
If you see an error about training_data.jsonl not being found, make sure the file from Project 3 is in the same folder as sft_train.py.
Help me debug my SFT training loop
- Confirm the checkpoint was saved by running:
ls checkpoints/sft/You should see files like config.json and model.safetensors in the output.
✔️ Awesome, I've got everything!
Great. Double check you have saved sft_train.py and that the checkpoints/sft/ directory contains your fine-tuned model files.
ⓧ I'd like to double check the full code
Here is the complete sft_train.py file:
import time
import torch
import torch.nn.functional as F
import tiktoken
from transformers import GPT2LMHeadModel
from data_loader import create_dataloader
def train_sft(
jsonl_path,
output_dir="checkpoints/sft",
batch_size=4,
max_length=256,
learning_rate=5e-5,
max_steps=100,
warmup_steps=10,
grad_clip=1.0,
log_every=10,
):
"""Supervised Fine-Tuning training loop from scratch."""
# Load tokenizer and model
print("Loading GPT-2 model...")
tokenizer = tiktoken.get_encoding("gpt2")
model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2")
model.train()
# Create data loader
print(f"Loading data from {jsonl_path}...")
dataloader = create_dataloader(jsonl_path, tokenizer, batch_size=batch_size, max_length=max_length)
# Set up optimizer (AdamW with weight decay)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
# Training loop
print(f"Starting SFT training for {max_steps} steps...")
step = 0
total_loss = 0.0
start_time = time.time()
while step < max_steps:
for batch in dataloader:
if step >= max_steps:
break
input_ids = batch["input_ids"]
labels = batch["labels"]
attention_mask = batch["attention_mask"]
# Learning rate warmup (linear)
if step < warmup_steps:
lr = learning_rate * (step + 1) / warmup_steps
else:
lr = learning_rate
for param_group in optimizer.param_groups:
param_group["lr"] = lr
# Forward pass
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Compute cross-entropy loss (shift logits and labels)
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
# Flatten for cross_entropy
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100,
)
# Backward pass
optimizer.zero_grad()
loss.backward()
# Gradient clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
# Optimizer step
optimizer.step()
total_loss += loss.item()
step += 1
# Logging
if step % log_every == 0:
avg_loss = total_loss / log_every
elapsed = time.time() - start_time
print(
f"Step {step}/{max_steps} | Loss: {avg_loss:.4f} | "
f"LR: {lr:.2e} | Time: {elapsed:.1f}s"
)
total_loss = 0.0
# Save checkpoint
import os
os.makedirs(output_dir, exist_ok=True)
model.save_pretrained(output_dir)
print(f"SFT model saved to {output_dir}")
return model
if __name__ == "__main__":
# Replace with your Project 3 JSONL file path
train_sft("training_data.jsonl", max_steps=100)Your SFT model is now saved and ready. Next up, you will implement DPO to align the model's outputs with human preferences using the fine-tuned checkpoint you just created.
Implement DPO from Scratch
Your SFT model has learned to mimic the patterns in your training data. But mimicry alone does not produce the kind of responses humans prefer.
Direct Preference Optimization (DPO) is the alignment technique that bridges this gap. It teaches a model to prefer one response over another by comparing log-probabilities under a trainable policy and a frozen reference model. In this step, you will implement the full DPO pipeline from scratch.
In this step, get ready to:
- Create synthetic preference pairs for DPO training.
- Implement the DPO loss function with log-probability computation.
- Run DPO training and save the aligned checkpoint.
Create synthetic preference pairs
DPO needs pairs of responses where one is "chosen" (preferred) and the other is "rejected" (dispreferred). Instead of collecting human annotations, you will generate these pairs synthetically using your SFT model at different temperatures.
A low temperature (0.7) produces focused, coherent text. A high temperature (1.5) produces scattered, random text. This difference creates a natural preference signal.
- Create a new file called dpo_train.py in your project folder.
- Add the imports and the create_preference_pairs function by copying the following code:
import copy
import time
import json
import torch
import torch.nn.functional as F
import tiktoken
from transformers import GPT2LMHeadModel
from utils import generate_text
def create_preference_pairs(model, tokenizer, prompts, num_pairs=20):
"""
Generate synthetic preference pairs using the SFT model.
'Chosen' = low temperature (more focused), 'Rejected' = high temperature (more random).
"""
preference_data = []
model.eval()
for prompt in prompts[:num_pairs]:
# Generate chosen response (low temperature = more coherent)
chosen = generate_text(model, tokenizer, prompt, max_new_tokens=30, temperature=0.7, top_k=40)
chosen_response = chosen[len(prompt):]
# Generate rejected response (high temperature = more random)
rejected = generate_text(model, tokenizer, prompt, max_new_tokens=30, temperature=1.5, top_k=50)
rejected_response = rejected[len(prompt):]
if chosen_response.strip() and rejected_response.strip():
preference_data.append({
"prompt": prompt,
"chosen": chosen_response,
"rejected": rejected_response,
})
model.train()
return preference_dataWhat does this code do?
The function takes the SFT model and generates two responses for each prompt using the generate_text helper you wrote in utils.py.
- temperature=0.7 produces the "chosen" response. Lower temperature sharpens the probability distribution, favoring higher-probability tokens.
- temperature=1.5 produces the "rejected" response. Higher temperature flattens the distribution, allowing low-probability tokens through.
- The function returns a list of dictionaries, each containing the prompt, chosen response, and rejected response.
- Save dpo_train.py.
You will not see output yet. The next chunks add the functions that use these preference pairs during training.
Why synthetic pairs instead of human annotations?
In production systems like ChatGPT, human labelers rank responses to create preference data. For this project, generating pairs at different temperatures creates a similar signal without requiring manual annotation.
The key insight is that DPO only needs a relative preference signal. It does not need perfect labels. Coherent text is reliably preferred over random text, which gives the DPO loss a clear gradient to follow.
Implement the DPO loss function
DPO works by comparing how a policy model and a frozen reference model score the same responses. The loss pushes the policy to assign higher probability to chosen responses relative to the reference.
- In dpo_train.py, add the compute_sequence_log_probs function between the imports and create_preference_pairs by copying the following code:
def compute_sequence_log_probs(model, input_ids, attention_mask, response_start_idx):
"""
Compute log probability of a response given a prompt.
Only computes log probs for tokens after response_start_idx.
"""
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Get log probs for response tokens only
log_probs = F.log_softmax(logits, dim=-1)
# Shift: logit at position i predicts token at position i+1
response_log_probs = []
for i in range(response_start_idx, input_ids.size(1) - 1):
token_id = input_ids[0, i + 1]
if token_id == 50256: # Skip padding
break
token_log_prob = log_probs[0, i, token_id]
response_log_probs.append(token_log_prob)
if len(response_log_probs) == 0:
return torch.tensor(0.0, requires_grad=True)
return torch.stack(response_log_probs).mean()What does this code do?
This function scores how likely a model thinks a specific response is, given a prompt. It only looks at response tokens (after response_start_idx) so the prompt tokens do not skew the score.
- F.log_softmax converts raw logits into log-probabilities across the vocabulary at each position.
- The loop gathers the log-probability of each actual response token (the token at position i+1 predicted by logit at position i).
- The mean of these log-probabilities gives a single "how likely is this response" score that DPO uses for comparison.
- Save dpo_train.py.
You will not see a change yet. The next chunk adds the DPO loss that uses these log-probabilities.
- Now add the dpo_loss function below create_preference_pairs by copying this code:
def dpo_loss(
policy_chosen_logps,
policy_rejected_logps,
reference_chosen_logps,
reference_rejected_logps,
beta=0.1,
):
"""
Compute the DPO loss.
L_DPO = -log(sigmoid(beta * (log(pi/pi_ref)(chosen) - log(pi/pi_ref)(rejected))))
"""
# Log ratios
chosen_log_ratios = policy_chosen_logps - reference_chosen_logps
rejected_log_ratios = policy_rejected_logps - reference_rejected_logps
# DPO logits
logits = beta * (chosen_log_ratios - rejected_log_ratios)
# DPO loss (negative log sigmoid)
loss = -F.logsigmoid(logits).mean()
# Implicit rewards (for monitoring)
chosen_rewards = beta * chosen_log_ratios.detach()
rejected_rewards = beta * rejected_log_ratios.detach()
return loss, chosen_rewards, rejected_rewardsWhat does the DPO loss formula do?
The DPO loss has three components working together to push the model toward preferred responses.
- Log ratios measure how much the policy diverges from the reference for each response. A positive ratio means the policy assigns higher probability than the reference.
- beta (set to 0.1) controls how much the policy can diverge from the reference. A smaller beta means tighter constraint.
- The -F.logsigmoid(logits) term creates a loss that decreases when the policy's preference for "chosen" over "rejected" grows relative to the reference.
- Save dpo_train.py.
You will not see output yet. The training loop in the next substep connects all these pieces together.
Seeing a syntax error?
- Check that compute_sequence_log_probs is placed between the imports and create_preference_pairs. Order matters for readability but not for Python execution.
- Make sure dpo_loss is placed after create_preference_pairs and before any train_dpo function you add next.
Help me fix a syntax error in my dpo_train.py file
Run DPO training and save the aligned checkpoint
The training loop ties everything together. It loads the SFT checkpoint as a trainable policy, creates a frozen copy as the reference, generates preference pairs, and optimizes the DPO loss.
- Add the train_dpo function and the __main__ block at the bottom of dpo_train.py by copying the following code:
def train_dpo(
sft_model_dir="checkpoints/sft",
output_dir="checkpoints/dpo",
jsonl_path="training_data.jsonl",
beta=0.1,
learning_rate=1e-5,
max_steps=50,
log_every=5,
):
"""DPO training loop from scratch."""
# Load tokenizer
tokenizer = tiktoken.get_encoding("gpt2")
# Load the SFT model as the policy (trainable)
print("Loading SFT model as policy...")
policy_model = GPT2LMHeadModel.from_pretrained(sft_model_dir)
policy_model.train()
# Create a frozen copy as the reference model
print("Creating frozen reference model...")
reference_model = copy.deepcopy(policy_model)
reference_model.eval()
for param in reference_model.parameters():
param.requires_grad = False
# Extract prompts from JSONL for preference pair generation
prompts = []
with open(jsonl_path, "r", encoding="utf-8") as f:
for line in f:
data = json.loads(line.strip())
if "instruction" in data:
prompts.append(f"### Instruction:\n{data['instruction']}\n\n### Response:\n")
elif "prompt" in data:
prompts.append(data["prompt"])
# Generate synthetic preference pairs
print("Generating synthetic preference pairs...")
preference_data = create_preference_pairs(policy_model, tokenizer, prompts, num_pairs=30)
print(f"Generated {len(preference_data)} preference pairs")
if len(preference_data) == 0:
print("No preference pairs generated. Check your dataset.")
return policy_model
# Set up optimizer
optimizer = torch.optim.AdamW(policy_model.parameters(), lr=learning_rate, weight_decay=0.01)
# DPO training loop
print(f"Starting DPO training for {max_steps} steps (beta={beta})...")
step = 0
total_loss = 0.0
start_time = time.time()
while step < max_steps:
for pair in preference_data:
if step >= max_steps:
break
prompt = pair["prompt"]
chosen_text = prompt + pair["chosen"]
rejected_text = prompt + pair["rejected"]
# Tokenize
prompt_tokens = tokenizer.encode(prompt)
chosen_tokens = tokenizer.encode(chosen_text)[:256]
rejected_tokens = tokenizer.encode(rejected_text)[:256]
chosen_ids = torch.tensor([chosen_tokens], dtype=torch.long)
rejected_ids = torch.tensor([rejected_tokens], dtype=torch.long)
chosen_mask = torch.ones_like(chosen_ids)
rejected_mask = torch.ones_like(rejected_ids)
response_start = len(prompt_tokens) - 1 # -1 for shift alignment
# Compute log probs under policy model
policy_chosen_logps = compute_sequence_log_probs(
policy_model, chosen_ids, chosen_mask, response_start
)
policy_rejected_logps = compute_sequence_log_probs(
policy_model, rejected_ids, rejected_mask, response_start
)
# Compute log probs under reference model (no grad)
with torch.no_grad():
ref_chosen_logps = compute_sequence_log_probs(
reference_model, chosen_ids, chosen_mask, response_start
)
ref_rejected_logps = compute_sequence_log_probs(
reference_model, rejected_ids, rejected_mask, response_start
)
# Compute DPO loss
loss, chosen_rewards, rejected_rewards = dpo_loss(
policy_chosen_logps.unsqueeze(0),
policy_rejected_logps.unsqueeze(0),
ref_chosen_logps.unsqueeze(0),
ref_rejected_logps.unsqueeze(0),
beta=beta,
)
# Backward pass
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(policy_model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
step += 1
# Logging
if step % log_every == 0:
avg_loss = total_loss / log_every
reward_margin = (chosen_rewards - rejected_rewards).mean().item()
elapsed = time.time() - start_time
print(
f"Step {step}/{max_steps} | Loss: {avg_loss:.4f} | "
f"Reward Margin: {reward_margin:.4f} | Time: {elapsed:.1f}s"
)
total_loss = 0.0
# Save checkpoint
import os
os.makedirs(output_dir, exist_ok=True)
policy_model.save_pretrained(output_dir)
print(f"DPO model saved to {output_dir}")
return policy_model
if __name__ == "__main__":
train_dpo()What does the training loop do?
- copy.deepcopy(policy_model) creates the frozen reference. This is a full copy of the SFT weights that never updates. DPO measures divergence against this baseline.
- For each preference pair, the loop computes log-probabilities of both the chosen and rejected responses under both models (4 forward passes per step).
- The reward_margin metric tracks how well the model separates preferred from dispreferred responses. A positive and growing margin means DPO is working.
- Gradient clipping at 1.0 prevents unstable updates from the log-probability computations.
- Save dpo_train.py.
- Run DPO training by executing this command in your terminal:
python dpo_train.pyThis will take a few minutes on CPU. You should see log output every 5 steps showing the loss decreasing and the reward margin.
What should I see?
The script first generates preference pairs (this takes a minute as it runs inference 60 times). Then it prints training logs like:
Generated 25 preference pairs
Starting DPO training for 50 steps (beta=0.1)...
Step 5/50 | Loss: 0.6931 | Reward Margin: 0.0012 | Time: 15.3s
Step 10/50 | Loss: 0.6890 | Reward Margin: 0.0045 | Time: 30.1s
...Your exact numbers will differ. The loss should start near 0.693 (log(2), the loss when the model has no preference) and decrease slightly. A growing reward margin confirms the model is learning to prefer chosen responses.
Seeing 0 preference pairs generated?
- Check that your training_data.jsonl file is in the same directory where you run the script. The script reads prompts from this file.
- Verify your JSONL entries have either an instruction field or a prompt field. Entries without either are skipped.
- Make sure the checkpoints/sft folder exists from Step 3. The DPO script loads the SFT model from that directory.
Help me debug why DPO training shows 0 preference pairs generated
✔️ Awesome, I've got everything!
Great. Double check you have saved dpo_train.py and that a checkpoints/dpo folder now exists.
ⓧ I'd like to double check the full code
import copy
import time
import json
import torch
import torch.nn.functional as F
import tiktoken
from transformers import GPT2LMHeadModel
from utils import generate_text
def compute_sequence_log_probs(model, input_ids, attention_mask, response_start_idx):
"""
Compute log probability of a response given a prompt.
Only computes log probs for tokens after response_start_idx.
"""
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
# Get log probs for response tokens only
log_probs = F.log_softmax(logits, dim=-1)
# Shift: logit at position i predicts token at position i+1
response_log_probs = []
for i in range(response_start_idx, input_ids.size(1) - 1):
token_id = input_ids[0, i + 1]
if token_id == 50256: # Skip padding
break
token_log_prob = log_probs[0, i, token_id]
response_log_probs.append(token_log_prob)
if len(response_log_probs) == 0:
return torch.tensor(0.0, requires_grad=True)
return torch.stack(response_log_probs).mean()
def create_preference_pairs(model, tokenizer, prompts, num_pairs=20):
"""
Generate synthetic preference pairs using the SFT model.
'Chosen' = low temperature (more focused), 'Rejected' = high temperature (more random).
"""
preference_data = []
model.eval()
for prompt in prompts[:num_pairs]:
# Generate chosen response (low temperature = more coherent)
chosen = generate_text(model, tokenizer, prompt, max_new_tokens=30, temperature=0.7, top_k=40)
chosen_response = chosen[len(prompt):]
# Generate rejected response (high temperature = more random)
rejected = generate_text(model, tokenizer, prompt, max_new_tokens=30, temperature=1.5, top_k=50)
rejected_response = rejected[len(prompt):]
if chosen_response.strip() and rejected_response.strip():
preference_data.append({
"prompt": prompt,
"chosen": chosen_response,
"rejected": rejected_response,
})
model.train()
return preference_data
def dpo_loss(
policy_chosen_logps,
policy_rejected_logps,
reference_chosen_logps,
reference_rejected_logps,
beta=0.1,
):
"""
Compute the DPO loss.
L_DPO = -log(sigmoid(beta * (log(pi/pi_ref)(chosen) - log(pi/pi_ref)(rejected))))
"""
# Log ratios
chosen_log_ratios = policy_chosen_logps - reference_chosen_logps
rejected_log_ratios = policy_rejected_logps - reference_rejected_logps
# DPO logits
logits = beta * (chosen_log_ratios - rejected_log_ratios)
# DPO loss (negative log sigmoid)
loss = -F.logsigmoid(logits).mean()
# Implicit rewards (for monitoring)
chosen_rewards = beta * chosen_log_ratios.detach()
rejected_rewards = beta * rejected_log_ratios.detach()
return loss, chosen_rewards, rejected_rewards
def train_dpo(
sft_model_dir="checkpoints/sft",
output_dir="checkpoints/dpo",
jsonl_path="training_data.jsonl",
beta=0.1,
learning_rate=1e-5,
max_steps=50,
log_every=5,
):
"""DPO training loop from scratch."""
# Load tokenizer
tokenizer = tiktoken.get_encoding("gpt2")
# Load the SFT model as the policy (trainable)
print("Loading SFT model as policy...")
policy_model = GPT2LMHeadModel.from_pretrained(sft_model_dir)
policy_model.train()
# Create a frozen copy as the reference model
print("Creating frozen reference model...")
reference_model = copy.deepcopy(policy_model)
reference_model.eval()
for param in reference_model.parameters():
param.requires_grad = False
# Extract prompts from JSONL for preference pair generation
prompts = []
with open(jsonl_path, "r", encoding="utf-8") as f:
for line in f:
data = json.loads(line.strip())
if "instruction" in data:
prompts.append(f"### Instruction:\n{data['instruction']}\n\n### Response:\n")
elif "prompt" in data:
prompts.append(data["prompt"])
# Generate synthetic preference pairs
print("Generating synthetic preference pairs...")
preference_data = create_preference_pairs(policy_model, tokenizer, prompts, num_pairs=30)
print(f"Generated {len(preference_data)} preference pairs")
if len(preference_data) == 0:
print("No preference pairs generated. Check your dataset.")
return policy_model
# Set up optimizer
optimizer = torch.optim.AdamW(policy_model.parameters(), lr=learning_rate, weight_decay=0.01)
# DPO training loop
print(f"Starting DPO training for {max_steps} steps (beta={beta})...")
step = 0
total_loss = 0.0
start_time = time.time()
while step < max_steps:
for pair in preference_data:
if step >= max_steps:
break
prompt = pair["prompt"]
chosen_text = prompt + pair["chosen"]
rejected_text = prompt + pair["rejected"]
# Tokenize
prompt_tokens = tokenizer.encode(prompt)
chosen_tokens = tokenizer.encode(chosen_text)[:256]
rejected_tokens = tokenizer.encode(rejected_text)[:256]
chosen_ids = torch.tensor([chosen_tokens], dtype=torch.long)
rejected_ids = torch.tensor([rejected_tokens], dtype=torch.long)
chosen_mask = torch.ones_like(chosen_ids)
rejected_mask = torch.ones_like(rejected_ids)
response_start = len(prompt_tokens) - 1 # -1 for shift alignment
# Compute log probs under policy model
policy_chosen_logps = compute_sequence_log_probs(
policy_model, chosen_ids, chosen_mask, response_start
)
policy_rejected_logps = compute_sequence_log_probs(
policy_model, rejected_ids, rejected_mask, response_start
)
# Compute log probs under reference model (no grad)
with torch.no_grad():
ref_chosen_logps = compute_sequence_log_probs(
reference_model, chosen_ids, chosen_mask, response_start
)
ref_rejected_logps = compute_sequence_log_probs(
reference_model, rejected_ids, rejected_mask, response_start
)
# Compute DPO loss
loss, chosen_rewards, rejected_rewards = dpo_loss(
policy_chosen_logps.unsqueeze(0),
policy_rejected_logps.unsqueeze(0),
ref_chosen_logps.unsqueeze(0),
ref_rejected_logps.unsqueeze(0),
beta=beta,
)
# Backward pass
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(policy_model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
step += 1
# Logging
if step % log_every == 0:
avg_loss = total_loss / log_every
reward_margin = (chosen_rewards - rejected_rewards).mean().item()
elapsed = time.time() - start_time
print(
f"Step {step}/{max_steps} | Loss: {avg_loss:.4f} | "
f"Reward Margin: {reward_margin:.4f} | Time: {elapsed:.1f}s"
)
total_loss = 0.0
# Save checkpoint
import os
os.makedirs(output_dir, exist_ok=True)
policy_model.save_pretrained(output_dir)
print(f"DPO model saved to {output_dir}")
return policy_model
if __name__ == "__main__":
train_dpo()You now have three model checkpoints: base GPT-2, your SFT-tuned model, and your DPO-aligned model. Next up, you will compare all three side by side to see how each training stage changed the model's behavior.
Evaluate and Compare Models
You now have two fine-tuned checkpoints sitting in your checkpoints/ folder. One was trained with SFT, the other aligned with DPO. The question is: did any of that training actually change how the model generates text?
In this step, you will write an evaluation script that loads all three model variants (base GPT-2, SFT, and DPO), measures perplexity on your training data, and prints a side-by-side generation comparison for the same prompts.
In this step, get ready to:
- Write a perplexity computation function that measures how surprised each model is by your dataset.
- Load base GPT-2, SFT, and DPO checkpoints and display perplexity scores.
- Generate and compare text completions from all three models on the same test prompts.
Write the perplexity computation function
Perplexity measures how well a language model predicts a sequence of tokens. A lower perplexity means the model assigns higher probability to the actual text. If SFT worked, the fine-tuned model should have lower perplexity on your dataset than base GPT-2.
- Create a new file called evaluate.py in your project folder.
- Add the imports and the compute_perplexity function by pasting this code:
import math
import torch
import torch.nn.functional as F
import tiktoken
from transformers import GPT2LMHeadModel
from utils import generate_text
from data_loader import create_dataloader
def compute_perplexity(model, dataloader, max_batches=10):
"""Compute perplexity on a dataset."""
model.eval()
total_loss = 0.0
total_tokens = 0
with torch.no_grad():
for i, batch in enumerate(dataloader):
if i >= max_batches:
break
input_ids = batch["input_ids"]
labels = batch["labels"]
attention_mask = batch["attention_mask"]
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100,
reduction="sum",
)
num_tokens = (shift_labels != -100).sum().item()
total_loss += loss.item()
total_tokens += num_tokens
avg_loss = total_loss / max(total_tokens, 1)
perplexity = math.exp(avg_loss)
return perplexityWhat does this code do?
The compute_perplexity function runs batches of tokenized text through a model and accumulates the total cross-entropy loss across all tokens.
- It uses reduction="sum" so it can correctly average loss over all non-padding tokens across batches.
- The final perplexity is exp(average_loss). Lower means the model assigns higher probability to the actual tokens in your data.
- The max_batches=10 limit keeps evaluation fast on CPU by only scoring a portion of the dataset.
- Save evaluate.py.
Load all three models and compute perplexity
Now you need a function that loads the base, SFT, and DPO checkpoints and runs the perplexity computation on each. This gives you a quantitative comparison before looking at generated text.
- In evaluate.py, add the evaluate_models function below compute_perplexity by pasting this code:
def evaluate_models(jsonl_path="training_data.jsonl"):
"""Compare base, SFT, and DPO models."""
tokenizer = tiktoken.get_encoding("gpt2")
# Load models
print("Loading models...")
base_model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2")
sft_model = GPT2LMHeadModel.from_pretrained("checkpoints/sft")
dpo_model = GPT2LMHeadModel.from_pretrained("checkpoints/dpo")
# Compute perplexity
print("\n--- Perplexity Evaluation ---")
dataloader = create_dataloader(jsonl_path, tokenizer, batch_size=4, max_length=256, shuffle=False)
base_ppl = compute_perplexity(base_model, dataloader)
sft_ppl = compute_perplexity(sft_model, dataloader)
dpo_ppl = compute_perplexity(dpo_model, dataloader)
print(f"{'Model':<15} {'Perplexity':<15}")
print(f"{'-'*30}")
print(f"{'Base GPT-2':<15} {base_ppl:<15.2f}")
print(f"{'SFT':<15} {sft_ppl:<15.2f}")
print(f"{'DPO':<15} {dpo_ppl:<15.2f}")What does this code do?
- It loads all three model variants using from_pretrained. The base model comes from Hugging Face, while the SFT and DPO models load from your local checkpoint directories.
- The dataloader is created with shuffle=False so all three models are evaluated on the same data in the same order.
- The formatted print statements create a neat table showing perplexity for each model stage.
What should I expect from the perplexity scores?
Base GPT-2 will have the highest perplexity because it has never seen your domain-specific data. SFT should be noticeably lower because it trained directly on that text. DPO may be slightly higher or lower than SFT depending on how the preference optimization shifted the distribution.
- Save evaluate.py.
Add the generation comparison and run the evaluation
Perplexity tells you how well each model fits your data numerically. But the real test is reading what the models actually generate. This section feeds the same prompts to all three models and prints the responses side by side.
- In evaluate.py, add the generation comparison section at the end of the evaluate_models function (still indented inside it), plus the if __name__ block at the bottom:
# Generation comparison
print("\n--- Generation Comparison ---")
test_prompts = [
"### Instruction:\nExplain what a pure function is.\n\n### Response:\n",
"### Instruction:\nWhat is formal verification?\n\n### Response:\n",
"### Instruction:\nDescribe property-based testing.\n\n### Response:\n",
]
for prompt in test_prompts:
print(f"\nPrompt: {prompt.strip()[:60]}...")
print(f"{'-'*60}")
base_output = generate_text(base_model, tokenizer, prompt, max_new_tokens=40, temperature=0.8)
sft_output = generate_text(sft_model, tokenizer, prompt, max_new_tokens=40, temperature=0.8)
dpo_output = generate_text(dpo_model, tokenizer, prompt, max_new_tokens=40, temperature=0.8)
base_response = base_output[len(prompt):]
sft_response = sft_output[len(prompt):]
dpo_response = dpo_output[len(prompt):]
print(f" Base: {base_response[:100]}")
print(f" SFT: {sft_response[:100]}")
print(f" DPO: {dpo_response[:100]}")
print("\n--- Evaluation Complete ---")
if __name__ == "__main__":
evaluate_models()What does this code do?
- Three test prompts are formatted in the same instruction/response template your training data used. This lets you see how each model handles the pattern it was (or wasn't) trained on.
- The generate_text function from utils.py uses top-k sampling at temperature 0.8 to produce varied but coherent completions.
- Each response is sliced to remove the prompt prefix, then truncated to 100 characters for clean display.
- Save evaluate.py.
✔️ Awesome, I've got everything!
Great. Double check you have saved your file before running it.
ⓧ I'd like to double check the full code
Here is the complete evaluate.py file:
import math
import torch
import torch.nn.functional as F
import tiktoken
from transformers import GPT2LMHeadModel
from utils import generate_text
from data_loader import create_dataloader
def compute_perplexity(model, dataloader, max_batches=10):
"""Compute perplexity on a dataset."""
model.eval()
total_loss = 0.0
total_tokens = 0
with torch.no_grad():
for i, batch in enumerate(dataloader):
if i >= max_batches:
break
input_ids = batch["input_ids"]
labels = batch["labels"]
attention_mask = batch["attention_mask"]
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100,
reduction="sum",
)
num_tokens = (shift_labels != -100).sum().item()
total_loss += loss.item()
total_tokens += num_tokens
avg_loss = total_loss / max(total_tokens, 1)
perplexity = math.exp(avg_loss)
return perplexity
def evaluate_models(jsonl_path="training_data.jsonl"):
"""Compare base, SFT, and DPO models."""
tokenizer = tiktoken.get_encoding("gpt2")
# Load models
print("Loading models...")
base_model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2")
sft_model = GPT2LMHeadModel.from_pretrained("checkpoints/sft")
dpo_model = GPT2LMHeadModel.from_pretrained("checkpoints/dpo")
# Compute perplexity
print("\n--- Perplexity Evaluation ---")
dataloader = create_dataloader(jsonl_path, tokenizer, batch_size=4, max_length=256, shuffle=False)
base_ppl = compute_perplexity(base_model, dataloader)
sft_ppl = compute_perplexity(sft_model, dataloader)
dpo_ppl = compute_perplexity(dpo_model, dataloader)
print(f"{'Model':<15} {'Perplexity':<15}")
print(f"{'-'*30}")
print(f"{'Base GPT-2':<15} {base_ppl:<15.2f}")
print(f"{'SFT':<15} {sft_ppl:<15.2f}")
print(f"{'DPO':<15} {dpo_ppl:<15.2f}")
# Generation comparison
print("\n--- Generation Comparison ---")
test_prompts = [
"### Instruction:\nExplain what a pure function is.\n\n### Response:\n",
"### Instruction:\nWhat is formal verification?\n\n### Response:\n",
"### Instruction:\nDescribe property-based testing.\n\n### Response:\n",
]
for prompt in test_prompts:
print(f"\nPrompt: {prompt.strip()[:60]}...")
print(f"{'-'*60}")
base_output = generate_text(base_model, tokenizer, prompt, max_new_tokens=40, temperature=0.8)
sft_output = generate_text(sft_model, tokenizer, prompt, max_new_tokens=40, temperature=0.8)
dpo_output = generate_text(dpo_model, tokenizer, prompt, max_new_tokens=40, temperature=0.8)
base_response = base_output[len(prompt):]
sft_response = sft_output[len(prompt):]
dpo_response = dpo_output[len(prompt):]
print(f" Base: {base_response[:100]}")
print(f" SFT: {sft_response[:100]}")
print(f" DPO: {dpo_response[:100]}")
print("\n--- Evaluation Complete ---")
if __name__ == "__main__":
evaluate_models()- Run the evaluation script by running this command in your terminal:
python evaluate.pyWhat should I see?
The script will first print a perplexity table with three rows (Base GPT-2, SFT, DPO). You should see the SFT perplexity significantly lower than the base model. DPO perplexity may be close to SFT or slightly different.
Below the table, you will see three test prompts with each model's generated response. The base model will likely produce generic or off-topic text, while SFT and DPO should generate responses more aligned with your training domain. This takes a few minutes on CPU since it loads three copies of GPT-2.
Seeing a FileNotFoundError?
- If the error mentions checkpoints/sft or checkpoints/dpo, make sure you ran both sft_train.py and dpo_train.py in the previous steps and that they completed without errors.
- If the error mentions training_data.jsonl, confirm that your JSONL file from Project 3 is in the same directory where you are running the script.
- If you see an out-of-memory error, close other applications to free up RAM. Loading three copies of GPT-2 requires roughly 1.5 GB of memory.
Still stuck? Why is my evaluate.py throwing an error when loading checkpoints?
Secret mission
Cosine LR Decay and Reward Margin Tracking
Your training pipeline works, but production pipelines use smarter learning rate schedules and track alignment metrics over time. In this extension, you will replace the flat post-warmup learning rate with cosine decay and add reward margin visualization to your DPO loop.
Clean Up Your Resources
Clean Up Your Resources
Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.
Resources you used:
- Python scripts (data_loader.py, utils.py, sft_train.py, dpo_train.py, evaluate.py).
- Model checkpoints (checkpoints/sft/ and checkpoints/dpo/).
- Generated files (reward_margin.png, training_data.jsonl).
- Python virtual environment.
Keep everything running
No action needed. Choose this if you plan to keep experimenting with your fine-tuned models or want to revisit the training loops later.
- Your project folder, checkpoints, and virtual environment stay on your Desktop.
- The cached GPT-2 model weights (~500MB) remain in your Hugging Face cache directory.
Pause - I'll come back to this later
Shut down any active terminal sessions and deactivate your virtual environment. Your files stay intact so you can pick up where you left off.
- Deactivate your virtual environment by running:
deactivate- Close VS Code.
- When you return, reactivate the virtual environment and you are ready to go.
Delete - I don't want to use this again
Remove all project files, checkpoints, and the virtual environment from your machine.
- Close VS Code.
- Deactivate your virtual environment if it is still active by running:
deactivate- Navigate to your Desktop and delete the project folder by running:
cd ~/Desktop
rm -rf [[PROJECT_FOLDER="your-project-folder-name"]]On Windows, use this command instead:
cd ~\Desktop
rmdir /s /q [[PROJECT_FOLDER="your-project-folder-name"]]- This removes all Python scripts, model checkpoints, reward_margin.png, training_data.jsonl, and the virtual environment in one step.
Nice Work!
Nice Work!
Nice work! You have built an end-to-end LLM fine-tuning pipeline from scratch, implementing the same post-training stages used in production language model alignment.
You've learned how to:
- Built a custom data loader that converts JSONL into tokenized training batches for causal language modeling, with padding, attention masks, and shifted input/target pairs.
- Implemented a complete SFT training loop with AdamW optimizer, gradient clipping, learning rate warmup, and cross-entropy loss on next-token predictions.
- Implemented DPO from scratch, including log-probability computation, frozen reference model management, synthetic preference pair generation, and the DPO loss function.
- Added cosine decay scheduling to the SFT loop and reward margin visualization to DPO training, plotting how the model progressively separates preferred from dispreferred responses.
Ready to quiz yourself?