Part 3 — Align a Language Model with DPO
A pretrained language model can predict text, and supervised fine-tuning teaches it to follow instructions. Alignment addresses a different question: when multiple responses are possible, which one should the model prefer?
SmallLM-Forge explores this stage with Direct Preference Optimization (DPO). The implementation focuses on preference optimization rather than a complete RLHF pipeline with a separately trained reward model and reinforcement learning.
The alignment path is:
SFT model
↓
preference pairs
↓
DPO objective
↓
aligned model
From Supervised Fine-Tuning to Preference Alignment
Supervised fine-tuning teaches one desired response per prompt. Many real tasks, however, do not have a single perfect answer. Preference learning instead compares two candidate responses and learns which one is better:
SFT optimizes one correct response per instruction, while preference learning compares chosen and rejected responses.
The model learns to increase the probability of preferred outputs relative to less preferred alternatives.
From Preferences to DPO
Traditional RLHF usually introduces an additional reward-model stage. DPO skips that stage and directly optimizes the policy from preference comparisons:
RLHF learns a reward model from human preferences before policy optimization; DPO optimizes the policy directly from chosen/rejected comparisons.
The key question becomes:
\[ \text{Does the current model increase the probability of chosen responses more than rejected responses?} \]From Data to the DPO Loss
The DPO objective is built from the same causal language-modeling machinery used in Part 1. Starting from the raw preference triple, this part prepares the inputs, defines a sequence score from token log-probabilities, and derives the loss that compares the policy with a frozen reference model.
Prepare Preference Data
The alignment dataset contains:
{
"prompt": "...",
"chosen": "...",
"rejected": "..."
}
Each of the three fields is tokenized separately, then padded so the batch is dense:
prompt → prompt_input_ids (left-padded)
chosen → chosen_input_ids (right-padded)
rejected → rejected_input_ids (right-padded)
These tensors provide the inputs needed by the DPO training step.
Compute Preference Scores
DPO compares sequence likelihoods. For a response \(y\) given a prompt \(x\), the model assigns the sequence score
\[ \log\pi_\theta(y|x) = \sum_t \log p_\theta(y_t \mid x, y_{\lt t}), \]the sum of the per-token log-probabilities produced by the model. The
implementation extracts these scores with get_log_prob().
The DPO Objective
DPO does not train a separate reward model. Instead, it directly optimizes the policy model from preference comparisons.
For each prompt \(x\), the dataset provides two possible responses:
- a preferred response \(y_c\) (chosen),
- a less preferred response \(y_r\) (rejected).
The objective is not simply to maximize the likelihood of the chosen response. The policy should increase the probability of the preferred response relative to a reference model.
For the chosen response:
\[ \Delta_c = \log\pi_\theta(y_c \mid x) - \log\pi_{ref}(y_c \mid x) \]For the rejected response:
\[ \Delta_r = \log\pi_\theta(y_r \mid x) - \log\pi_{ref}(y_r \mid x) \]The DPO objective encourages:
\[ \Delta_c > \Delta_r \]through:
\[ L_{DPO} = -\log\sigma \left( \beta(\Delta_c-\Delta_r) \right) \]The implementation also tracks:
- Reward accuracy: how often the chosen response receives a higher preference score than the rejected response.
- Reward margin: the average gap between chosen and rejected preference scores.
The Training Loop
With the objective defined, training becomes a compact loop: evaluate the chosen and rejected responses with both the policy and a frozen reference model, extract their scores, and step the optimizer through the DPO loss. The same loop composes with the LoRA and DoRA adapters from the previous part.
Implement DPO Training
For each preference pair, the policy model and the reference model evaluate the same responses:
\[ (x,y_c),\qquad(x,y_r) \]The policy model is trainable, while the reference model remains frozen.
DPO training: prompt and responses run through the trainable policy and the frozen reference model; the DPO loss updates only the policy side.
The implementation first obtains logits:
chosen_logits = model(chosen).logits
rejected_logits = model(rejected).logits
with torch.no_grad():
ref_chosen_logits = ref_model(chosen).logits
ref_rejected_logits = ref_model(rejected).logits
The sequence log probabilities are extracted:
chosen_logps
rejected_logps
ref_chosen_logps
ref_rejected_logps
Then they are passed into:
loss, acc, margin = dpo_loss(
chosen_logps,
rejected_logps,
ref_chosen_logps,
ref_rejected_logps,
)
Finally, the loss updates the policy model.
Aligning with Parameter-Efficient Updates
DPO can be combined with the LoRA and DoRA adapters introduced in the previous part.
The alignment pipeline combines instruction tuning, parameter-efficient adapters, and preference optimization.
The three stages optimize different aspects of model behavior:
| Stage | Objective |
|---|---|
| Pretraining | Learn general language patterns |
| SFT | Learn how to answer instructions |
| DPO | Learn which answers are preferred |
When LoRA or DoRA is used, DPO updates only the adapter parameters while the pretrained base model remains frozen.
Example: Preference Alignment
A preference example contains one prompt and two candidate responses.
A preference example: one prompt branches into a chosen and a rejected response.
The workflow is:
- Format preference pairs.
- Tokenize prompt, chosen, and rejected responses.
- Compute policy and reference likelihoods.
- Optimize the DPO objective.
- Update the alignment parameters.
The example demonstrates preference optimization. The same pipeline applies whenever feedback can be represented as comparisons between candidate outputs.
