Assignment 5 · Alignment & Reasoning RL

What you build

The first four assignments built a base model; this one turns a base model into something that solves problems. The testbed is fixed: OLMo-2-0425-1B (trained on 4T tokens, so unusually capable for its size) on GSM8K grade-school math word problems. The work comes in four waves.

Prompting baselines. Evaluate the base model zero-shot with a bare question_only prompt, with the DeepSeek r1_zero chain-of-thought prompt (the model reasons inside think tags, answers inside answer tags), and few-shot with a 3-shot variant. You categorize failures by format reward versus answer reward and characterize how each prompt shapes the model’s behavior.

On-policy GRPO. The core of the assignment. You derive the algorithm step by step in the handout — REINFORCE, then a group-mean baseline, then std normalization, then sequence-length normalization — and implement it as composable pieces: prompt/response tokenization with a response mask, per-token log-probs and entropies, rollout rewards, group-normalized advantages, the per-token policy gradient loss, loss aggregation, and finally a full train step with gradient accumulation and grad-norm clipping. Then you write the training loop and run it: with the suggested hyperparameters (256 rollouts per batch, group size 8, 200 steps), validation accuracy should climb from single digits to at least 25%, averaged over 4 random seeds.

Variants and ablations. GRPO’s design choices are contested, and you test them empirically: constant versus sequence loss normalization, Dr. GRPO (no std normalization), rejection fine-tuning (keep only correct rollouts, do SFT on them), and MaxRL (normalize by the group mean). Written problems have you derive what each normalizer implicitly does — each one reweights prompts by difficulty in a different way.

Off-policy RL. Take multiple gradient steps per inference batch (32x off-policy) and watch what breaks. You implement importance reweighting with no clipping, PPO/GRPO-style token-level clipping, and GSPO’s sequence-level geometric mean, then compare stability and clip fractions across seeds. A final open-ended problem asks you to propose and test your own estimator — change exactly one thing and defend it.

Which lectures feed it

This assignment leans on the post-training arc of the course (Lectures 15–17): what alignment is and why pretraining alone doesn’t give you a task-solver, the policy gradient / RLHF-to-RLVR lineage, and GRPO and its modern descendants. When the derivations in Section 4 feel dense, the RL lectures are the place to rewatch — the handout’s math (baselines preserve expectations, normalization changes the objective, importance weights correct staleness) mirrors the lecture presentation closely, and it also recommends OpenAI’s Spinning Up and Lambert’s RLHF book for anything still shaky. Earlier course material comes back too: mixed precision (the model loads in bfloat16), gradient accumulation from the systems lectures, and the inference-is-different lesson — generation runs on a dedicated vLLM server on a second GPU with explicit weight syncing, exactly the memory-bound/compute-bound split from Lecture 2.

Order of attack

The handout’s order is the right order, because the pieces compose. Do the prompting section first — it forces you to get vLLM generation, the r1_zero prompt, stop strings, and the provided drgrpo_grader reward functions working before any training exists. Then build the GRPO components bottom-up in the order the tests are listed: tokenization, log-probs, rewards, group normalization, loss, aggregation, train step. Each has a small adapter hook and a pytest target, so you get tight feedback loops before committing GPU hours. Only then write the training loop, run it for about 50 steps to sanity-check that validation reward moves, and scale to the full 4-seed runs. Do the written derivation problems (baseline_calcs, the difficulty-reweighting and surrogate-objective derivations) before the corresponding experiments — they tell you what to expect in the plots.

Budget matters here: the handout prices the big experiments explicitly — about 2 B200-hours for the main GRPO runs, 4 each for the learning-rate sweep and prompt ablation, and 8 each for the on-policy variants and off-policy comparisons. Note the reward often sits near zero for the first few steps before takeoff; don’t kill a run early.

What’s hard

The hard part is not any single function — most components are worth 0.5–1 point and a few dozen lines — it’s that RL is slow, unstable, and high-variance across seeds, and the handout says so up front. Small implementation details matter disproportionately: using torch.std’s default rather than a hand-rolled formula, adding advantage_eps before dividing, getting the response mask aligned with shifted labels, and making microbatch gradient accumulation exactly equivalent to the full-batch gradient. Logging is half the assignment — loss, grad norm, token entropy, train and validation rewards, response lengths, clip fractions, and periodic rollout dumps are what let you tell a bug from bad luck. And because every experiment needs 4 seeds with variance shown, sloppy experiment management gets expensive fast.

Logistics

Clone the repo; starter code is minimal by design — vllm_utils.py (vLLM server with weight sync), drgrpo_grader.py (reward functions), prompt files, and the test suite. Required tests live in tests/test_grpo.py and connect to your code through tests/adapters.py; run them with uv run pytest -k <test_name> as you go. There is no leaderboard for this one — the headline empirical bar is the 25% validation accuracy deliverable. Submit writeup.pdf (typeset answers to all written problems, with plots) and code.zip via Gradescope, using test_and_make_submission.sh. An optional supplement covers instruction tuning/safety with its own tests.