RL Post-Training — Concepts & Diagnostic Guide
Core Concepts
RL post-training optimizes a language model's policy using reward signals. The standard pipeline:
prompt → generate completions → score with reward → compute advantages → policy gradient update
Each stage has distinct failure modes. When a model "shows no improvement," the bug could be anywhere in this pipeline.
Diagnostic Methodology
When RL training produces no improvement, work through these stages in order. Each stage depends on the previous one being correct.
Stage 1: Verify Reward Signal
- Are rewards non-constant? If all rewards are identical, there is no learning signal.
- Do rewards correlate with completion quality? Spot-check decoded completions against their scores.
- Is the reward function being called on the correct text? Check that decoding/stripping preserves the content the reward function needs to evaluate.
Stage 2: Verify Advantage Computation
- Are advantages non-zero when rewards vary? If they collapse to ~0, the policy gradient vanishes.
- Check the magnitude and dtype of every numerical-stability constant in the advantage path (additive epsilons, clipping bounds). Compare each to what the math requires.
- Check the group size
G.G ≤ 2makesstdeither undefined or extremely noisy.
Stage 3: Verify Log-Probability Computation
- Verify bounds: log-probs of valid tokens must be non-positive.
- Compare your implementation against
F.log_softmaxon a small deterministic input — a numerical match rules out sign errors, wrong gathering axis, and off-by-one subtraction. - On a near-one-hot input, confirm the dominant token's log-prob is close to 0 (not close to the min).
Stage 4: Verify Loss Computation
- Is the loss changing across steps? Flat loss suggests zero gradients upstream.
- Log the KL term and the policy-gradient term separately. Either dominating the other is diagnostic.
- Check the fraction of clipped samples. Near-100% clipping means the clip range is starving the signal.
Stage 5: Verify Generation and Decoding
- Are padding and decoder artefacts stripped from the text the reward function sees?
- Does every completion shape your model can emit survive the decoding path with non-empty output where a human would expect non-empty output? Include degenerate cases (no formatting markers, only a prefix, only a suffix) in the round-trip test.
- Print or log a sample of the actual strings handed to the reward function — mismatches between "what the model generated" and "what the reward saw" are often visible at a glance.
Fixing, Not Rewriting
A bug in a branched function is a bug in one branch, not a verdict on the whole function. Before editing:
- Enumerate the input shapes the function handles today and the output each branch produces. The branches exist because callers rely on them.
- Identify which input-output pairs violate the intended contract. Those are the only branches you need to change.
- If your diff collapses a multi-branch function to a one-liner, you have almost certainly broken a contract a different caller depends on. Re-read the call sites before committing.
The same principle applies to epsilon values, sign conventions, and clipping bounds — if a constant looks wrong, replace it with a correct constant; don't remove the surrounding numerical-stability logic.
Key Mathematical Invariants
These invariants should always hold. If any is violated, there is a bug.
| Invariant | What it means | How to check |
|---|---|---|
log_prob <= 0 | Log of a probability is non-positive | assert (log_probs <= 1e-6).all() |
log-probs match F.log_softmax | Manual implementation equals the reference | torch.allclose(manual, F.log_softmax(...).gather(...)) |
sum(softmax(logits)) == 1 | Probabilities sum to 1 | assert torch.allclose(softmax.sum(-1), ones) |
0 < epsilon << 1 | Additive epsilons exist for numerical stability only | assert 0 < epsilon < 1 |
advantages != 0 when rewards vary | Non-constant rewards yield non-zero advantages | assert advantages.abs().max() > 0.1 |
| Decoding preserves content | Every generation shape the model emits survives the decode path with the content a human reader would expect | Round-trip representative samples; confirm non-empty where non-empty is expected |
Common Pitfall Categories
Detailed pitfall catalog with examples is in references/common-pitfalls.md. Summary:
- Sign errors in log-space — log-softmax, KL divergence, DPO log-ratio
- Numerical stability constants out of range — additive epsilons, clip bounds, temperature
- String processing in decoding — cleanup that blanks shapes the rules didn't anticipate
- Reward / decoding format mismatch — reward function sees text with the shape it needs removed
- Reference model drift — reference model not frozen, or gradients flowing through it
- Gradient flow breakage — detached tensors, in-place ops,
.item()/numpy round-trips in the loss path - Configuration misuse — SFT-scale LR, KL coefficient dominating, clip range too tight,
Gtoo small
Available References
| File | Contents | When to load |
|---|---|---|
references/common-pitfalls.md | Catalog of pitfall categories with symptoms and detection strategies | When you have a suspicious area and want to match it against known pattern shapes |
references/diagnostic-workflow.md | Stage-by-stage diagnostic procedures with verification snippets | When you need guidance on which invariant to check in which order |
scripts/verify_pipeline.py | Runnable diagnostic that prints log-prob / advantage / decoding values on small fixed inputs so you can compare them against the invariants above | Before declaring a fix done: run this and inspect the output, paying attention to lines marked ? |
