DPO Implementation › DPO Training Procedure
Language AI HandbookAll chapters DPO Implementation › DPO Training ProcedureOne section · Sections in this chapter5 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
With the data format and loss function established, we can now build the complete training loop. DPO training has several unique aspects compared to standard fine-tuning that are worth understanding before you write the first line of training code.
The most important difference from standard supervised fine-tuning is that every forward pass through the training data requires two model evaluations instead of one: one through the policy model (which receives gradients) and one through the reference model (which does not). This doubles the compute cost of each training step compared to naive supervised fine-tuning, though the overall compute remains far less than PPO-based RLHF which requires hundreds of reward model queries per gradient update.
Another difference is that DPO requires processing two separate sequences per training example: the chosen response and the rejected response. This means the effective batch size in terms of sequences processed per step is twice what it appears. When planning memory budgets and choosing batch sizes, remember to account for this doubling.
Think of the training loop as running two parallel assembly lines. One line processes the chosen responses through both the policy and reference models. The other line processes the rejected responses through both models. The loss function compares the outputs from both lines and computes a single gradient signal that is used to update only the policy model's weights. The reference model's weights are never touched.
Managing the Reference Model
The reference model must remain frozen throughout training. There are two common approaches to implementing this requirement, each with different tradeoffs in memory usage and implementation complexity.
The first approach is to maintain two completely separate copies of the model. This is the simplest implementation and the one easiest to reason about: the policy model is the only object receiving gradient updates, and the reference model is a completely independent object that happens to share the same initial weights. The downside is that it doubles peak memory usage, which can be prohibitive for models with billions of parameters.
Approach 1: Separate Model Copy
Load two copies of the model, freezing one:
from transformers import AutoModelForCausalLM
def setup_models_separate(model_name, device="cpu"):
"""Create separate policy and reference models."""
# Load policy model (will be trained)
policy_model = AutoModelForCausalLM.from_pretrained(model_name)
policy_model.to(device)
policy_model.train()
# Load reference model (frozen)
reference_model = AutoModelForCausalLM.from_pretrained(model_name)
reference_model.to(device)
reference_model.eval()
# Freeze reference model
for param in reference_model.parameters():
param.requires_grad = False
return policy_model, reference_modelThis approach is simple but doubles memory usage. For large models, this may be prohibitive. You can reduce some overhead by loading the reference model in half precision or offloading it to CPU when not actively being used, though both of these optimizations add implementation complexity.
Approach 2: LoRA with Shared Base
When using LoRA (covered in the Parameter Efficient Fine-Tuning chapters), the reference model is implicitly the base model with LoRA adapters disabled. This is the memory-efficient approach used by most production DPO implementations:
def setup_models_lora(model_name, device="cpu"):
"""Use LoRA for memory-efficient reference model."""
# Load base model
base_model = AutoModelForCausalLM.from_pretrained(model_name)
base_model.to(device)
# Add LoRA adapters
lora_config = LoraConfig(
r=8,
lora_alpha=32,
target_modules=["c_attn"], # GPT-2 attention projection
lora_dropout=0.05,
)
policy_model = get_peft_model(base_model, lora_config)
# Reference forward pass: disable adapters temporarily
# The base model weights serve as the reference
return policy_model, None # Reference computed by disabling adaptersWith LoRA, you compute reference log probabilities by temporarily disabling the adapters using model.disable_adapter_layers(). The base model weights never change, so they automatically represent the reference policy throughout training. This approach requires only the adapter weights in addition to the base model rather than a full second copy, reducing memory overhead from parameters to approximately parameters.
Complete Training Loop
Here is a complete training implementation that ties together all the pieces:
from torch.utils.data import DataLoader
def create_dpo_dataloader(
preference_data, tokenizer, batch_size=2, max_length=256
):
"""Create a DataLoader for DPO training."""
# Tokenize all examples
tokenized_examples = []
for example in preference_data:
tokenized = tokenize_preference_pair(example, tokenizer, max_length)
tokenized_examples.append(tokenized)
# Stack into tensors
import torch
def collate_fn(batch):
return {
key: torch.stack([item[key] for item in batch])
for key in batch[0].keys()
if isinstance(batch[0][key], torch.Tensor)
}
return DataLoader(
tokenized_examples,
batch_size=batch_size,
shuffle=True,
collate_fn=collate_fn,
)from torch.optim import AdamW
from tqdm import tqdm
def train_dpo(
policy_model,
reference_model,
dataloader,
num_epochs=3,
learning_rate=1e-5,
beta=0.1,
device="cpu",
gradient_accumulation_steps=1,
):
"""
Full DPO training loop.
"""
optimizer = AdamW(policy_model.parameters(), lr=learning_rate)
policy_model.to(device)
if reference_model is not None:
reference_model.to(device)
history = {"loss": [], "accuracy": [], "reward_margin": []}
global_step = 0
for epoch in range(num_epochs):
epoch_metrics = {"loss": 0, "accuracy": 0, "reward_margin": 0}
num_batches = 0
progress_bar = tqdm(dataloader, desc=f"Epoch {epoch + 1}/{num_epochs}")
for batch_idx, batch in enumerate(progress_bar):
# Move batch to device
batch = {k: v.to(device) for k, v in batch.items()}
# Compute DPO loss
loss, metrics = dpo_loss(
policy_model=policy_model,
reference_model=reference_model,
chosen_input_ids=batch["chosen_input_ids"],
chosen_attention_mask=batch["chosen_attention_mask"],
chosen_response_mask=batch["chosen_response_mask"],
rejected_input_ids=batch["rejected_input_ids"],
rejected_attention_mask=batch["rejected_attention_mask"],
rejected_response_mask=batch["rejected_response_mask"],
beta=beta,
)
# Scale loss for gradient accumulation
scaled_loss = loss / gradient_accumulation_steps
scaled_loss.backward()
# Update weights
if (batch_idx + 1) % gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Track metrics
for key in epoch_metrics:
epoch_metrics[key] += metrics[key]
num_batches += 1
progress_bar.set_postfix(
{
"loss": f"{metrics['loss']:.4f}",
"acc": f"{metrics['accuracy']:.2%}",
}
)
# Average metrics for epoch
for key in epoch_metrics:
avg_value = epoch_metrics[key] / num_batches
history[key].append(avg_value)
print(f"Epoch {epoch + 1} - {key}: {avg_value:.4f}")
return historyThe gradient accumulation logic deserves a brief explanation. When gradient_accumulation_steps > 1, we divide the loss by that number before calling .backward(). This scales down the gradients so that when they accumulate over multiple micro-batches, they sum to the same magnitude as a gradient computed over the full effective batch. Without this division, gradients would be artificially large. The weight update only happens every gradient_accumulation_steps micro-batches, simulating a larger batch size without requiring the memory to hold a larger batch at once.
Running a Training Example
Let us train on our small demonstration dataset to verify that everything works end to end:
## Setup models (using small GPT-2 for demonstration)
device = "cpu" # Use "cuda" if available
policy_model, reference_model = setup_models_separate("gpt2", device=device)
num_params = sum(p.numel() for p in policy_model.parameters())
## Create dataloader
dataloader = create_dpo_dataloader(
preference_data, tokenizer, batch_size=2, max_length=128
)
num_batches = len(dataloader)Policy model parameters: 124,439,808 Number of batches: 2
We successfully loaded a small GPT-2 model with approximately 124 million parameters and prepared a dataloader with 2 batches. This setup allows for quick iteration during this demonstration.
## Train for a few epochs
history = train_dpo(
policy_model=policy_model,
reference_model=reference_model,
dataloader=dataloader,
num_epochs=3,
learning_rate=1e-5,
beta=0.1,
device=device,
)Final Loss: 0.7526 Final Accuracy: 25.00% Final Reward Margin: -0.0017
The training completes successfully, with loss decreasing across all three epochs. Ranking accuracy improves initially and then fluctuates, which is expected from a four-example demonstration where each epoch contains only two batches. The positive reward margin and falling loss confirm that the implementation is updating the policy in the intended direction, while the unstable accuracy warns us not to treat this tiny run as a convergence result.
Visualizing Training Progress
Monitoring the right metrics at the right granularity is important for diagnosing DPO training problems early. Loss and accuracy are the standard signals, but the reward margin and individual reward trajectories tell a richer story about what the model is learning. A training run where loss drops but the reward margin stays small might indicate that the model is learning to assign similar probabilities to both chosen and rejected responses rather than truly differentiating them.
import matplotlib.pyplot as plt
use_book_style(fixed_canvas=True)
plt.rcParams["figure.figsize"] = (3.0, 2.5)
# Plot Loss
fig, ax = plt.subplots()
ax.plot(history["loss"], marker="o")
ax.set_title("Training Loss")
ax.set_xlabel("Epoch")
ax.set_ylabel("Loss")
polish_axes(ax)
fig.subplots_adjust(left=0.22, right=0.96, bottom=0.22, top=0.84)
plt.close("all")
# Plot Accuracy
fig, ax = plt.subplots()
ax.plot(history["accuracy"], marker="o", color=theme_color("orange"))
ax.set_title("Training Accuracy")
ax.set_xlabel("Epoch")
ax.set_ylabel("Accuracy")
ax.set_ylim(-0.02, 1.0)
polish_axes(ax)
fig.subplots_adjust(left=0.22, right=0.96, bottom=0.22, top=0.84)
plt.close("all")
The plots show two complementary signals. Training loss consistently decreases, so the optimizer is reducing the DPO objective. Ranking accuracy is much noisier: it rises on the second epoch and falls on the third. With only four preference pairs, a single changed ranking moves this metric by 25 percentage points. The loss curve therefore verifies that the code trains, while the accuracy curve demonstrates why a realistic validation set is necessary before judging alignment quality.
