Direct Preference Optimization (DPO) Implementation Guide: From Theory to Code in train-llm-from-scratch

Direct Preference Optimization (DPO) is a reinforcement learning objective that aligns language models to human preferences by optimizing the log-probability ratio between a policy and a frozen reference model, eliminating the need for an explicit reward model.

Direct Preference Optimization (DPO) provides a streamlined alternative to traditional RLHF pipelines by deriving implicit rewards directly from the policy's probability distribution. This article examines the complete implementation of DPO in the FareedKhan-dev/train-llm-from-scratch repository, covering the mathematical foundations, core utility functions, and the end-to-end training orchestration.

What is Direct Preference Optimization?

DPO reformulates preference alignment as a binary classification problem over paired responses. Instead of training a separate reward model to estimate human preferences, DPO treats the log-probability ratio between the learned policy and a static reference model (typically the SFT checkpoint) as an implicit reward:

$$r = \beta \bigl(\log p_{\text{policy}} - \log p_{\text{ref}}\bigr)$$

The training objective maximizes the likelihood that the policy assigns higher probability to chosen responses over rejected ones. The loss function is defined as:

$$\mathcal{L}{\text{DPO}} = -\frac{1}{B}\sum{i}\log\sigma!\bigl(\beta(\Delta\pi_i - \Delta r_i)\bigr)$$

Where $\Delta\pi_i$ represents the log-probability difference between chosen and rejected pairs for the policy, $\Delta r_i$ is the same difference for the reference model, and $\beta$ acts as a temperature parameter controlling the divergence from the reference.

Core Implementation Components

The dpo_loss Function

Located in src/post_training/dpo.py, the dpo_loss function implements the core optimization objective. It accepts summed log-probabilities for chosen and rejected sequences from both the policy and reference models, then computes the binary cross-entropy loss along with detached implicit reward values for monitoring.

from src.post_training.dpo import dpo_loss
import torch

# Log-probabilities from forward passes (batch size = 2)

policy_chosen_logps = torch.tensor([-5.2, -6.1])
policy_rejected_logps = torch.tensor([-7.4, -8.0])
ref_chosen_logps = torch.tensor([-5.5, -6.3])
ref_rejected_logps = torch.tensor([-7.8, -8.2])

loss, chosen_reward, rejected_reward = dpo_loss(
    policy_chosen_logps,
    policy_rejected_logps,
    ref_chosen_logps,
    ref_rejected_logps,
    beta=0.1,
)

Sequence Log-Probability Computation

The sequence_logprobs function in src/post_training/rollout.py calculates the total log-probability of a response by summing token-level log-probabilities across the sequence. This utility is essential for both the active policy and the frozen reference model, ensuring consistent comparison dimensions.

Preference Data Pipeline

The preference_dataset module in data_loader/preference_dataset.py provides get_preference_iterator, which yields batches containing chosen_ids, rejected_ids, and their corresponding attention masks. This iterator handles distributed sampling, padding, and device placement for DDP training.

Training Orchestration

The scripts/train_dpo.py script manages the full training lifecycle:

• Loads the SFT checkpoint as both the initialization point and the frozen reference
• Creates a non-trainable copy of the reference model
• Sets up Distributed Data Parallel (DDP) training
• Implements the optimization loop with gradient clipping
• Logs implicit_accuracy metrics every 20 steps

Configuration and Monitoring

Hyperparameters are controlled via the DPOConfig dataclass in config/post_training_config.py, which specifies beta, batch size, learning rate schedules, and checkpoint paths. The implicit_accuracy utility (lines 93-96 in src/post_training/dpo.py) calculates the fraction of preference pairs where the chosen response receives a higher implicit reward than the rejected one, providing a lightweight alignment metric.

The DPO Training Pipeline

The implementation follows a strict six-stage workflow:

1. Data preparation – Preference pairs are generated and stored as JSONL files via scripts/prepare_preference_data.py
2. Batch loading – get_preference_iterator creates distributed batches of tokenized sequences
3. Log-probability extraction – sequence_logprobs computes summed log-probs for both policy and reference models
4. Loss calculation – dpo_loss, orpo_loss, or kto_loss is selected based on DPOConfig.loss_type
5. Optimization – Gradients flow only through the policy network while the reference remains frozen
6. Evaluation – implicit_accuracy monitors preference alignment without requiring separate reward model inference

Practical Code Examples

Integrating the Training Loop

This snippet from scripts/train_dpo.py demonstrates the core optimization step:

# Inside the main training loop

loss, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx)
optimizer.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(policy.parameters(), cfg.grad_clip)
optimizer.step()

# Log implicit-reward accuracy on main process

if ctx.is_main and step % 20 == 0:
    acc = implicit_accuracy(cr, rr).item()
    print(f"step {step} | loss {loss.item():.4f} | acc {acc:.3f}")

Evaluating Alignment Quality

To validate model performance on held-out preference data:

from src.post_training.dpo import implicit_accuracy
from data_loader.preference_dataset import get_preference_iterator

def eval_implicit_acc(policy, ref, cfg, ctx, max_batches=100):
    policy.eval()
    it = get_preference_iterator(
        "/ephemeral/data/preferences_test.jsonl",
        cfg.batch_size,
        cfg.max_len,
        device=ctx.device,
        rank=ctx.rank,
        world_size=ctx.world_size,
        shuffle=False,
        infinite=False,
    )
    acc, n = 0.0, 0
    for batch in it:
        _, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx)
        acc += implicit_accuracy(cr, rr).item()
        n += 1
        if n >= max_batches:
            break
    return acc / max(1, n)

Alternative Loss Variants

The repository also implements reference-free variants in src/post_training/dpo.py:

• orpo_loss – Odds Ratio Preference Optimization that eliminates the need for a reference model entirely
• kto_loss – Kahneman-Tversky Optimization that uses a KL-regularized objective

Summary

• Direct Preference Optimization eliminates the reward model training phase by using log-probability ratios between policy and reference as implicit rewards
• The implementation in train-llm-from-scratch centers on the dpo_loss function in src/post_training/dpo.py, which processes log-probabilities from sequence_logprobs
• Training is orchestrated by scripts/train_dpo.py, which handles DDP setup, frozen reference model maintenance, and implicit_accuracy logging
• The codebase supports DPO variants including ORPO and KTO through the loss_type parameter in DPOConfig
• All preference data flows through get_preference_iterator in data_loader/preference_dataset.py, ensuring proper masking and distributed sampling

Frequently Asked Questions

What is the difference between DPO and traditional RLHF?

Traditional RLHF requires training a separate reward model to score completions, then optimizing the policy against that reward using PPO or similar algorithms. DPO removes the reward model entirely by deriving the optimization objective directly from the policy's probability distribution relative to a frozen reference model. This reduces memory overhead and training instability while maintaining alignment quality.

How does the beta parameter affect DPO training?

The beta parameter in DPOConfig acts as a temperature scaling factor controlling the divergence penalty between the policy and reference model. Higher beta values (e.g., 0.5) enforce stronger adherence to the reference distribution, preventing the policy from deviating too far from the SFT checkpoint. Lower values (e.g., 0.1) allow more aggressive optimization toward preference data but may increase the risk of overfitting or mode collapse.

Can DPO be used alongside other preference optimization methods?

Yes. The train-llm-from-scratch repository implements DPO, ORPO, and KTO in the same src/post_training/dpo.py module, allowing users to switch between objectives via the loss_type configuration field. While DPO requires a reference model, ORPO operates reference-free, and KTO introduces KL-regularization. These variants share the same data pipeline and evaluation metrics, making it straightforward to compare approaches.

Why does the training script freeze the reference model?

The reference model in DPO serves as a static baseline representing the SFT-trained policy. By freezing its parameters (as implemented in scripts/train_dpo.py), the implementation ensures that the implicit reward calculation remains stable throughout training. This stability is crucial because the DPO loss depends on the difference between policy and reference log-probabilities; if both models updated simultaneously, the target would shift, destabilizing convergence.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}