Post-training — SFT, RLHF, DPO

A pre-trained language model computes P(next_token | preceding_tokens) — it will happily continue any prefix with statistically plausible text, regardless of whether the continuation is helpful, toxic, factually wrong, or dangerous. GPT-3's base model, given the prompt "How do I make a bomb?", would produce an answer indistinguishable in tone from a chemistry tutorial, because its training corpus (Common Crawl, WebText2, Books, Wikipedia) contains exactly that kind of text. Post-training is the set of techniques that transform this raw next-token predictor into an aligned, instruction-following assistant.

Supervised Fine-Tuning (SFT)

SFT trains the pre-trained model on curated (instruction, response) pairs using the same next-token prediction objective — cross_entropy(model(prompt + response[:t]), response[t]) — but on data that demonstrates ideal assistant behavior. The model's weights shift so that the distribution over next tokens, conditioned on an instruction prefix, favors helpful completions.

The training data format:

python
sft_examples = [
    {
        "instruction": "Explain quantum entanglement in two sentences.",
        "response": "Quantum entanglement is a phenomenon where two particles become correlated such that measuring one instantaneously determines the state of the other, regardless of distance. This correlation persists even when the particles are separated by arbitrary distances, violating classical locality assumptions."
    },
    {
        "instruction": "Write a Python function that checks if a number is prime.",
        "response": "def is_prime(n):\n    if n < 2:\n        return False\n    for i in range(2, int(n**0.5) + 1):\n        if n % i == 0:\n            return False\n    return True"
    },
]

Dataset sizes for effective SFT are surprisingly small. InstructGPT (Ouyang et al. 2022) used ~13,000 demonstration examples. Alpaca (Taori et al. 2023) used 52,000 instruction-response pairs generated by GPT-4 to fine-tune Llama 7B. LIMA (Zhou et al. 2023) demonstrated that just 1,000 carefully curated examples could produce a competitive instruction-following model from Llama 65B — suggesting that data quality dominates data quantity for SFT.

The training loop is standard supervised learning:

python
import torch
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B")

def format_example(example):
    return f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n{example['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n{example['response']}<|eot_id|>"

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)

for batch in dataloader:
    input_ids = batch["input_ids"].to("cuda")
    labels = input_ids.clone()
    # Mask the instruction tokens — only compute loss on the response
    labels[:, :batch["response_start"]] = -100
    
    outputs = model(input_ids=input_ids, labels=labels)
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

The critical detail: the loss is computed only on the response tokens. The instruction prefix is masked out (set to -100, which PyTorch's CrossEntropyLoss ignores). This ensures the model learns to generate good responses given instructions, not to generate instructions themselves.

SFT costs are modest. Fine-tuning Llama 3 8B on 50K examples for 3 epochs takes ~4 hours on 8×A100 GPUs. Fine-tuning Llama 3 70B on the same data takes ~24 hours on 8×A100.

Reinforcement Learning from Human Feedback (RLHF)

SFT teaches the model what a good response looks like by imitation. RLHF teaches it what humans prefer when there are multiple plausible responses. The procedure, formalized by Ouyang et al. (2022) for InstructGPT/ChatGPT, has three stages.

Stage 1: Preference data collection

Human annotators see a prompt and two model-generated responses (A and B). They mark which response is better — or tie. A single annotation session might label 30-50 comparisons per hour. InstructGPT collected ~33,000 comparison labels. Anthropic's HH-RLHF dataset (Bai et al. 2022) contains ~170,000 preference pairs.

Stage 2: Reward model training

A reward model r(prompt, response) → scalar is trained to predict human preferences. It's initialized from the SFT model (same architecture, minus the language modeling head, plus a scalar value head) and trained with a pairwise ranking loss:

loss=log(sigmoid(r(x,ychosen)r(x,yrejected)))\text{loss} = - \log(\text{sigmoid}(r(x, y_{\text{chosen}}) - r(x, y_{\text{rejected}})))

This is the Bradley-Terry model of preferences — the probability that response A is preferred over response B is sigmoid(r(A) - r(B)). The reward model learns to assign higher scores to responses humans prefer.

python
import torch
import torch.nn as nn

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.backbone = base_model
        self.value_head = nn.Linear(base_model.config.hidden_size, 1)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden = outputs.last_hidden_state[:, -1, :]
        reward = self.value_head(last_hidden).squeeze(-1)
        return reward

def reward_loss(reward_chosen, reward_rejected):
    return -torch.log(torch.sigmoid(reward_chosen - reward_rejected)).mean()

Stage 3: PPO optimization

The SFT model (now called the "policy") is optimized to maximize the reward model's score, subject to a KL-divergence penalty that prevents it from diverging too far from the SFT model:

objective=E[r(x,y)]βKL(policysft_model)\text{objective} = E[r(x, y)] - \beta \cdot \text{KL}(\text{policy} \| \text{sft}\_\text{model}\|)

The KL penalty is critical. Without it, the policy collapses to producing degenerate outputs that exploit the reward model (reward hacking) — e.g., producing long, repetitive text that happens to score high but is clearly unhelpful. The beta coefficient is typically 0.01-0.2.

PPO (Schulman et al. 2017) is the standard algorithm for this optimization. It collects a batch of (prompt, response) pairs from the current policy, scores them with the reward model, computes advantages, and updates the policy with a clipped surrogate objective. A typical RLHF training run on a 7B model processes 512-1024 prompts per batch, generating 256-512 tokens per response, for 200-500 PPO steps. This requires 4 models in GPU memory simultaneously: the policy (being trained), the reference model (frozen SFT model for KL), the reward model, and the value model (critic). For a 7B model, this means ~28B parameters across 4 models, requiring 8×A100-80GB GPUs.

Direct Preference Optimization (DPO)

Rafailov et al. (2023) showed that the RLHF objective can be reformulated as a supervised loss directly on preference pairs, eliminating the reward model and PPO entirely. The key insight: the optimal policy under the KL-constrained reward maximization objective has a closed-form relationship with the reward function:

r(x,y)=βlog(policy(yx)reference(yx))+constantr(x, y) = \beta \cdot \log(\frac{\text{policy}(y|x|)}{\text{reference}(y|x|)}) + \text{constant}

Substituting this into the Bradley-Terry preference model gives the DPO loss:

python
import torch
import torch.nn.functional as F

def dpo_loss(policy_chosen_logprobs, policy_rejected_logprobs,
             reference_chosen_logprobs, reference_rejected_logprobs,
             beta=0.1):
    """
    All inputs are log-probabilities of the full response sequence.
    policy_chosen_logprobs: sum of log P_policy(token | prefix) for chosen response
    reference_chosen_logprobs: same under the frozen reference model
    """
    chosen_rewards = beta * (policy_chosen_logprobs - reference_chosen_logprobs)
    rejected_rewards = beta * (policy_rejected_logprobs - reference_rejected_logprobs)
    
    loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
    return loss

DPO requires only two models in memory (the policy being trained and the frozen reference), uses a standard supervised training loop (no RL, no reward model, no value model), and converges in fewer steps. Training Llama 3 8B with DPO on 60K preference pairs takes ~2 hours on 8×A100, versus ~8-12 hours for the equivalent RLHF run.

DPO has become the dominant alignment method for open-weight models. Zephyr (Tunstall et al. 2023), Neural Chat, and most community fine-tunes use DPO or its variants (IPO, KTO, ORPO).

The alignment tax

RLHF and DPO improve helpfulness and safety at a cost to raw benchmark performance. Ouyang et al. (2022) documented this: InstructGPT scored lower than the base GPT-3 on certain NLP benchmarks while dramatically outperforming it on human preference evaluations. The aligned model becomes more cautious — it may refuse to answer questions that the base model would happily attempt (even when the questions are benign). This tension between safety and capability is the "alignment tax."

Constitutional AI (CAI)

Bai et al. (2022) at Anthropic proposed using the model itself to generate preference data. The process: generate a response, ask the model to critique it against a set of principles (the "constitution"), then ask it to revise. The original and revised responses form preference pairs for DPO/RLHF training. This removes the need for expensive human preference annotation while scaling to arbitrary volumes of training data. Claude's training uses CAI extensively.

GRPO — Group Relative Policy Optimization

DeepSeek-V2 (DeepSeek 2024) introduced GRPO: for each prompt, sample K responses (typically 4-16), score them with a reward function (which can be a rule-based verifier, not necessarily a learned reward model), rank them, and use the relative ranking as the policy gradient signal. Responses above the group mean get positive advantage; those below get negative advantage. No separate reward model is needed, and no value model (critic) is needed — the group mean serves as the baseline.

GRPO is particularly effective for math and code where you have verifiable reward signals (the answer is either correct or not, the code either passes tests or not). DeepSeek-R1 used GRPO with code execution feedback to train its reasoning capabilities.

← Previous