Reward Model Training on GPU, in depth: the paired training step, last-token pooling, sizing, evaluation and the bugs that silently ruin a reward model

What the model is

Start from a decoder-only language model, usually the SFT checkpoint of the policy family or something close to it, so the reward model already understands the distribution it will judge. Remove the vocabulary projection and attach a linear layer from the hidden size to one output. Run the prompt and response through the transformer, take the hidden state at the final token of the response, and project it to a scalar. Because the model is causal, the last token is the only position whose state has attended to the entire sequence, which is why pooling happens there.

Training data are comparisons: a prompt, a chosen response and a rejected response, collected from human raters or from a stronger model acting as judge. Nothing in the data says how good either response is in absolute terms, only which one was preferred.

The loss, with numbers

The Bradley-Terry model says the probability that the chosen response wins is the sigmoid of the reward difference. Training maximises that probability, so the loss for one pair is minus the log of sigmoid(r_chosen minus r_rejected).

Work through one pair. Suppose the model currently scores the chosen answer 1.3 and the rejected one 0.5. The margin is 0.8, sigmoid(0.8) is 0.690, and the loss is 0.371. If the model had them backwards, a margin of minus 0.8, the loss would be 1.171 and the gradient would push hard to swap them. At a margin of zero the loss is ln 2, about 0.693. That last number is your first sanity check: a freshly initialised reward model should start near 0.693, because its scores are close to equal. A first-step loss far from it means the head was initialised with large weights or the pooling is reading something other than what you think.

Because only differences enter the loss, adding a constant to every reward changes nothing. The absolute level drifts freely during training, which is harmless for ranking but confusing for RL, where the reward feeds into advantages and a KL penalty. Two common fixes are a small penalty on the squared mean reward and normalising rewards after training with a mean and standard deviation measured on a reference set.

The training step in PyTorch

The step below concatenates chosen and rejected sequences into one batch so the backbone runs once, then splits the scalars. It includes an optional per-pair margin, which Llama 2 used to push the reward gap wider for pairs that raters marked as clearly better, and an optional centring term.

import torch
import torch.nn.functional as F

class RewardModel(torch.nn.Module):
    def __init__(self, backbone, hidden_size):
        super().__init__()
        self.backbone = backbone                    # decoder-only LM without its lm_head
        self.head = torch.nn.Linear(hidden_size, 1)
        torch.nn.init.normal_(self.head.weight, std=1 / (hidden_size + 1) ** 0.5)
        torch.nn.init.zeros_(self.head.bias)

    def forward(self, input_ids, attention_mask):
        h = self.backbone(input_ids=input_ids,
                          attention_mask=attention_mask).last_hidden_state   # [N, T, d]
        # Index of the last real token, correct for right AND left padding.
        pos = torch.arange(h.size(1), device=h.device)
        last = (attention_mask * pos).argmax(dim=1)                           # [N]
        pooled = h[torch.arange(h.size(0), device=h.device), last]           # [N, d]
        return self.head(pooled).squeeze(-1).float()                         # [N]

def train_step(model, batch, optimizer, margin=None, center_coef=0.0):
    ids = torch.cat([batch["chosen_ids"], batch["rejected_ids"]])
    mask = torch.cat([batch["chosen_mask"], batch["rejected_mask"]])
    with torch.autocast("cuda", dtype=torch.bfloat16):
        r = model(ids, mask)
    r_c, r_r = r.chunk(2)
    diff = r_c - r_r
    if margin is not None:                  # optional per-pair margin, as in Llama 2
        diff = diff - margin
    loss = -F.logsigmoid(diff).mean()
    if center_coef:                         # keep rewards near zero mean
        loss = loss + center_coef * ((r_c + r_r) ** 2).mean()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step(); optimizer.zero_grad(set_to_none=True)
    acc = (r_c > r_r).float().mean()
    return loss.item(), acc.item(), diff.detach().mean().item()

Track three numbers every step: loss, pairwise accuracy within the batch and mean margin. Loss and accuracy should move together. A mean margin that keeps growing while held-out accuracy stalls is the reward model becoming overconfident, which RL will later exploit.

Last-token pooling and the padding bug

The most common silent failure is reading the hidden state at the wrong position. Batches are padded to a common length, and whether padding goes on the right or the left depends on the tokenizer and on whether the code was written for training or for generation. A pooling rule that is correct for one layout is wrong for the other:

mask (right padded)  = [1, 1, 1, 1, 0, 0]   last real token is index 3
mask.sum() - 1       = 3                    correct
position -1          = 5                    a pad token: wrong

mask (left padded)   = [0, 0, 1, 1, 1, 1]   last real token is index 5
mask.sum() - 1       = 3                    a prompt token: wrong
(mask * pos).argmax  = 5                    correct for both layouts

Reading a pad position gives every sequence of the same length nearly the same reward, so the model learns length and padding patterns instead of quality. Reading a prompt token makes the reward ignore the response entirely. Neither crashes, and loss still falls a little, because the model finds whatever shortcut is available. The code above uses the argmax of mask times position, which is correct for both layouts. Two more guards help: make sure every sequence ends with an end-of-sequence token, so the pooled position is a consistent boundary rather than the last word of a truncated answer, and drop pairs where truncation cut off the response, because the model would be judging text the rater never saw.

Batching, packing and data hygiene

Each pair contributes two sequences that share a prompt, and lengths vary widely, so pad-to-longest batches waste a lot of compute. Bucket pairs by length, and size batches by token count rather than by pair count. Packing several sequences into one row is possible but needs per-sequence attention masking and per-sequence pooling indices, so most teams use length bucketing instead.

Data hygiene matters more than architecture:

• Randomise which response appears in which slot during collection. Raters and judges have position bias, and a model can learn it.
• Deduplicate prompts across train and evaluation splits, and keep all comparisons for one prompt in the same split.
• When one prompt has several ranked responses, InstructGPT's approach was to put all of its comparisons into one batch element rather than scattering them, which reduced overfitting to repeated prompts.
• Check length statistics: if chosen responses are much longer on average, the model will learn length. Balance the data or measure the effect explicitly.
• Drop pairs with identical or near-identical responses, and pairs where annotators disagreed, unless you model the disagreement.

Sizing memory and throughput

Full fine-tuning with AdamW in mixed precision costs about 16 bytes per parameter for weights, gradients, FP32 master weights and two optimizer moments. A 7-billion-parameter reward model therefore needs about 112 GB before activations, which is why it is sharded with FSDP or ZeRO across eight GPUs, about 14 GB of model state each, leaving most of an 80 GB GPU for activations. Activation checkpointing is usually worth it, because sequences are long and the batch holds two of them per pair.

Training compute is about 6 FLOPs per parameter per token. Take 64 pairs per step at an average of 1,024 tokens per sequence: 131,072 tokens per step, or 6 x 7e9 x 131,072, about 5.5e15 FLOPs. Eight H100 SXM GPUs have about 989 TFLOPS of dense BF16 peak each; at a realistic 40 percent utilisation that is about 3.2e15 FLOP/s, so a step takes about 1.7 seconds. One epoch over 100,000 pairs is about 1,560 steps, roughly 45 minutes. Reward models train for about one epoch, because they overfit quickly, so the expensive part of the project is data, not GPU time.

Hyperparameters that matter

• Epochs. One, sometimes two. Both InstructGPT and Llama 2 reported training for a single epoch; held-out accuracy typically drops after that while training accuracy keeps rising.
• Learning rate. Lower than SFT, commonly in the 1e-6 to 1e-5 range for full fine-tuning of billion-parameter models, with warm-up and cosine or linear decay. Higher rates destabilise the backbone.
• Head initialisation. Small weights and zero bias, so the starting loss sits near 0.693.
• Size. A reward model should not be much weaker than the policy it judges. A small reward model judging a large policy is easier to exploit.
• LoRA. Adapters cut memory sharply and are fine for experiments, but full fine-tuning usually gives better accuracy on hard comparisons; measure before choosing.

Evaluating a reward model

Held-out pairwise accuracy is the headline number, but it hides the failures that matter. Report it per category (safety, coding, reasoning, chat), because a model at 72 percent overall can be at 55 percent on the category RL will push hardest. Measure the correlation between reward and response length on a neutral set; a strong positive correlation predicts verbose policies. Check calibration by bucketing pairs by predicted probability and comparing with the observed win rate. Public benchmarks such as RewardBench give a comparison point across models, but your own held-out data from your own policy is what predicts RL behaviour.

The decisive test is downstream: run a short RL or best-of-n experiment and have humans or a stronger judge compare the outputs. A reward model that gains a point of accuracy but makes best-of-16 outputs worse is not better.

Serving the reward model inside RL

During PPO the reward model scores every rollout, so it becomes an inference service. It needs only a forward pass, no KV cache across steps, so it can run in BF16 with large batches. Freeze it, put it in eval mode, and pin its tokenizer and chat template to exactly what it saw in training; a template mismatch is invisible in logs and shifts every score. Normalise rewards with statistics from a reference set, and clip extremes. The RLHF pipeline on GPU covers whether to co-locate it with the policy or give it a dedicated pool.

Watch reward and policy behaviour together. A reward that keeps climbing while output length grows, or while a held-out judge's rating falls, is reward hacking. Fixes are a stronger KL penalty, an ensemble of reward models whose disagreement is penalised, or retraining the reward model on fresh comparisons from the current policy.

Failure modes

Trade-offs

A separate reward model costs data, GPU time and a serving component, and it can be hacked. In exchange it can score any response, including ones the policy will generate tomorrow, which is what online RL needs. DPO skips the reward model by training the policy directly on pairs; that is cheaper and simpler, but it learns only from the fixed dataset. For verifiable tasks such as maths or code with tests, rule-based rewards avoid learned-reward hacking entirely. Many teams combine them: verifiable rewards where possible, a learned reward model for open-ended quality, and periodic refreshes from on-policy data.