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

> Learn how to implement Direct Preference Optimization DPO in train-llm-from-scratch. Understand DPO theory and see code examples for aligning language models with human preferences easily.

- Repository: [Fareed Khan/train-llm-from-scratch](https://github.com/FareedKhan-dev/train-llm-from-scratch)
- Tags: tutorial
- Published: 2026-06-11

---

**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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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.

```python
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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_dpo.py) demonstrates the core optimization step:

```python

# 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:

```python
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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/dpo.py), which processes log-probabilities from `sequence_logprobs`
- Training is orchestrated by [`scripts/train_dpo.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/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.