★ Part 4 checkpoint · covers Lecture 21

HW4 · LLM RL

Modern LLM post-training, built by hand. You implement the algorithmic core of fine-tuning a base language model with verifiable outcome rewards — per-token log-probabilities, completion masking, group-relative advantages, an approximate KL penalty — then run controlled experiments comparing GR-REINFORCE against GRPO on a debugging task and on hard MATH problems.

Starter code on GitHub ↗ Handout (PDF) ↗

What you build

This assignment ports the course’s policy-gradient machinery onto language models. A prompt $x$ gets a sampled completion $y = (y_1, \ldots, y_T)$ from the trainable policy $\pi_\theta$ ; a frozen reference policy $\pi_{\text{ref}}$ anchors a KL penalty; the reward $r(x, y)$ is verifiable — computed automatically from the completion text. For each prompt you sample a group of $G$ completions and normalize rewards within the group, GRPO-style:

$A_{i,j} = \frac{r_{i,j} - \mu_i}{\sigma_i + \varepsilon}$

one scalar advantage per completion, positive when it beats its group-mates. That group-relative baseline powers both algorithms you compare. GR-REINFORCE weights each completion’s sequence-averaged log-probability by $A_i$ and takes exactly one on-policy pass per rollout batch. GRPO reuses each rollout for two PPO epochs, with a per-token importance ratio $\rho_{i,t}(\theta) = \pi_\theta / \pi_{\text{old}}$ inside a clipped objective $\min\!\left(\rho A,\ \text{clip}(\rho, 1-\epsilon, 1+\epsilon) A\right)$ . Both add $\beta\, \widehat{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})$ using the sampled-token estimator $\hat{k}(a) = e^{\Delta(a)} - \Delta(a) - 1$ . Section 2 of the handout derives why this estimator is unbiased for the KL yet nonnegative per token — read it carefully, because report question 1 asks you to explain it in your own words.

Two tasks carry the experiments. format_copy — copy an integer into an XML tag — is intentionally trivial, a fast pipeline validator whose reward decomposes into correctness, tag presence, and clean-XML components. math_hard is the real task: a filtered level-5 subset of the MATH dataset, a deterministic 512-example held-out eval, and a reward built around extracting a final $\boxed{\text{NUMBER}}$ answer, with small shaping components the handout notes are there to help exploration.

Which lectures feed it

Lecture 21 (RL with Sequence Models and LLMs) is the direct source: rewatch it for why outcome-reward RL on language models is just a policy gradient over token sequences, and where the reference-policy KL comes from. The underlying machinery is older — Lecture 5 gives you the log-gradient identity and why baselines (here, the group mean) cut variance without bias, and Lecture 9 explains why the clipped importance ratio is the bounded-error way to reuse off-policy samples. If the clipping objective feels arbitrary, that lecture is the one to revisit.

Order of attack

Read the code before writing any — the handout prescribes an order (train script, log-probs, rollout buffer, the two algorithms, the task files) that mirrors data flow through one training iteration. Then implement the TODOs in the stated order: log-probs, completion mask, and the KL estimator first; minibatching; group advantages and normalization; then the REINFORCE update; then GRPO. The order matters — even the GR-REINFORCE baseline depends on group advantages, so build those before debugging either algorithm.

The implementation checklist flags where bugs hide: logits at position $t-1$ score target token $t$ ; only completion tokens are scored, never prompt or padding; rewards reshape to groups before normalization; minibatch slicing must keep every tensor field aligned. Debug everything on format_copy — its runs finish in about five minutes and the handout says to run those commands first for exactly this reason.

The four required runs

Run the four commands from Section 7 unmodified — the Gradescope autograder expects them. Everything runs detached on Modal with an H100; you can close the terminal and watch WandB. Budget: the format-copy runs take about five minutes each, math with GR-REINFORCE about two hours, math with GRPO about five — roughly eight H100 hours total, and the handout asks you to stay under a third of your Modal credits (about 40 hours). Start early.

The handout gives concrete debugging thresholds: format_copy reward typically jumps sharply around step 10 and reaches about 1.3, with eval exact match approaching 1.0 — materially below 0.8 means debug. On math_hard, the pre-RL baseline evals around 0.23–0.26; GR-REINFORCE should land near 0.31 final eval, GRPO above 0.37 with train reward above 0.55. Watch the two eval exact-match metrics and the mean rollout reward first, then KL, clip fraction, and gradient norm; rollout reward fluctuates a lot, so compare algorithms on the smoother eval curves. The comparison is deliberately controlled — identical batch, minibatch, and learning-rate settings, differing mainly in GRPO’s two epochs of sample reuse — which is why GRPO should be more sample-efficient over the first 200 iterations. Also open WandB’s sampled-generations panel: reading actual completions with their rewards and advantages is often the fastest sanity check.

Ablations, report, submission

Beyond the required runs, do a GRPO hyperparameter study on format_copy: at least five extra runs varying ppo_epochs, minibatch_size together with grad_accum_steps, kl_coef (smaller and larger), and clip_eps, pushing values until performance degrades. The report answers five questions — the KL estimator, your implementation order and one bug, the math-run comparison, the ablation summary, and one or two informative generations. Submission is two Gradescope uploads: a report PDF, and a code zip built by bundling run artifacts on Modal, then zipping your hw4 code folder next to the unzipped bundle with no wrapper directory. Partial run submissions are accepted and scored for what’s present.

Back of the chalk tray — what to remember

You can implement the core of LLM RL post-training from scratch: masked per-token log-probs, group-relative advantages, a sampled-token KL estimator, and both on-policy and clipped sample-reuse updates.
You know why $e^{\Delta} - \Delta - 1$ is the right per-token KL surrogate — unbiased in expectation, nonnegative per sample — and can defend it in a paragraph.
You can read a controlled experiment: same data budget, one optimizer pass versus two clipped epochs, and interpret the sample-efficiency gap between GR-REINFORCE and GRPO.
You can diagnose an RL fine-tuning run from WandB — eval exact match, rollout reward, KL, clip fraction — and know which curves are trustworthy and which just fluctuate.
You’ve mapped how hyperparameters like KL coefficient, clipping threshold, and epochs move a run from stable to degenerate, with evidence from your own ablations.