LLM training, for someone who serves them

Pre-training, SFT, RLHF/PPO, DPO, GRPO, LoRA, distillation, evaluation; and, for every stage, what it changes at inference time. Every number and snippet here was computed or run (python3 + PyTorch); outputs are pasted verbatim.

1. The big picture pipeline

A modern chat model is built in stages. Each stage changes what the weights are optimized for while reusing the previous stage’s weights as the starting point.

2. Data & tokenization, scaling laws, compute

Tokenization: BPE

Models do not see characters or words; they see integer token ids. Byte-pair encoding (BPE) builds the vocabulary greedily: start from single characters (or bytes), repeatedly find the most frequent adjacent pair in the corpus and merge it into a new symbol. After ~100k merges you get a vocabulary where common words are single tokens and rare words split into pieces. Byte-level BPE starts from the 256 bytes so any input is representable (no unknown token).

Copyfrom collections import Counter
corpus = "low low low low low lower lower newest newest newest widest widest"
words = [tuple(w) + ("</w>",) for w in corpus.split()]
vocab = Counter(words)

def pair_counts(vocab):
    c = Counter()
    for w, n in vocab.items():
        for a, b in zip(w, w[1:]):
            c[(a, b)] += n
    return c

def merge(vocab, pair):
    out = Counter()
    for w, n in vocab.items():
        i, new = 0, []
        while i < len(w):
            if i < len(w) - 1 and (w[i], w[i + 1]) == pair:
                new.append(w[i] + w[i + 1]); i += 2
            else:
                new.append(w[i]); i += 1
        out[tuple(new)] += n
    return out

for step in range(6):
    pc = pair_counts(vocab)
    best = max(pc, key=lambda p: (pc[p], p))
    vocab = merge(vocab, best)
    print(f"merge {step+1}: {best} (count {pc[best]})")
for w, n in vocab.items():
    print(n, w)

merge 1: ('o', 'w') (count 7)
merge 2: ('l', 'ow') (count 7)
merge 3: ('t', '</w>') (count 5)
merge 4: ('s', 't</w>') (count 5)
merge 5: ('low', '</w>') (count 5)
merge 6: ('e', 'st</w>') (count 5)
5 ('low</w>',)
2 ('low', 'e', 'r', '</w>')
3 ('n', 'e', 'w', 'est</w>')
2 ('w', 'i', 'd', 'est</w>')

Read the output: o+w merges first (7 occurrences), then l+ow, and the end-of-word marker lets the algorithm learn the suffix est shared by “newest” and “widest”. (Ties between equal counts are broken arbitrarily; real implementations fix a deterministic rule.)

Why inference people care: tokens, not characters, are the unit of cost. Latency and price scale with token count; a tokenizer with a bigger vocab (128k vs 32k) packs more text per token (fewer decode steps for the same text) but enlarges the embedding and LM-head matrices (2 × 128,256 × 4096 = 1.05B params in our config) and the final-logits matmul. Languages and code that tokenize poorly cost more to serve for the same content.

Data mixture, dedup, quality filtering

• Mixture: web crawl (the bulk), code, books/papers, math, multilingual, and curated high-quality sources, with sampling weights. Upweighting code and math early or late in training is a common lever; the weights are tuned with small proxy runs.
• Dedup: exact (hash) and fuzzy (MinHash/LSH over n-grams) dedup, at document and sometimes paragraph level. Duplicated text is memorized rather than generalized, wastes tokens, and inflates benchmark leakage.
• Quality filtering: heuristics (length, symbol ratios, language ID), perplexity filters using a small reference LM, and classifiers trained to recognize “textbook-like” text. A strong recent pattern is using an existing LLM to score or rewrite documents.
• Safety/PII filtering and decontamination: remove personal data, toxic content, and any n-gram overlap with evaluation sets (see section 11).
• Annealing / data curriculum: many recipes shift the mixture toward higher-quality, code/math-heavy data near the end of training with a decaying LR. Treat the specific proportions as lab-private.

Scaling laws and compute

Training cost for a dense transformer is well approximated by C ≈ 6 · N · D FLOPs, where N = parameters and D = training tokens. Derivation: forward pass ≈ 2 FLOPs per parameter per token (one multiply, one add); backward ≈ 2× the forward (gradients w.r.t. activations and weights), so 2 + 4 = 6. (Attention’s sequence-length term is extra and small at short context.)

Chinchilla (Hoffmann et al., 2022) found that for a fixed compute budget the loss-optimal split grows N and D together, at roughly D ≈ 20 · N tokens. Combine with C = 6ND: N_opt = √(C / 120).

Copyimport math
def flops(N, D): return 6 * N * D
N, D = 70e9, 20 * 70e9
C = flops(N, D)
print(f"70B @ 20 tok/param: D={D/1e12:.1f}T tokens, C={C:.2e} FLOPs")
h100, mfu = 989e12, 0.40                  # dense bf16 peak, assumed utilisation
print(f"  = {C/(h100*mfu)/3600:,.0f} H100-hours; {C/(h100*mfu)/4096/86400:.1f} days on 4096 GPUs")
Cb = 1e24; Nopt = math.sqrt(Cb / (6 * 20))
print(f"compute-optimal for 1e24 FLOPs: N={Nopt/1e9:.1f}B, D={20*Nopt/1e12:.2f}T")
print(f"8B on 15T tokens: {15e12/8e9:.0f} tok/param, C={flops(8e9,15e12):.2e} FLOPs")
for N in (7e9, 70e9):
    gb = lambda b: b * N / 1e9
    print(f"{N/1e9:.0f}B: bf16 w {gb(2):.0f} + grad {gb(2):.0f} + fp32 master {gb(4):.0f} + m {gb(4):.0f} + v {gb(4):.0f} = {gb(16):.0f} GB")
print(f"70B state sharded over 64 GPUs: {16*70e9/64/1e9:.1f} GB/GPU (before activations)")
print(f"RLHF-PPO 7B: policy 112 + critic 112 + ref 14 + RM 14 = {(112+112+14+14)} GB")

70B @ 20 tok/param: D=1.4T tokens, C=5.88e+23 FLOPs
  = 412,875 H100-hours; 4.2 days on 4096 GPUs
compute-optimal for 1e24 FLOPs: N=91.3B, D=1.83T
8B on 15T tokens: 1875 tok/param, C=7.20e+23 FLOPs
7B: bf16 w 14 + grad 14 + fp32 master 28 + m 28 + v 28 = 112 GB
70B: bf16 w 140 + grad 140 + fp32 master 280 + m 280 + v 280 = 1120 GB
70B state sharded over 64 GPUs: 17.5 GB/GPU (before activations)
RLHF-PPO 7B: policy 112 + critic 112 + ref 14 + RM 14 = 252 GB

3. Transformer architecture essentials (and what each costs at inference)

One decoder block: x → RMSNorm → attention → +residual → RMSNorm → FFN (SwiGLU or MoE) → +residual, repeated L times, then a final norm and an LM head that projects to vocabulary logits.

Attention

Each token is projected into a query q, key k, value v. For token t: out_t = Σ_{s≤t} softmax_s(q_t · k_s / √d_h) · v_s (causal mask: only past tokens). Multiple heads run this in parallel on slices of the hidden size and are concatenated and projected by W_o.

During decode, the keys and values of all previous tokens are cached (the KV cache) so each new token only computes its own q, k, v. The cache is read in full every step, so attention decode is memory-bandwidth-bound and the cache competes with weights for VRAM.

MHA vs GQA vs MQA

KV bytes per token = 2 (K and V) × layers × kv_heads × head_dim × bytes_per_element. For our config in bf16: 2 × 32 × 8 × 128 × 2 = 131,072 B = 128 KiB.

Copyd, L, H, KV, hd, ffn, V = 4096, 32, 32, 8, 128, 14336, 128256   # an 8B-class Llama-style shape
attn  = d*H*hd + 2*d*KV*hd + H*hd*d       # Wq, Wk+Wv (GQA), Wo
mlp   = 3*d*ffn                           # SwiGLU: gate, up, down
layer = attn + mlp + 2*d                  # + 2 RMSNorm gains
total = L*layer + 2*V*d + d               # untied in/out embeddings + final norm
print(f"per layer: attn {attn/1e6:.1f}M  mlp {mlp/1e6:.1f}M   total {total/1e9:.2f}B")
print(f"MHA attn would be {4*d*d/1e6:.1f}M/layer (GQA saves {(4*d*d-attn)/1e6:.1f}M)")
for name, kv in [("MHA", 32), ("GQA-8", 8), ("MQA", 1)]:
    b = 2*L*kv*hd*2                       # K and V, bf16
    print(f"{name:6s} KV cache {b/1024:5.0f} KiB/token | 8k ctx {b*8192/2**30:6.3f} GiB | 128k ctx {b*131072/2**30:5.1f} GiB")

per layer: attn 41.9M  mlp 176.2M   total 8.03B
MHA attn would be 67.1M/layer (GQA saves 25.2M)
MHA    KV cache   512 KiB/token | 8k ctx  4.000 GiB | 128k ctx  64.0 GiB
GQA-8  KV cache   128 KiB/token | 8k ctx  1.000 GiB | 128k ctx  16.0 GiB
MQA    KV cache    16 KiB/token | 8k ctx  0.125 GiB | 128k ctx   2.0 GiB

That 128k-context line is the whole argument for GQA: 64 GiB of cache for one MHA sequence versus 16 GiB with 8 KV heads (4× smaller), with 25.2M fewer attention parameters per layer as a bonus. Because decode reads the cache every step, a 4× smaller cache is also up to ~4× less attention bandwidth per token and more concurrent sequences per GPU.

RoPE (rotary position embeddings)

Instead of adding a position vector to the input, RoPE rotates each (q, k) pair of dimensions by an angle proportional to position, so the dot product q_m · k_n depends on the relative offset m − n. Inference consequences: positions are applied to q and k inside attention, so cached keys already contain their rotation (cache stays valid as you append); and context extension works by changing the rotation frequencies (position interpolation, NTK-aware scaling, YaRN-style methods) plus a short fine-tune on long sequences. The trained context length is therefore a training-time property you can partly stretch afterwards, with quality risk.

RMSNorm and SwiGLU

• RMSNorm: y = x / √(mean(x²) + ε) · g. Like LayerNorm without mean-subtraction or bias: cheaper, and trains as well. At inference it is a small memory-bound op that serving stacks fuse into neighbors (one of the “fused kernels” you see in profiles).
• SwiGLU FFN: down( silu(gate(x)) ⊙ up(x) ). Three matrices instead of two, so hidden width is set to ~(8/3)·d (here 14336 ≈ 3.5·4096, a larger multiple used by some families) to keep parameters comparable. The FFN holds ~80% of each layer’s parameters in our config (176.2M of 218.1M), so weight bytes read per decoded token are dominated by FFN weights.

Mixture of Experts (MoE)

Replace the single FFN by E expert FFNs plus a small router (a linear layer d→E). For each token the router picks the top-k experts and mixes their outputs by router weights. Training adds a load-balancing auxiliary loss (or bias-based balancing) so experts get similar traffic and none collapses.

Copyimport math
# speculative decoding: expected tokens per target forward, accept prob a, k draft tokens
for a, k in [(0.6, 4), (0.8, 4), (0.9, 8)]:
    print(f"a={a} k={k}: {(1-a**(k+1))/(1-a):.2f} tokens / target step")
# decode is bandwidth-bound: tokens/s <= BW / bytes_read_per_token
BW, N = 3.35e12, 8.03e9
for name, bpp in [("bf16", 2), ("int8", 1), ("int4", 0.5)]:
    print(f"{name}: weights {N*bpp/1e9:5.2f} GB -> <= {BW/(N*bpp):4.0f} tok/s at batch 1")
# MoE: 8 experts, top-2, same dims as above
d, L, H, KV, hd, ffn, V, E, k = 4096, 32, 32, 8, 128, 14336, 128256, 8, 2
attn = d*H*hd + 2*d*KV*hd + H*hd*d; mlp = 3*d*ffn
tot = L*(attn + E*mlp + 2*d + d*E) + 2*V*d + d
act = L*(attn + k*mlp + 2*d + d*E) + 2*V*d + d
print(f"MoE 8x top-2: total {tot/1e9:.1f}B params in memory, {act/1e9:.1f}B active per token")
print(f"perplexity = exp(loss): exp(2.0)={math.exp(2.0):.2f}, exp(1.0)={math.exp(1.0):.2f}")

a=0.6 k=4: 2.31 tokens / target step
a=0.8 k=4: 3.36 tokens / target step
a=0.9 k=8: 6.13 tokens / target step
bf16: weights 16.06 GB -> <=  209 tok/s at batch 1
int8: weights  8.03 GB -> <=  417 tok/s at batch 1
int4: weights  4.01 GB -> <=  834 tok/s at batch 1
MoE 8x top-2: total 47.5B params in memory, 13.7B active per token
perplexity = exp(loss): exp(2.0)=7.39, exp(1.0)=2.72

(The same script covers sections 3, 11, and 12; the MoE line is the one to look at here.) The 8-expert top-2 version of our config has 47.5B parameters to store but only 13.7B to compute per token: capacity of a ~47B model at roughly the per-token FLOPs of a ~14B model.

4. Pre-training mechanics

Objective: cross-entropy on the next token

Given a token sequence x_1…x_T, the model outputs at each position a distribution over the vocabulary. Loss = −(1/T) Σ_t log p_θ(x_{t+1} | x_1…x_t). It is self-supervised: the labels are the input shifted by one. All T positions train in parallel in one forward pass (causal masking prevents cheating), which is why training is far more parallel than decoding. A uniform guess over V tokens scores ln V (so perplexity V); real models drive the loss down to a few nats (perplexity single digits to tens) on web text.

Optimizer: AdamW

Adam keeps two running statistics per parameter: first moment m (mean gradient) and second moment v (mean squared gradient), and updates θ ← θ − lr · m̂ / (√v̂ + ε). AdamW decouples weight decay: θ ← θ − lr·λ·θ applied separately rather than folded into the gradient. Typical choices: β₁ = 0.9, β₂ = 0.95 (lower than the 0.999 default, for stability at large batch), weight decay ~0.1, gradient clipping at norm 1.0. Treat these as common published defaults, not universal.

Learning-rate schedule: warmup + cosine decay

Linear warmup from 0 to the peak LR over the first ~0.1-1% of steps (Adam’s second-moment estimates are noisy at the start, and large early updates can blow up), then cosine decay to ~10% of peak. Some recent recipes use warmup-stable-decay (flat LR, then a short fast decay) so you can branch off checkpoints cheaply.

Copyimport math
def lr(step, total=1000, warm=100, peak=3e-4, floor=3e-5):
    if step < warm: return peak * step / warm
    p = (step - warm) / (total - warm)
    return floor + 0.5 * (peak - floor) * (1 + math.cos(math.pi * p))
for s in (0, 50, 100, 300, 550, 1000):
    print(f"step {s:4d}  lr {lr(s):.2e}")

step    0  lr 0.00e+00
step   50  lr 1.50e-04
step  100  lr 3.00e-04
step  300  lr 2.68e-04
step  550  lr 1.65e-04
step 1000  lr 3.00e-05

Mixed precision: bf16 and fp8

bf16 has fp32’s exponent range with 8 bits of mantissa precision, so unlike fp16 it rarely needs loss scaling. Standard recipe: forward and backward matmuls in bf16, but keep an fp32 master copy of the weights and fp32 optimizer states, because tiny updates (lr × grad) would be rounded away in bf16.

Memory per parameter: where ~16 bytes comes from

So a 7B model needs 112 GB of training state (vs 14 GB to serve it in bf16, 8× more) and 70B needs 1,120 GB: far beyond one GPU, and activations come on top. Variants exist (fp32 grads → 20 B; 8-bit optimizers or bf16 optimizer states → ~10-12 B; no master copy → less, with quality risk), so “16” is the standard accounting, not a law.

Activation memory and checkpointing

The backward pass needs activations from the forward pass; stored for all layers, they scale with batch × sequence × hidden × layers and can dwarf the weights at long context. Activation (gradient) checkpointing stores only layer-boundary activations and recomputes the rest during backward. Cost: one extra forward pass, i.e. 4 units of work instead of 3 (forward 1 + backward 2), about +33% compute, for a large activation-memory saving. FlashAttention-style fused attention also avoids materializing the T × T score matrix.

Parallelism: what each axis shards and why

Large runs combine them (“3D/4D parallelism”): TP inside a node, PP across nodes, FSDP/DP across replicas, EP for MoE. Sharding the 70B state (1,120 GB) over 64 GPUs gives 17.5 GB/GPU before activations.

5. Supervised fine-tuning (SFT)

SFT continues training the base model with the same cross-entropy loss on curated (prompt, ideal response) pairs. It does not teach much new knowledge; it teaches the model to behave like an assistant: follow instructions, adopt the chat format, answer then stop. Quality over quantity: tens of thousands of carefully written or filtered examples can already give a usable instruct model (the “LIMA” hypothesis: alignment is mostly surface format), though modern pipelines use millions, heavily synthetic.

Loss masking on prompt tokens

Compute the loss only on response tokens. Prompt tokens are the user’s text, not something we want the model to learn to generate, and including them spends gradient on predicting user inputs. Implementation: set the labels of prompt positions to -100 (PyTorch’s ignore_index), and normalize by the count of unmasked tokens.

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# SFT cross-entropy with prompt masking (-100 = ignored)
V = 10
logits = torch.randn(1, 6, V)                       # [batch, seq, vocab]
tokens = torch.tensor([[3, 1, 4, 1, 5, 9]])         # 3,1,4 = prompt; 1,5,9 = response
labels = tokens.clone(); labels[:, :3] = -100
# position t predicts token t+1
loss = F.cross_entropy(logits[:, :-1].reshape(-1, V), labels[:, 1:].reshape(-1), ignore_index=-100)
manual = -torch.stack([F.log_softmax(logits[0, t], -1)[tokens[0, t + 1]] for t in (2, 3, 4)]).mean()
print("SFT masked loss", round(loss.item(), 4), "manual", round(manual.item(), 4))
print("uniform-guess loss ln(V) =", round(torch.log(torch.tensor(float(V))).item(), 4), "-> perplexity", V)

SFT masked loss 2.9529 manual 2.9529
uniform-guess loss ln(V) = 2.3026 -> perplexity 10

Takeaway: the masked loss equals the hand-computed mean of −log p over just the three response positions (2.9529). It sits above the uniform-guess baseline ln 10 = 2.3026 only because these logits are random noise, not a trained model.

Chat templates

The conversation is serialized with special tokens marking roles and turn ends, e.g. <|system|> … <|user|> … <|assistant|> … <|end|> (the exact tokens are model-specific). The end-of-turn token is part of the response label so the model learns to emit it and stop. The serving stack must apply the same template the model was trained with: a wrong or missing template is one of the most common causes of mysteriously degraded output, and a missing stop token means generations never terminate.

Packing

Examples are short and variable; padding them wastes compute. Packing concatenates several examples into one fixed-length sequence (e.g. 8k tokens). To avoid cross-contamination, use a block-diagonal attention mask (each example attends only to itself) and reset position ids per example; naive packing without this lets one example attend to its neighbor, which is a mild train/serve mismatch. Packing can improve token throughput substantially for short-example data.

Overfitting and catastrophic forgetting

• SFT uses small LR (order 1e-5, 10-100× below pre-training peak), 1-3 epochs. More epochs on small data memorizes the examples.
• Catastrophic forgetting: aggressive fine-tuning on a narrow domain erodes general capabilities. Mitigations: lower LR, mixing in general/replay data, LoRA (which constrains the update), and model merging/weight averaging.
• SFT teaches imitation: it cannot distinguish a “good” from a “slightly worse” answer, and it can teach hallucination (training the model to answer questions whose facts it lacks). That gap motivates preference optimization.

6. Reward models and RLHF / PPO

RLHF (as published in InstructGPT) has three phases: SFT (section 5), train a reward model from human rankings, then optimize the policy against it with RL.

Step 1: the reward model (Bradley-Terry)

Humans see a prompt x and two responses and pick the better one: y_w (chosen) over y_l (rejected). The Bradley-Terry model says the probability that y_w beats y_l is a logistic function of the reward difference:

P(y_w ≻ y_l | x) = σ( r(x, y_w) − r(x, y_l) )

Train r by maximum likelihood, i.e. minimize L_RM = − E log σ( r(x,y_w) − r(x,y_l) ). Architecturally, r is the SFT model with the LM head replaced by a linear layer producing one scalar read at the last token. Only reward differences are identified (adding a constant to r changes nothing).

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# Bradley-Terry reward-model loss
def rm_loss(r_chosen, r_rejected):
    return -F.logsigmoid(r_chosen - r_rejected).mean()
rc, rr = torch.tensor([2.0, 0.5, -1.0]), torch.tensor([0.0, 0.5, 1.0])
print("RM loss", round(rm_loss(rc, rr).item(), 4), " equal rewards -> ln2 =", round(rm_loss(rr, rr).item(), 4))
print("P(chosen>rejected) =", [round(x, 3) for x in torch.sigmoid(rc - rr).tolist()])

RM loss 0.9823  equal rewards -> ln2 = 0.6931
P(chosen>rejected) = [0.881, 0.5, 0.119]

Sanity checks that held: equal rewards give ln 2 = 0.6931 (a coin flip); a margin of +2 means the chosen response wins with probability σ(2) = 0.881.

Step 2: policy optimization with a KL leash

Objective: maximize expected reward while staying close to a frozen reference policy π_ref (the SFT model):

max_π E_{x~D, y~π(·|x)} [ r(x,y) ] − β · KL( π(·|x) ‖ π_ref(·|x) )

The KL term is not decoration. The reward model is an imperfect proxy trained on limited data, so an unconstrained optimizer finds adversarial outputs that score high but are garbage (reward hacking, e.g. length exploitation, sycophancy). The KL keeps the policy in the region where the reward model is trustworthy and preserves fluency. In practice the penalty is applied per token: r_t = −β (log π(a_t|s_t) − log π_ref(a_t|s_t)) for every token, plus the RM score added at the final token.

PPO: advantage and clipping

• A value network (critic) V(s) estimates expected future reward from a partial response. The advantage A_t = how much better token t was than the critic expected. Generalized Advantage Estimation: δ_t = r_t + γ V(s_{t+1}) − V(s_t), A_t = Σ_l (γλ)^l δ_{t+l}.
• Policy ratio ρ_t = π_θ(a_t|s_t) / π_old(a_t|s_t) compares the current policy to the one that generated the samples.
• Clipped objective: L = E[ min( ρ_t A_t , clip(ρ_t, 1−ε, 1+ε) A_t ) ], ε ≈ 0.2. It removes the incentive to move the policy more than ~ε per update, so you can reuse a batch of samples for several gradient steps without the policy running away.

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# PPO clipped objective for one token
def ppo_obj(logp_new, logp_old, A, eps=0.2):
    ratio = (logp_new - logp_old).exp()
    return torch.min(ratio * A, ratio.clamp(1 - eps, 1 + eps) * A), ratio
for A in (1.0, -1.0):
    for lr_ in (0.5, 0.0, -0.5):
        o, r = ppo_obj(torch.tensor(lr_), torch.tensor(0.0), torch.tensor(A))
        print(f"A={A:+.0f} ratio={r.item():.3f} objective={o.item():+.3f}")

A=+1 ratio=1.649 objective=+1.200
A=+1 ratio=1.000 objective=+1.000
A=+1 ratio=0.607 objective=+0.607
A=-1 ratio=1.649 objective=-1.649
A=-1 ratio=1.000 objective=-1.000
A=-1 ratio=0.607 objective=-0.800

Read it: with A = +1, once the ratio exceeds 1.2 the objective stops growing (1.649 → 1.200), so pushing probability up further earns nothing. With A = −1, the min keeps the pessimistic value: for ratio 0.607 the objective is −0.800 (clipped, no further reward for lowering it), but for ratio 1.649 it stays −1.649 (unclipped, so a mistake that made a bad token more likely is still fully penalized).

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# KL-shaped reward and k3 estimator
lp_pol, lp_ref = torch.tensor([-1.0, -2.0, -0.5]), torch.tensor([-1.2, -1.5, -0.9])
kl_tok = lp_pol - lp_ref
print("per-token log-ratio", [round(x, 2) for x in kl_tok.tolist()], "sum", round(kl_tok.sum().item(), 2),
      "-> RLHF reward = r - beta*sum, beta=0.1, r=3.0 ->", round(3.0 - 0.1 * kl_tok.sum().item(), 3))

per-token log-ratio [0.2, -0.5, 0.4] sum 0.1 -> RLHF reward = r - beta*sum, beta=0.1, r=3.0 -> 2.99

Why it’s expensive and unstable

• Generation is inside the training loop. Every PPO step must sample full responses from the current policy (autoregressive decode, slow and memory-bound) and then do a training step on them. Real RLHF systems therefore embed an inference engine (vLLM-style) next to the trainer and resync weights each iteration. This is a big reason frontier labs hire inference-optimization engineers for RL.
• Instability: sensitivity to hyperparameters, critic quality, reward hacking, KL coefficient tuning (too low: hacking; too high: no learning), and noisy rewards. Many practitioners moved to simpler methods (DPO, GRPO) for these reasons.

7. DPO: Direct Preference Optimization

DPO (Rafailov et al., 2023) shows you can skip the explicit reward model and the RL loop: the KL-regularized objective from section 6 has a closed-form optimum, which lets you rewrite the Bradley-Terry loss directly in terms of the policy.

Derivation

1. The KL-regularized reward-maximization problem has the optimal policy

   π*(y|x) = (1/Z(x)) · π_ref(y|x) · exp( r(x,y)/β )

2. Solve for the reward: r(x,y) = β · log( π*(y|x) / π_ref(y|x) ) + β · log Z(x). The reward is the (scaled) log-ratio of policy to reference, up to a prompt-dependent constant.
3. Substitute into Bradley-Terry P(y_w ≻ y_l) = σ(r(x,y_w) − r(x,y_l)). The intractable β log Z(x) is the same for both responses and cancels.
4. Replace π* by the trainable π_θ and take the negative log-likelihood:

L_DPO = − E_{(x,y_w,y_l)} log σ( β [ log π_θ(y_w|x)/π_ref(y_w|x) − log π_θ(y_l|x)/π_ref(y_l|x) ] )

Here log π(y|x) is the sum of token log-probs over the response. In words: increase the chosen response’s log-prob relative to the reference, decrease the rejected one’s, with the “implicit reward” β log(π_θ/π_ref) playing the role of r. The gradient weights each pair by σ(−margin), i.e. by how wrong the implicit reward currently is, so already-correct pairs fade out.

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# DPO loss
def dpo_loss(pol_c, pol_r, ref_c, ref_r, beta=0.1):
    margin = beta * ((pol_c - ref_c) - (pol_r - ref_r))
    return -F.logsigmoid(margin).mean(), margin
# sequence log-probs (sum over response tokens)
ref_c, ref_r = torch.tensor([-42.0]), torch.tensor([-40.0])
l0, _ = dpo_loss(ref_c.clone(), ref_r.clone(), ref_c, ref_r)
print("DPO loss at policy==ref:", round(l0.item(), 4), "(= ln2)")
pol_c = ref_c.clone().requires_grad_(); pol_r = ref_r.clone().requires_grad_()
l, m = dpo_loss(pol_c, pol_r, ref_c, ref_r, beta=0.1); l.backward()
print("grad wrt logp(chosen) =", round(pol_c.grad.item(), 4), " wrt logp(rejected) =", round(pol_r.grad.item(), 4),
      " (= -/+ beta*sigmoid(-margin) = -/+", round(0.1 * torch.sigmoid(-m).item(), 4), ")")
l2, m2 = dpo_loss(torch.tensor([-40.0]), torch.tensor([-45.0]), ref_c, ref_r)
print("after chosen up 2, rejected down 5: margin", round(m2.item(), 3), "loss", round(l2.item(), 4))

DPO loss at policy==ref: 0.6931 (= ln2)
grad wrt logp(chosen) = -0.05  wrt logp(rejected) = 0.05  (= -/+ beta*sigmoid(-margin) = -/+ 0.05 )
after chosen up 2, rejected down 5: margin 0.7 loss 0.4032

Checks that held: at initialization (policy = reference) the loss is exactly ln 2 = 0.6931; the gradient on the chosen log-prob is −β·σ(0) = −0.05 and +0.05 on the rejected one (equal and opposite, scaled by β = 0.1); after the policy moves (chosen +2, rejected −5 nats vs the reference) the margin is 0.1 × (2 + 5) = 0.7 and the loss falls to 0.4032.

Beta, data, and practicalities

• β (typically ~0.1 to 0.5 in published setups) is the KL strength: larger β keeps the policy tighter to the reference; smaller β lets it move farther. It is also the scale of the implicit reward.
• Data: triples (prompt, chosen, rejected). Human-labeled, or AI-labeled (RLAIF, section 10), or constructed from model samples ranked by a reward model or a judge.
• Memory: two models (policy trainable + frozen reference) instead of four. Reference log-probs can even be precomputed once offline. No generation in the loop for the vanilla version.
• Offline vs online: vanilla DPO is offline: the preference pairs come from some other model’s distribution, so as the policy drifts, the data goes off-policy and gains can plateau. Online/iterative DPO regenerates pairs from the current policy (and labels them with a judge or RM) each round, recovering much of RL’s on-policy benefit at lower cost.
• Known quirks: DPO can reduce the log-prob of both chosen and rejected (only their gap is optimized), can lengthen responses, and is sensitive to data quality.

8. GRPO and RL with verifiable rewards (reasoning models)

For tasks where correctness is checkable (math answer matches, code passes unit tests, format constraints) you do not need a learned reward model: use a programmatic verifiable reward (RLVR). This is the core ingredient of most “reasoning model” training: the model samples a long chain of thought, the final answer is verified, and RL reinforces chains that led to correct answers.

GRPO (Group Relative Policy Optimization, introduced in the DeepSeekMath work) removes PPO’s critic: sample a group of G responses per prompt, score each, and use the group-normalized reward as the advantage: A_i = (r_i − mean(r)) / std(r). The same clipped-ratio objective and a KL term to a reference are then used. This drops one of the four models (the value network: 112 GB of the 252 GB in the 7B example of section 6).

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# GRPO group-relative advantages
rewards = torch.tensor([1., 0., 0., 1., 1., 0., 0., 0.])
adv = (rewards - rewards.mean()) / (rewards.std() + 1e-6)
print("GRPO rewards", rewards.tolist(), "-> advantages", [round(a, 3) for a in adv.tolist()])

GRPO rewards [1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0] -> advantages [1.208, -0.725, -0.725, 1.208, 1.208, -0.725, -0.725, -0.725]

3 of 8 samples correct: correct ones get advantage +1.208, wrong ones −0.725 (they sum to zero). Edge cases worth knowing: if all samples in a group are right or all wrong, every advantage is 0 and that prompt yields no gradient, so prompt difficulty must be matched to the model’s current ability. With a single correct sample out of 4 the correct one gets +1.5 and the others −0.5.

9. PEFT, LoRA, QLoRA

Parameter-efficient fine-tuning (PEFT) freezes the backbone and trains a tiny number of added parameters. The dominant method is LoRA (Hu et al., 2021). The observation: the weight update ΔW during fine-tuning has low intrinsic rank, so represent it as a product of two thin matrices.

W' = W + (α/r) · B A,  W ∈ R^{d_out×d_in} frozen,  B ∈ R^{d_out×r},  A ∈ R^{r×d_in},  r ≪ d

• Init: A random (small Gaussian/Kaiming), B = 0, so at step 0 the adapter contributes exactly nothing and the model starts as the base.
• Rank r sets capacity (common: 8-64). alpha sets the scale α/r of the update; with α fixed and r changed the effective step size changes, which is why people often set α = 2r or keep α constant and tune LR. (rsLoRA proposes α/√r.)
• Typically applied to attention projections, increasingly all linear layers including the FFN.

Copyimport torch, torch.nn as nn
torch.manual_seed(0)
class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r=16, alpha=32):
        super().__init__()
        self.base, self.scale = base, alpha / r
        for p in base.parameters(): p.requires_grad_(False)
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)   # r x d_in
        self.B = nn.Parameter(torch.zeros(base.out_features, r))         # d_out x r, zero init
    def forward(self, x):
        return self.base(x) + self.scale * (x @ self.A.T @ self.B.T)
    @torch.no_grad()
    def merge(self):
        self.base.weight += self.scale * (self.B @ self.A)
        return self.base

base = nn.Linear(4096, 4096, bias=False)
m = LoRALinear(base, r=16, alpha=32)
x = torch.randn(2, 4096)
print("B=0 at init -> same output as base:", torch.allclose(m(x), base(x)))
trainable = sum(p.numel() for p in m.parameters() if p.requires_grad)
print("trainable", trainable, "frozen", base.weight.numel(), f"-> {base.weight.numel()/trainable:.0f}x fewer")
with torch.no_grad(): m.B.normal_(0, 0.01)           # pretend training happened
y = m(x).clone()
merged = m.merge()
print("merged W' = W + (alpha/r)BA matches adapter forward:", torch.allclose(merged(x), y, atol=1e-4))
print("scale alpha/r =", 32 / 16, "; rank(BA) <= 16:", torch.linalg.matrix_rank(m.B @ m.A).item())

B=0 at init -> same output as base: True
trainable 131072 frozen 16777216 -> 128x fewer
merged W' = W + (alpha/r)BA matches adapter forward: True
scale alpha/r = 2.0 ; rank(BA) <= 16: 16

Parameter-count example (our 8B config, r = 16)

(From my parameter script: per(i,o) = r·(i+o).) Training memory: the frozen backbone needs only weights (2 B/param) with no grads or Adam states for it; the 16 B/param applies only to the adapter (tens of MB), plus activations. So fine-tuning a 7B model needs roughly the 14 GB of weights plus activations, versus 112 GB for full fine-tuning.

QLoRA

QLoRA (Dettmers et al., 2023) quantizes the frozen base to 4-bit NF4 (NormalFloat4: 16 levels placed at quantiles of a normal distribution, suited to bell-shaped weights) with double quantization of the per-block scale constants, keeps the LoRA adapters in bf16, dequantizes blocks on the fly for the matmul, and uses paged optimizers to survive memory spikes. The base weights of a 7B model drop to 7e9 × 0.5 B = 3.5 GB (plus small quantization-constant overhead), so a 7B fine-tune fits a 24 GB card and a 65-70B fine-tune fits ~48 GB (the paper’s headline claim). Gradients flow through the dequantized weights to the adapters only.

Merging and multi-LoRA serving

• Merge: W ← W + (α/r) B A is a one-time add (verified above). Merged model = zero extra inference cost, identical architecture. Downside: one merged model per adapter, and merging into an already-quantized base needs care (merge in higher precision then re-quantize).
• Unmerged / multi-LoRA serving: keep one base model resident and many small adapters (each ~84 MB at r=16 all-linear in our example, versus 16 GB for a full copy). Each request names its adapter; the engine computes y = xW + (α/r)(xAᵀ)Bᵀ per request. Batching requests with different adapters needs grouped/batched-gather matmul kernels (S-LoRA / Punica-style “BGMV/SGMV”), plus adapter paging between CPU and GPU. This is how a platform serves thousands of tenant fine-tunes on a few GPUs; vLLM and others expose it as a feature.
• Costs of unmerged serving: extra small matmuls and kernel launches per layer, and lower effective batch efficiency when many distinct adapters are in a batch.

10. RLAIF, Constitutional AI, synthetic data, distillation

RLAIF and Constitutional AI

Human preference labels are slow and expensive. RLAIF replaces (some) human labelers with an LLM judge that picks the better response. Constitutional AI (Anthropic, 2022) is the best-known recipe: write a short list of principles (the “constitution”); in a supervised phase the model critiques and revises its own outputs against the principles and is fine-tuned on the revisions; in an RL phase an AI feedback model, prompted with the principles, produces the preference labels that train the reward model (or DPO pairs). Benefits: scalability, cheap iteration on policy, and auditable principles. Risks: the labeler’s biases and blind spots are inherited and amplified; labels are only as good as the judge; and self-preference (a model favoring its own style) can create feedback loops.

Synthetic data

Models generate training data for models: instruction/response pairs (self-instruct), chain-of-thought solutions filtered by a verifier (rejection sampling / best-of-n then SFT), code with executable tests, and rewritten web text. It is now a main source of post-training data. Failure modes: homogeneity, model collapse when models repeatedly train on their own unfiltered outputs, and contamination of eval sets. Verification (tests, answer checking) and diversity controls are what separate useful synthetic data from noise.

Distillation: how small fast models are made

Train a small student to mimic a large teacher. Two forms: sequence-level (SFT the student on teacher-generated outputs; the common approach, works when only an API is available) and logit-level (match the teacher’s full token distribution via KL, which carries far more information per token: “dark knowledge” about near-miss tokens). With temperature T softening both distributions, the loss is T² · KL(softmax(z_t/T) ‖ softmax(z_s/T)) (the T² keeps gradient magnitudes comparable across T).

Copyimport torch, torch.nn.functional as F
torch.manual_seed(0)
# distillation KL with temperature
t_logits, s_logits = torch.tensor([[4.0, 1.0, 0.0]]), torch.tensor([[2.0, 1.5, 0.0]])
for T in (1.0, 4.0):
    kl = F.kl_div(F.log_softmax(s_logits / T, -1), F.softmax(t_logits / T, -1), reduction="batchmean") * T * T
    print(f"distill KL(teacher||student) T={T}: {kl.item():.4f}")

distill KL(teacher||student) T=1.0: 0.3382
distill KL(teacher||student) T=4.0: 0.6658

(That KL is ≥ 0 and is 0 only when the student matches the teacher. Note the T² factor; the larger value at T = 4 here partly reflects it.) Logit distillation requires teacher and student to share a tokenizer, otherwise use sequence-level distillation.

11. Evaluation

Perplexity

Perplexity = exp(average cross-entropy per token in nats). It is the effective branching factor: a loss of 2.0 nats means the model is as uncertain as a uniform choice among e² ≈ 7.39 tokens; loss 1.0 → 2.72. Good for tracking pre-training progress and comparing models with the same tokenizer on the same text. It is not comparable across tokenizers (different tokens per word; use bits-per-byte), and it says little about instruction following, safety, or reasoning. For quantization it is the standard cheap regression test: a small perplexity increase after INT4 on a held-out set is a necessary (not sufficient) check.

Benchmarks and contamination

• Knowledge/reasoning multiple-choice (MMLU-style), math (GSM8K/MATH-style), code (HumanEval/MBPP-style, pass@k), long-context retrieval, instruction following, and agentic/task suites. Exact benchmark choices change quickly; check current ones.
• Contamination: if benchmark questions (or near-copies) are in the training set, scores overstate ability. Defenses: n-gram overlap decontamination, held-out/private test sets, fresh time-split benchmarks (problems published after the training cutoff), and canary strings. Saturated benchmarks stop discriminating.
• Always compare under identical prompting, few-shot count, and decoding settings; small prompt-format changes can shift scores by several points.

LLM-as-judge and win-rate

For open-ended quality, an LLM compares two responses (pairwise) or scores one (pointwise); aggregate to a win-rate. Convention: win-rate = (wins + 0.5 · ties) / total, so 60 wins, 30 ties, 10 losses of 100 gives (60 + 15) / 100 = 75%. Pitfalls:

• Position bias: judges favor the first (or second) slot; fix by running both orders and counting only consistent verdicts.
• Verbosity/length bias: longer answers win; use length-controlled win-rate.
• Self-preference: judges favor their own family’s style.
• Noisy on correctness: judges miss subtle factual or code errors; prefer verifiable checks where possible. Report confidence intervals, sample size, and judge model/version.
• Goodhart: optimizing a policy against a judge (or RM) eventually exploits it; hold out a different judge and human spot-checks.

12. How training choices show up at inference

Numbers for the headline levers

• Weight bytes → decode speed. At batch 1, decode reads all weights once per token, so tokens/s ≲ bandwidth ÷ weight bytes. With 3.35 TB/s (an H100-SXM-class figure) and 8.03B params: bf16 ≤ 209 tok/s, INT8 ≤ 417, INT4 ≤ 834 (upper bounds ignoring KV reads, scales/zeros overhead, and kernel inefficiency; see the serve.py output above).
• Speculative decoding. With acceptance probability α per drafted token (assumed independent) and k drafted tokens, the expected tokens emitted per target forward is (1 − α^{k+1}) / (1 − α): α=0.6,k=4 → 2.31; α=0.8,k=4 → 3.36; α=0.9,k=8 → 6.13. The speedup is less than that because the draft costs time, so the draft must be much cheaper and well aligned with the target.
• KV cache sizing. Concurrent 128k-token sequences per GPU on a 141 GB-class card with ~16 GB (bf16) of weights: (141 − 16 − some headroom) / 16 GiB ≈ 6-7 with GQA-8 bf16, versus 1 with MHA. (Rough, ignores activations/workspace.)

13. Interview prep

Try to answer aloud before opening each one.