RLHF take-home aligning a small open model

Question: RLHF take-home aligning a small open model

Companies: Hugging Face

Problem statement

A take-home over a few days: align a small open model with preference data and a parameter-efficient adapter on a single GPU, push the result to the Hugging Face Hub, and write up what happened, including what failed and why.

Choose the model, the data and the method yourself and justify them. The review covers your evaluation, your failure analysis and your reasoning about the loss.

As a warm-up, fill in the core of a DPO step without a framework:

"""The core of the DPO step, without a framework, so it can be unit-tested.
Fill in the three functions; the training script calls them per batch."""


def sequence_logp(token_logps, loss_mask):
    """Sum of log-probabilities over the response tokens of one sequence.
    loss_mask is 1 on response tokens and 0 on prompt and padding tokens."""
    raise NotImplementedError


def dpo_loss_and_grad(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1):
    """DPO loss for one preference pair from sequence log-probabilities.
    Returns (loss, d_loss/d_policy_chosen, d_loss/d_policy_rejected).
    Must stay finite for very large margins."""
    raise NotImplementedError


def dpo_metrics(pairs, beta=0.1):
    """pairs: [(policy_chosen, policy_rejected, ref_chosen, ref_rejected)].
    Returns {"loss", "accuracy", "margin", "chosen_reward", "rejected_reward"}
    averaged over pairs; a reward is beta * (policy - ref) and accuracy is the
    share of pairs whose chosen reward beats the rejected one."""
    raise NotImplementedError

Scope it before you start

The take-home is graded on judgment more than on the final score: did you pick a setup that fits one GPU and a few days, did you measure the right things, and is the write-up honest about what failed. Decide these four things in the first hour and write them at the top of your README.

Why DPO, not PPO. PPO needs a reward model, a value head, rollouts and four models in memory. DPO turns the same objective into a classification loss on pairs you already have. For a take-home, the honest answer is "DPO, because it is cheap and stable; PPO or GRPO if I had a verifiable reward or wanted on-policy data", and then say what DPO gives up: it never sees the model's own outputs, so it can't fix mistakes the data doesn't cover.

The pipeline

{"type":"excalidraw","elements":[
 {"type":"ellipse","id":"data","x":0,"y":-5,"width":260,"height":80,"label":{"text":"Preference pairs\n10k, filtered"}},
 {"type":"rectangle","id":"prep","x":340,"y":0,"width":220,"height":70,"label":{"text":"Chat template\n+ loss mask"}},
 {"type":"rectangle","id":"train","x":640,"y":0,"width":240,"height":70,"label":{"text":"DPO trainer\nbase + LoRA r=16"}},
 {"type":"rectangle","id":"eval","x":960,"y":0,"width":180,"height":70,"label":{"text":"Evals"}},
 {"type":"rectangle","id":"hub","x":1220,"y":0,"width":200,"height":70,"label":{"text":"Hub: adapter\n+ model card"}},
 {"type":"ellipse","id":"ref","x":640,"y":170,"width":240,"height":80,"label":{"text":"Reference\n(adapter off)"}},
 {"type":"arrow","id":"a1","x":260,"y":35,"width":80,"height":0,"start":{"id":"data"},"end":{"id":"prep"}},
 {"type":"arrow","id":"a2","x":560,"y":35,"width":80,"height":0,"start":{"id":"prep"},"end":{"id":"train"}},
 {"type":"arrow","id":"a3","x":880,"y":35,"width":80,"height":0,"start":{"id":"train"},"end":{"id":"eval"}},
 {"type":"arrow","id":"a4","x":1140,"y":35,"width":80,"height":0,"start":{"id":"eval"},"end":{"id":"hub"}},
 {"type":"arrow","id":"a5","x":760,"y":70,"width":1,"height":100,"start":{"id":"train"},"end":{"id":"ref"},"strokeStyle":"dashed"}
]}

The dashed arrow is the trick that halves memory: with LoRA, the reference model is the same weights with the adapter switched off, so you never load a second copy.

1. Filter the data. Drop pairs where chosen and rejected scores tie or differ by less than 1 point (noise), pairs longer than the length limit, and exact duplicates of eval prompts. Log how many you dropped and why.
2. Format. Apply the model's own chat template. Build a loss mask that is 1 only on response tokens. Check one example by decoding it with the mask shown; this catches most bugs before training.
3. Train (settings below), logging loss, reward accuracy, reward margin, chosen and rejected log-probabilities, and mean response length every 10 steps.
4. Evaluate base versus tuned on held-out prompts.
5. Publish the adapter with a model card: data and filtering, hyperparameters, eval results with intervals, and known failures.

The numbers

LoRA size. For a 0.5B model with hidden size 896, 24 layers, an MLP width of 4,864 and grouped-query attention (key and value projections 896 → 128), rank 16 on every linear layer adds per layer: query and output 2 × 16 × (896 + 896) = 57,344; key and value 2 × 16 × (896 + 128) = 32,768; the three MLP matrices 3 × 16 × (896 + 4,864) = 276,480. That is 366,592 per layer, 8.8M in total, under 2% of the model.

Memory. Base weights in bf16: about 1 GB. Adapter, its gradients and Adam states: under 0.2 GB. The rest is activations: with 1,024-token sequences and both responses of a pair in one batch, gradient checkpointing keeps a micro-batch of 4 pairs well under 16 GB.

Compute. 10k pairs × 2 responses × about 600 tokens = 12M tokens per epoch. A LoRA step costs about 4 × params × tokens (forward, plus backward through the activations but not the frozen weight gradients): 4 × 0.5B × 12M = 2.4 × 10^16 FLOPs. The reference pass adds 2 × 0.5B × 12M = 1.2 × 10^16. Total ≈ 3.6 × 10^16. On one A100 at 30% of its 312 TFLOP/s peak, that is about 6 minutes per epoch; on a T4 at about 15 TFLOP/s effective, about 40 minutes. Either way you can afford three or four runs, which you need.

Settings to start from.

The DPO step, tested

The training library computes this for you, but the review will ask you to explain it, and writing it once catches the most common bug (summing log-probs over the prompt too). The starter in the prompt has the three functions to fill in.

"""The core of the DPO step, without a framework, so it can be unit-tested."""

import math


def sequence_logp(token_logps, loss_mask):
    """Sum of log-probabilities over the response tokens of one sequence.
    loss_mask is 1 on response tokens and 0 on prompt and padding tokens."""
    return sum(lp for lp, m in zip(token_logps, loss_mask) if m)


def softplus(x):
    """log(1 + e^x) without overflow."""
    return max(x, 0.0) + math.log1p(math.exp(-abs(x)))


def sigmoid(x):
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    e = math.exp(x)
    return e / (1.0 + e)


def dpo_loss_and_grad(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1):
    """loss = -log sigmoid(beta * ((pc - rc) - (pr - rr))) = softplus(-z).
    Returns (loss, d_loss/d_policy_chosen, d_loss/d_policy_rejected)."""
    z = beta * ((policy_chosen - ref_chosen) - (policy_rejected - ref_rejected))
    weight = sigmoid(-z)  # how wrong the pair still is: 0.5 at z = 0, -> 0 once it is learned
    return softplus(-z), -beta * weight, beta * weight


def dpo_metrics(pairs, beta=0.1):
    n = len(pairs)
    totals = {"loss": 0.0, "accuracy": 0.0, "margin": 0.0, "chosen_reward": 0.0, "rejected_reward": 0.0}
    for pc, pr, rc, rr in pairs:
        chosen, rejected = beta * (pc - rc), beta * (pr - rr)
        totals["loss"] += dpo_loss_and_grad(pc, pr, rc, rr, beta)[0]
        totals["accuracy"] += chosen > rejected
        totals["margin"] += chosen - rejected
        totals["chosen_reward"] += chosen
        totals["rejected_reward"] += rejected
    return {k: v / n for k, v in totals.items()}

What the gradient says: each pair pushes the chosen response up and the rejected one down with weight β · σ(−z). That weight is 0.5β for a pair the model has not learned yet and goes to 0 once the implicit reward margin is large. DPO is a self-weighting classifier: it stops spending effort on pairs it already gets right.

import math

from dpo import dpo_loss_and_grad, dpo_metrics, sequence_logp


def close(a, b, tol=1e-9):
    return abs(a - b) < tol


def test_only_response_tokens_count():
    logps = [-5.0, -4.0, -0.5, -0.25, -9.0]  # prompt, prompt, response, response, padding
    assert close(sequence_logp(logps, [0, 0, 1, 1, 0]), -0.75)


def test_untrained_policy_loss_is_log_two():
    loss, d_chosen, d_rejected = dpo_loss_and_grad(-12.0, -15.0, -12.0, -15.0, beta=0.1)
    assert close(loss, math.log(2))
    assert close(d_chosen, -0.05) and close(d_rejected, 0.05)


def test_loss_falls_as_the_margin_grows():
    losses = [dpo_loss_and_grad(-10.0 + m, -10.0, -10.0, -10.0)[0] for m in (0, 1, 5, 20)]
    assert losses == sorted(losses, reverse=True)


def test_huge_margins_stay_finite():
    good = dpo_loss_and_grad(0.0, -1e4, -100.0, -100.0)
    bad = dpo_loss_and_grad(-1e4, 0.0, -100.0, -100.0)
    assert close(good[0], 0.0, 1e-6) and math.isfinite(bad[0])
    assert close(bad[0], 0.1 * 1e4, 1e-6)  # softplus(z) ~ z for large z


def test_metrics():
    pairs = [(-10.0, -20.0, -12.0, -18.0), (-9.0, -8.0, -10.0, -10.0)]
    m = dpo_metrics(pairs, beta=0.5)
    # pair 1: chosen reward 1.0, rejected -1.0; pair 2: chosen 0.5, rejected 1.0
    assert close(m["accuracy"], 0.5)
    assert close(m["chosen_reward"], 0.75) and close(m["rejected_reward"], 0.0)
    assert close(m["margin"], 0.75)

from dpo import dpo_loss_and_grad


def test_gradient_matches_finite_differences():
    h = 1e-6
    for pc, pr, rc, rr, beta in [(-10, -12, -11, -11, 0.1), (-3, -40, -20, -20, 0.5), (-50, -5, -9, -9, 0.05)]:
        _, d_chosen, d_rejected = dpo_loss_and_grad(pc, pr, rc, rr, beta)
        num_c = (dpo_loss_and_grad(pc + h, pr, rc, rr, beta)[0] - dpo_loss_and_grad(pc - h, pr, rc, rr, beta)[0]) / (2 * h)
        num_r = (dpo_loss_and_grad(pc, pr + h, rc, rr, beta)[0] - dpo_loss_and_grad(pc, pr - h, rc, rr, beta)[0]) / (2 * h)
        assert abs(d_chosen - num_c) < 1e-6 and abs(d_rejected - num_r) < 1e-6

The level 3 test runs DPO end to end on a toy: one prompt, four responses with hidden quality 0–3, 400 noisy labels, and a reference that prefers the worst response. Training on the labels alone recovers the true order.

import math
import random

from dpo import dpo_loss_and_grad


def softmax(xs):
    top = max(xs)
    es = [math.exp(x - top) for x in xs]
    return [e / sum(es) for e in es]


def test_dpo_recovers_the_preference_order():
    # One prompt, four possible responses with hidden quality 3 > 2 > 1 > 0.
    # Labelers prefer the better response with Bradley-Terry probability.
    rng = random.Random(0)
    quality = [0.0, 1.0, 2.0, 3.0]
    pairs = []
    for _ in range(400):
        a, b = rng.sample(range(4), 2)
        p_a = 1 / (1 + math.exp(quality[b] - quality[a]))
        pairs.append((a, b) if rng.random() < p_a else (b, a))
    ref = [math.log(p) for p in softmax([0.5, 0.0, 0.0, -0.5])]  # reference likes the worst one
    theta = [0.5, 0.0, 0.0, -0.5]
    for _ in range(200):
        logp = [math.log(p) for p in softmax(theta)]
        probs = softmax(theta)
        grad = [0.0] * 4
        for chosen, rejected in pairs:
            _, gc, gr = dpo_loss_and_grad(logp[chosen], logp[rejected], ref[chosen], ref[rejected], beta=0.5)
            for k in range(4):  # d log p(y) / d theta_k = [y == k] - p_k
                grad[k] += gc * ((chosen == k) - probs[k]) + gr * ((rejected == k) - probs[k])
        theta = [t - 0.5 * g / len(pairs) for t, g in zip(theta, grad)]
    final = softmax(theta)
    assert final[3] > final[2] > final[1] > final[0]

Evaluation

• Held-out preference accuracy and margin. The cheapest check, from the same metrics as training. Expect 65–75% on held-out pairs from this kind of data; 95% means you overfit or leaked.
• Win rate against the base model. 300 held-out prompts, generate with both models at the same sampling settings, ask a strong judge model which answer is better, and run each comparison twice with the order swapped (judges favor the first position). Count a split decision as a tie. With 300 prompts, the 95% interval on a win rate near 50% is about ±5.7 points (1.96 × √(0.25 / 300)), so a 54% result is not a win. Say so.
• Length control. Report mean response length for both models and the win rate within length buckets. If the tuned model is 40% longer and only wins in the long bucket, the judge is rewarding length.
• Capability check. A few hundred items from a math or knowledge benchmark the base model already handles. A drop of more than a couple of points means β is too low or the learning rate too high.
• Read 20 outputs. Pick at random and write down what changed. Graders weight this heavily, because it is where you find the failure the metrics missed.

Failure analysis to expect and write up

Reviewer rubric

Interviewer expectations

Mid-level (L4)

• Gets a DPO run with LoRA working end to end and publishes the adapter with a model card.
• Explains the DPO loss and why the reference model is there.
• Reports a win rate and a few example outputs.

Senior (L5)

• Scopes model, data and settings from a compute estimate, and checks the template, loss mask and step-0 loss before training.
• Evaluates with position-swapped judging, a confidence interval, length control and a capability check.
• Writes up at least two real failures (length growth, falling log-probs) with the evidence and the fix tried.

Staff (L6+)

• Compares DPO with PPO, GRPO and online DPO for this data and says what each would change and cost.
• Designs the ablation that matters (β, data filtering, length balancing) inside the budget, with seeds.
• Connects the result to what a team would do next: on-policy preference data, a reward model to label it, and the eval that would catch reward hacking.