Implementing DPO
The DPO loss fits on one line. Getting it right takes more care than the line suggests, because almost every way of getting it wrong still trains. The loss still falls, the reward accuracy still climbs, and the model that comes out is quietly worse, or degenerate, or no different from the one you started with. A reference model left in training mode, a mean where there should be a sum, an off-by-one in a mask, a reference that is secretly the policy itself. None of these crash. They show up only as numbers in the training log that look slightly off, if you know what the right numbers are.
This section builds a DPO training step from scratch and explains every piece. We compute masked sequence log-probabilities and check them by hand, handle the frozen reference model and what it costs, choose and a learning rate from published recipes, train a toy model end to end while logging what real trainers log, and then break the run on purpose, one bug at a time, so you can recognize each failure from its symptoms.
The anatomy of one step
Recall the loss from §17.2. For a pair with chosen response and rejected response , the policy , the frozen reference , and the strength ,
Everything the loss needs is four numbers per pair, two sequence log-probabilities under each model. A training step computes them in six stages.
- Build the batch. Take triples. Render each prompt with the chat template (§14.3), tokenize prompt plus chosen and prompt plus rejected, and record which tokens belong to the response. Stack all sequences into one padded batch, chosen first and rejected second, so one forward pass serves both.
- Policy forward. Run the policy on the sequences to get logits of shape .
- Sequence log-probabilities. Turn the logits into for each sequence, summing log-probabilities over response tokens only. Split into chosen and rejected halves, each of shape .
- Reference forward. Do the same with the reference model, without gradients and with dropout off, or read precomputed values.
- Loss and metrics. Compute , the implicit rewards, and the margins.
- Backward and update. Gradients flow only through the policy's log-probabilities.
The rest of the section takes these in turn.
Sequence log-probabilities with masks
A causal language model's logits at position predict the token at position (§6.3). Write a tokenized sequence as , its logits as , and a response mask that is 1 when token belongs to the response (including its end-of-sequence token) and 0 for prompt tokens and padding. Then
In plain words, for every response token, look up the log-probability the model gave it at the previous position, and add them up. The two things that go wrong are the shift (the log-probability of token comes from logits ) and the mask (prompt tokens and padding must not count).
A worked example small enough to check in your head. A sequence of six positions holds a two-token prompt, a two-token response with its end token, and one pad.
| position | 0 | 1 | 2 | 3 | 4 | 5 |
|---|---|---|---|---|---|---|
| token | <bos> | Q | A1 | A2 | <eos> | <pad> |
| mask | 0 | 0 | 1 | 1 | 1 | 0 |
| probability of from logits | 0.3 | 0.5 | 0.25 | 0.8 | 0.9 |
The response log-probability uses the entries at , which are . Forgetting the shift reads probabilities from the wrong rows. Applying the mask to input positions rather than target positions (a common slip when the labels are shifted inside the model, as Hugging Face models do) counts the predictions made at positions 2, 3, 4, which are the targets at 3, 4, 5. That drops A1, the first and often most informative response token, and adds the prediction of <pad>.
Masked sequence log-probabilities from logits
Implement sequence_logps(logits, input_ids, response_mask, average=False). logits has shape (N, T, V) where logits[n, t] predicts token t + 1, input_ids has shape (N, T), and response_mask has shape (N, T) with 1 for response tokens of input_ids and 0 for prompt and padding. Return the (N,) array of , or, with average=True, that sum divided by the number of response tokens (assume at least one). It must be stable for logits as large as , and padding may hold any token id.
Precision
The sum runs over hundreds of tokens, and DPO cares about differences of these sums of a few nats. That makes precision matter more than in SFT. Each term computed in 16-bit floating point carries a relative rounding error, and accumulating the sum in 16 bits adds a much larger error, because every addition rounds a running total that is growing in magnitude. The cell sums 400 per-token log-probabilities three ways. (NumPy has float16 but not bfloat16, which keeps 8 bits of precision instead of float16's 11, so bfloat16 is worse than what is shown.)
Rounding the terms costs about 0.005 nats per response. Accumulating in half precision costs 0.84 nats, which is larger than the whole implicit margin of many pairs early in training. The rule is to compute the log-softmax and the sum in float32 even when the model runs in bfloat16, which is what careful trainers do.
The prompt mask, and a surprise
Here is a question that separates people who have implemented DPO from people who have read about it. What happens to the DPO loss if you forget to mask the prompt?
Write the sequence log-probability of prompt plus response as a prompt part and a response part, , where is the sum over prompt tokens. The chosen and rejected sequences share the same prompt, and a causal model's predictions for prompt tokens depend only on earlier prompt tokens. So is identical in the two sequences, and the same holds for the reference. In it appears four times,
and cancels. Forgetting the prompt mask does not change the DPO loss or its gradient at all. So why mask? Because the cancellation is fragile, and several things that depend on the mask do not cancel.
- The logged implicit rewards change. The chosen and rejected rewards each pick up , which drifts as training changes the model's opinion of prompts. The margin is unaffected, but the individual rewards you monitor are wrong.
- Length normalization breaks. Averaging over tokens (SimPO, length-normalized DPO, §17.4) divides by the token count, so the prompt's log-probability no longer cancels, and long prompts dilute the signal.
- Auxiliary terms break. An NLL term on the chosen response (as in Llama 3's recipe, below) would train on prompt tokens, which is the SFT mistake of §14.3.
- Unpaired losses break. KTO (§17.4) scores single responses, with no rejected partner to cancel against.
- Truncation breaks it. If long examples are truncated so that the chosen and rejected sequences keep different amounts of the prompt, the prompt terms differ and no longer cancel.
The log-ratio difference is the same to every printed digit, while the chosen response's log-ratio, and therefore its logged implicit reward, is not.
The reference model
The reference is the model DPO measures every change against, normally the SFT checkpoint the policy starts from. Four rules keep it honest.
Freeze it, and compute it without gradients. The reference's log-probabilities are constants in the loss. In PyTorch that means torch.no_grad() (or inference mode) around its forward pass and requires_grad_(False) on its parameters, so no optimizer ever sees them.
Make it compute exactly what the policy computed at step 0. Same weights, same dtype, same attention implementation, same tokenization and chat template, and dropout off in both models. With any difference, at initialization, the loss does not start at , and the implicit rewards carry a bias from the first step. Dropout is the most common culprit. It adds noise to every token's log-probability, which accumulates over the response, so the reference's numbers change every time they are computed. TRL's DPOConfig disables dropout in both models by default for this reason.
Precompute it if memory is tight. Since the reference is frozen and the dataset is fixed, its two log-probabilities per pair can be computed in one pass before training and stored as two floats. TRL calls this precompute_ref_log_probs. Training then holds only the policy.
With LoRA, the reference is free. If the policy is the base model plus a LoRA adapter (§14.5), the reference is the same network with the adapter switched off, so no second copy of the weights is needed.
What does it cost? For a model with parameters, a forward pass costs about FLOPs per token and a forward plus backward pass about (§6.7). A DPO step on a pair with prompt length and responses of length and processes both sequences through the policy (training) and the reference (forward only),
or with precomputed reference values. SFT on the chosen response alone costs . With equal lengths, a DPO step costs about times an SFT step on the same examples, or twice with a precomputed reference. Memory is the policy's training state (about 16 bytes per parameter for bfloat16 weights with AdamW and float32 master weights, §10.1) plus 2 bytes per parameter for a bfloat16 reference, against PPO's policy, reference, reward model, and value model, two of them in training.
Put real numbers in. Zephyr-7B was trained with DPO on UltraFeedback, about 61,000 pairs. Suppose prompts average 200 tokens and responses 300 (an assumption for the estimate, not a measured statistic). Then each pair is tokens through both models.
An epoch is about FLOPs, roughly two and a half H100-hours at 40% MFU (§6.7), so three epochs take about an hour on one 8-GPU node. Compute is not what makes DPO hard at this scale. Memory is. The policy's training state alone is 112 GB, more than one 80 GB GPU holds, so even a 7B DPO run shards it across GPUs with ZeRO or FSDP (§10.3), and the 14 GB reference rides along, sharded the same way or replaced by precomputed values.
Choosing beta and the learning rate
Section 17.2 explained what does. Smaller values let the policy move further from the reference on each pair before the gradient fades. In practice and the learning rate are chosen together, and DPO learning rates are usually much smaller than SFT's (Zephyr used for SFT and for DPO), although Llama 3 reports for both. Here are published settings.
| recipe | learning rate | other details | |
|---|---|---|---|
| DPO paper (Rafailov et al. 2023) | 0.1, and 0.5 for TL;DR summarization | , RMSprop | batch 64, linear warmup over 150 steps |
| Zephyr-7B (Tunstall et al. 2023) | 0.1 | , linear decay | batch 32, 10% warmup, 1 to 3 epochs (3 for the final model) |
| Llama 3 (Dubey et al. 2024) | 0.1 | formatting tokens masked, NLL term with weight 0.2 on chosen responses | |
| Tülu 3 8B (Lambert et al. 2024) | 5, length-normalized log-probabilities | ( for 70B) | batch 128, 1 epoch, maximum length 2,048 |
| Iterative DPO (Dong et al. 2024) | 0.1 | , cosine | batch 128, 3% warmup, 2 epochs per iteration |
The Tülu 3 row looks out of line until you remember that it averages log-probabilities over tokens. An average is about times smaller than a sum, so a around 5 on averages plays the role of something like on sums, which is in the usual range for responses of a few dozen to a few hundred tokens. Two more practical notes. Most recipes run one to three epochs, because DPO overfits the preference data quickly. Zephyr's authors saw strong overfitting after one epoch that did not hurt their benchmark scores, which is a reminder that the training loss is a poor guide to when to stop. And the KTO paper gives a useful diagnostic that applies to DPO as well. If the learning rate is too high, the implicit rewards of both chosen and rejected responses fall steeply.
A complete training step, from scratch
To see a whole run, we need a model small enough to train in a fraction of a second but real enough to have prompts, variable-length responses, end tokens, and shared parameters. A bigram language model fits. Its parameters are a matrix , and the next-token distribution after token is . A prompt is a start token followed by one of two question tokens. A response is up to six tokens from a small vocabulary, ending in an end token. A hidden "true" reward scores each response by its tokens (one token is worth , one , one , one ), and preference labels are sampled from Bradley-Terry with that reward, which makes them noisy like human labels.
The gradient is the one place a bigram model needs a formula. For one sequence, , the familiar "one-hot minus probabilities" for every response token, placed in the row of the token that preceded it. DPO weights each sequence's gradient by .
The dataset is built the way the DPO paper recommends, by sampling two responses per prompt from the reference model itself and labeling them, so the data lie exactly where the reference puts its mass.
Only 73% of labels agree with the true reward's ranking, and 13% of pairs are exact ties that the labels split at random. Remember that number. It caps the reward accuracy of a model that has learned the true reward (a model that memorized the labels could score higher, which is the wrong goal). Now the training loop, with the metrics a real trainer logs, plus two that only a toy allows us to measure, the KL from the reference on the policy's own samples and the true reward of those samples.
Read the log row by row, because it is what a healthy run looks like.
- Step 0. The loss is exactly and every implicit reward is zero. Always check this first.
- Accuracy. It jumps to about 0.70 and stays there. That is not a failure to learn. The labels agree with the true ranking on only 73% of pairs (counting the ties, which the labels split at random, as any model would), so even a model that had learned the true reward exactly would score about 0.73, and 0.70 is close to that ceiling. Training accuracy near 100% on real preference data would mean memorization.
- Rewards. The rejected reward falls steadily. The chosen reward first rises a little, then turns negative, so the policy ends up assigning the chosen responses lower probability than the reference did (their mean log-probability falls from to ). The margin keeps growing all the same. This is likelihood displacement, and it is normal in DPO. Probability mass has moved to responses the dataset never contains, here short responses full of high-reward words, which beat even the chosen responses. Section 17.5 explains when this is benign, as here, and when it is catastrophic.
- KL and true reward. The true reward of the policy's own samples climbs from about to while the KL grows to about 1.1 nats, which is the reward-KL trade of §17.1 playing out.
The metrics every DPO run logs
Implement dpo_metrics(pc, pr, rc, rr, beta) for arrays of shape (B,) (policy chosen, policy rejected, reference chosen, reference rejected sequence log-probabilities). Return a dict with the batch means under the keys TRL uses, "rewards/chosen", "rewards/rejected", "rewards/margins", "rewards/accuracies", "logps/chosen", and "logps/rejected", all as Python floats. A pair counts as accurate only if its chosen reward is strictly greater than its rejected reward, so ties (including every pair at initialization) count as wrong.
Breaking the run on purpose
Each bug below is a common one in real DPO code. The function above takes a bug argument that injects it. Three of them first, compared with the healthy run after the same number of steps.
Dropout in the reference. Simulated here as unit-variance noise on the reference's sequence log-probabilities, recomputed every step. The loss starts at 0.697 instead of 0.693, the telltale sign. In this toy the final model is barely affected, because the noise averages out over 512 pairs. In a real model the noise is larger (it accumulates over hundreds of tokens), the starting loss can be well above , and the implicit rewards are biased in a way that the gradient weights inherit. The fix is model.eval() for the reference, and dropout off in the policy too.
Mean instead of sum. Averaging the log-probabilities over tokens with the same divides each log-ratio by its response length, about four tokens here, and hundreds in practice. The loss barely moves from 0.693, the margin is about a sixth of the healthy one, and the true reward improves far less. Averaging is not wrong in itself. It is a different objective (length-normalized DPO, §17.4), and it needs a tens to hundreds of times larger, like Tülu 3's . The bug is switching one without the other, which happens easily, because some trainers average for some loss types and sum for others.
The fourth bug is the most dangerous, and needs a longer run to show itself.
The reference is the policy. This happens when the "reference" is the same model object as the policy, evaluated under no_grad, so it changes every step along with the policy. Then always, the logged loss sits at exactly , the logged rewards are zero, and the accuracy is zero. But the gradient is not zero. With the weight is on every pair forever, which is the unweighted objective that §17.2 warned degenerates. The log above measures the rewards against the true reference, which a buggy trainer would not show you. The chosen log-probability falls from to , the rejected one to , and the true reward of the policy's samples rises for a while and then falls. A loss that is exactly flat at 0.6931 is never good news.
A checklist of common bugs
| symptom in the log | likely cause | fix |
|---|---|---|
| loss at step 0 is not 0.6931 | reference differs from the policy (dropout, dtype, attention kernel, template, wrong checkpoint) | eval() both, match configs, load the same weights |
| loss exactly 0.6931 forever, rewards exactly 0 | the reference is the policy | a frozen separate copy, or precomputed reference log-probabilities |
| loss barely moves, tiny margins | averaging log-probs with a tuned for sums, or learning rate far too low | sum, or rescale |
| chosen and rejected rewards both fall fast | learning rate too high | lower it, add warmup |
| accuracy near 100%, margins growing without bound | overfitting near-deterministic pairs | fewer epochs, larger , IPO or label smoothing (§17.4) |
| response length grows every eval | length bias in the data, exploited by DPO (Park et al. 2024) | length-normalized or regularized losses, length-balanced pairs |
| implicit rewards look plausible but the margin is wrong sign on a held-out set | chosen and rejected swapped in the data pipeline | check a few raw examples by eye |
| repetition or abrupt endings after DPO | shared formatting tokens pushed both up and down | mask special formatting tokens, as Llama 3 does |
| noisy or drifting losses with long sequences | log-softmax or sum in 16-bit precision | upcast to float32 before the log-softmax |
The last-but-one row comes from the Llama 3 report, which masks header and termination tokens out of both chosen and rejected responses because, the authors write, these shared tokens give the contrastive loss a conflicting objective, and leaving them in led to tail repetition or abruptly generated end tokens. The same report adds an NLL term with weight 0.2 on the chosen responses to keep their likelihood from falling, the remedy for likelihood displacement that §17.5 examines. Exercise 17.3c implements that combined objective.
A DPO plus NLL training step for a bigram model
Implement dpo_nll_loss_and_grad(W, W_ref, c_ids, c_mask, r_ids, r_mask, beta, alpha) for the bigram model of this section. W and W_ref have shape (V, V), the id arrays have shape (B, T), and the masks mark response tokens of the ids (as in Exercise 17.3a, the log-probability of token comes from row ids[t-1] of the log-softmax of W). The loss is
DPO on summed log-probabilities plus a length-normalized NLL term on the chosen responses, where is the number of response tokens. Return (loss, dW) with dW of shape (V, V). The tests check dW against finite differences.
In PyTorch
Here is the same step for a real nn.Module, written the way production trainers structure it, with chosen and rejected concatenated into one batch, a frozen reference in eval mode under no_grad, the log-softmax upcast to float32, and the logged metrics. The model is a toy (an embedding, a dropout layer, and an output projection, so a bigram model again), but nothing in the step depends on that. It reuses the toy dataset built above.
The first line evaluates the untrained policy with its dropout active. Every pair already has a nonzero log-ratio difference, so the implicit rewards are noise before a single update, and the loss sits above (only slightly here, because the loss is flat near zero, so the excess grows with the square of the noise). With dropout off, the first logged loss is exactly 0.6931, and the run then shows the same pattern as the NumPy version, accuracy near the label-noise ceiling and a growing margin.
Read more
- Rafailov, Sharma, Mitchell, Ermon, Manning & Finn (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. Appendix B has the authors' reference implementation and hyperparameters.
- Hugging Face. DPO Trainer. The TRL documentation, with the logged metrics, reference handling, precomputation, and the supported loss variants.
- Tunstall et al. (2023). Zephyr: Direct Distillation of LM Alignment. A complete open recipe, including the observation that DPO overfits after one epoch without hurting benchmarks.
- Llama Team, AI @ Meta (2024). The Llama 3 Herd of Models. Section 4.1.4 describes DPO at frontier scale, with formatting tokens masked and an NLL term on chosen responses.
- Lambert et al. (2024). Tülu 3: Pushing Frontiers in Open Language Model Post-Training. Section 5 documents an open DPO pipeline in detail, with ablations of loss variants, learning rates, and data.
- Park, Rafailov, Ermon & Finn (2024). Disentangling Length from Quality in Direct Preference Optimization. How DPO exploits length bias in preference data, and a regularizer that controls it.