Alignment and Reasoning RL

Lectures behind it: L15 Alignment: SFT and RLHF, L16 Alignment: RL 1, L17 Alignment: RL 2 · supporting: L10 Inference (vLLM, paged KV cache), L12 Evaluation (MATH, MMLU, AlpacaEval, judges)

A5 is where the course stops building the model and starts building the behaviour. Every prior assignment optimized cross-entropy against a fixed corpus; here cross-entropy is explicitly abandoned as the objective, because the thing you want — a model that solves competition math problems and formats the answer where a grader can find it — is not a likelihood. You get a downstream metric instead, and you optimize it directly. The course also draws a hard boundary around what counts as yours: you may call vLLM to generate and HuggingFace to run forward passes, but the RL is yours. No Trainer, no TRL, no PPO library. That constraint is the point — GRPO is about two hundred lines once you strip the framework, and the assignment exists so you find that out by writing them.

We expect you to build most of the RL related components from scratch. You may use tools like vLLM to generate text from language models (§3.1). In addition, you may use HuggingFace Transformers to load the Qwen 2.5 Math 1.5B model and tokenizer and run forward passes (§4.1), but you may not use any of the training utilities (e.g., the Trainer class).§1, handout

Map of the assignment

Twenty-six graded problems across nine sections, 66 points by my count (the handout never prints a total, and grpo_off_policy carries no point value of its own — it is an implementation step the sweep that follows depends on). Eleven of them are unit-tested; the other fifteen are scripts, runs and write-ups. Only two test files are mandatory.

You are only expected to pass the tests in tests/test_sft.py and tests/test_grpo.py—the rest of the tests are for the non-mandatory parts of the assignment.§1, on what tests/*.py contains

The optional supplement adds sixteen more problems and 61 more points on a different model and a different set of benchmarks; it has its own map further down. The two halves share only tests/adapters.py.

Setup: environment, data, tests

Environment. Same uv pattern as the rest of the course, with one wrinkle the README calls out in so many words: FlashAttention-2 is weird. It needs torch present at build time, so you install everything else first and then install it. Run uv sync --no-install-package flash-attn, then uv sync. The pyproject.toml pins vllm==0.7.2 and flash-attn==2.7.4.post1 exactly, sets no-build-isolation-package = ["flash-attn"], and pulls alpaca-eval from a course fork on a branch called forward_kwargs_to_vllm — do not swap in the PyPI release, it cannot pass sampling kwargs through to a local vLLM annotator.

Models and data. Nothing is downloaded. Everything lives on the Together cluster and the handout asks you explicitly not to re-download it — a 70B judge model replicated across a class is a lot of disk.

MATH itself is Hendrycks et al. 2021 and is not publicly redistributable — a copyright claim pulled it. If you are following along outside Stanford the handout names four substitutes: Countdown (the synthetic arithmetic task TinyZero made popular), GSM8K (easier, best for debugging the pipeline), Tulu 3 SFT Math (synthetic, some labels are wrong), and DART-Math's comparison table of math SFT sets. If your substitute has long-form solutions rather than short labels, run the answer column through Math-Verify to get something a string comparison can grade.

The test harness. Same adapter pattern as A1–A4: tests/adapters.py holds sixteen run_* stubs that all raise NotImplementedError, and the tests import those, never your code. You write your implementation wherever you like inside cs336_alignment/ and wire it up in one line per adapter. Run everything with uv run pytest, or one problem at a time with uv run pytest -k test_masked_mean.

Almost every A5 test is a snapshot test, which is different from A1's reference-implementation tests and changes how you debug. conftest.py defines a numpy_snapshot fixture that loads a stored .npz from tests/_snapshots/ named after the test function and compares your returned arrays to it at rtol=1e-4, atol=1e-2 (pass --snapshot-exact to demand bit equality). Two consequences worth internalizing before you start. First, the tolerance is loose enough that a subtly wrong reduction can still pass — check shapes yourself. Second, the snapshot compares a dict of named arrays, so your return keys must match the reference exactly: "log_probs" and "token_entropy", not "logprobs".

Three fixtures are worth knowing about before the first failure. The tokenizer and model fixtures both load from the hard-coded cluster path /data/a5-alignment/models/Qwen2.5-Math-1.5B — so test_tokenize_prompt_and_output and test_get_response_log_probs simply cannot run off-cluster without editing conftest. The reward_fn fixture is not the real grader: it is a deterministic SHA-256 hash of the response string mapped into {0.0, 0.1, …, 0.9}, chosen so the group-normalization test has spread in its rewards without needing sympy. And group_size is 4 over 8 rollouts, so the group test really does exercise two groups rather than degenerating to one.

2 · Reasoning with language models

Two things change relative to every prior assignment, and the handout is unusually explicit about both. You stop using your own model, because a 100M-parameter transformer trained on TinyStories has no mathematical reasoning to elicit — you cannot RL your way to competition math from a model that never had it. And you stop using cross-entropy as the yardstick, because the whole premise now is the gap between likelihood and downstream behaviour. Qwen 2.5 Math 1.5B Base (Yang et al. 2024) was continually pretrained on synthetic math, which is exactly why it has latent ability that post-training can surface.

The intellectual arc packed into §2.2 runs: scratchpads (Nye et al. 2021) showed a model can be finetuned to externalize intermediate steps; chain-of-thought prompting (Wei et al. 2023) showed you can get the same for free from a big enough model; STaR turned it into a bootstrap loop — sample chains, keep the ones that reach the right answer, finetune on those, repeat — which is expert iteration with a string-match verifier standing in for a search oracle; and then o1, R1 and Kimi k1.5 replaced the filter-and-finetune step with an actual policy gradient. You will implement that whole ladder in order, which is the best argument for doing §5 before §7 even though GRPO is the headline.

The R1-Zero prompt. Everything in the required assignment uses the prompt in cs336_alignment/prompts/r1_zero.prompt — copy it from the file, not from the PDF, which is exactly why the file exists. It instructs the assistant to think inside <think> </think> and answer inside <answer> </answer>, and it pre-fills the opening <think> so the model starts mid-reasoning. Two tags, two jobs: <answer> gives the grader something to parse, and </answer> gives vLLM a stop string, which is what keeps rollouts from running to 1024 tokens every time.

The handout then undercuts its own choice, and this is worth reading carefully because it sets up the last ablation. Liu et al. 2025 (Dr. GRPO) found that prompting Qwen 2.5 Math with nothing but the bare question starts far higher and takes the R1-Zero prompt 100+ RL steps just to catch up — evidence the model was pretrained on question-answer pairs in roughly that shape. The staff pick R1-Zero anyway, on the grounds that it improves visibly fast and therefore makes a correct implementation obvious within a handful of steps. You verify the claim yourself in grpo_prompt_ablation.

3 · Measuring zero-shot MATH performance

3.1 · vLLM for offline inference. RL is generation-bound: every GRPO step samples hundreds of complete solutions before a single gradient flows. Writing your own sampler (as you did in A1) is fine for a demo and hopeless here, so the assignment hands you vLLM — continuous batching, paged KV cache (Kwon et al. 2023), the whole machine L10 takes apart. The starter snippet in the handout is lifted from vLLM's own offline-inference example: build SamplingParams, build LLM(model=...), call llm.generate(prompts, sampling_params), read output.outputs[0].text.

That last row is the single most common day-one bug in this assignment. stop strings are excluded from vLLM output by default; r1_zero_reward_fn requires "</answer>" in response. Forget the flag and you get a uniformly zero baseline and conclude the model is broken.

3.2 · The reward function. Math answers cannot be string-compared: the gold label 0.5 and the model's 1/2 are the same answer. The course hands you a parser rather than making you write one — cs336_alignment/drgrpo_grader.py, about a thousand lines lifted from Dr. GRPO, with LaTeX normalization, sympy equivalence and a fast path. The entry point is r1_zero_reward_fn(response, ground_truth), which returns a dict with "format_reward", "answer_reward" and "reward". Read its first line before you run anything: the format gate is the literal substring "</think> <answer>", one space, plus "</answer>" somewhere after. Also note the asymmetry that the write-up asks you to explain — a formatted-but-wrong answer gets format_reward 1.0 and total reward 0.0, deliberately, so a model cannot farm reward by emitting well-formed nonsense.

Problem (math_baseline): 4 points

Deliverable: an evaluation script, plus a three-part write-up on where the base model and the parser each fail.

Part (a) is the script: load validation.jsonl, format each example with the R1-Zero prompt, generate, score, and — the part people skip and regret — serialize examples, generations and per-example scores to disk. Parts (b) and (c) require you to go back and read those generations, and re-running 5,000 rollouts because you only kept the aggregate is a waste of an hour. The handout suggests factoring the work into evaluate_vllm(vllm_model, reward_fn, prompts, eval_sampling_params), and it is worth taking that suggestion literally: you will call this function again in sft_experiment, in the EI loop and inside the GRPO train loop's validation pass.

Part (b) is the real content. Bucket the 5,000 generations three ways — format 1 / answer 1, format 1 / answer 0, format 0 / answer 0 — then read at least ten from each of the failure buckets and argue whether the fault is the model's or the parser's. Both happen. Format-0 is usually the base model genuinely never closing </think>, or drifting into a fresh User: turn, or hitting the 1024-token cap mid-derivation; but it is also sometimes the parser's exact-space requirement rejecting </think><answer> or </think>\n<answer>. Format-1/answer-0 splits between honest arithmetic errors and equivalence the grader misses — an interval written differently, an unsimplified radical, a unit tacked on. The point of the exercise is calibration: your entire RL signal is this function, and you should know its failure modes before you spend 50 H100-hours optimizing against it.

• cs336_alignment/drgrpo_grader.py — r1_zero_reward_fn(response, ground_truth, fast=True) and question_only_reward_fn, verified against the file
• cs336_alignment/prompts/r1_zero.prompt — the prompt template, three lines, ends with Assistant: <think>
• No unit test. This problem is graded entirely on the write-up and the script you submit in code.zip.

4 · Supervised finetuning for MATH

Algorithm 1 in the handout is four lines and holds no surprises: sample a batch of (question, response) pairs, cross-entropy on the response, step. What makes it interesting is what the responses are. You are not finetuning the model to emit correct answers; you are finetuning it to emit a chain of thought and then an answer, using traces distilled from DeepSeek R1 in MATH/sft.jsonl. Distillation of a reasoning process, not of an answer key.

The handout also explains why SFT and RL are treated as separate phases here rather than composed as they are in practice: at 1.5B parameters the effects do not stack visibly, so chaining them would only add compute and confusion. Two motivations for the ordering survive anyway — SFT needs annotated traces where RL needs only an answer key, and RL can find policies better than anything in the SFT set.

4.1 · HuggingFace mechanics. Load with torch_dtype=torch.bfloat16 and attn_implementation="flash_attention_2", forward with model(input_ids).logits, save with save_pretrained(save_directory=output_dir) under /data/yourusername (checkpoints are large), and save the tokenizer alongside so the directory is self-contained. The section closes with a page on gradient accumulation — divide the loss by gradient_accumulation_steps before backward(), call optimizer.step() and zero_grad() every k microbatches — which matters because both *_microbatch_train_step problems are graded on getting exactly this scaling right.

Problem (tokenize_prompt_and_output): Prompt and output tokenization · 2 points

Deliverable: tokenize prompt and response separately, concatenate, pad, and build a response mask.

This is the boring function that everything else silently depends on, and it is where most of the assignment's off-by-one bugs live. You tokenize prompt and output separately — not the concatenated string — because you need to know where the boundary is; a joint tokenization can merge the last prompt token with the first response token and you lose the split. Then you concatenate, pad to the batch max, and produce three tensors of width max(prompt_and_output_lens) - 1: input_ids is the concatenation with the final token dropped, labels is the same shifted left by one, and response_mask is 1 exactly on the response positions of the labels. That last clause is the trap. The mask indexes labels, not input_ids, so it is shifted too; mask the input positions instead and every loss you compute for the rest of the assignment is off by one token, which is subtle enough to pass a smoke test and wrong enough to ruin a training run. Padding must be masked to 0 as well.

• tests/adapters.py · run_tokenize_prompt_and_output — (prompt_strs: list[str], output_strs: list[str], tokenizer: PreTrainedTokenizerBase) -> dict[str, Tensor] with keys input_ids, labels, response_mask
• tests/test_sft.py — test_tokenize_prompt_and_output snapshots all three tensors against three short prompt/output pairs; needs the Qwen tokenizer from the cluster path

Problem (compute_entropy): Per-token entropy · 1 point

Deliverable: entropy of the next-token distribution at every position, from logits.

One line of math — H(p) = −Σ p(x) log p(x) over the vocabulary — and one instruction that carries the point: do it in a numerically stable way with logsumexp. Naively exponentiating raw logits over a 150K-token vocabulary in bfloat16 overflows. The stable form computes logZ = logsumexp(logits, dim=-1), then H = logZ - sum(softmax(logits) * logits, dim=-1), which never materializes an unnormalized exponential. Input is (batch, seq, vocab), output is (batch, seq) — you reduce the vocabulary axis only. Entropy is not part of any loss; it is the diagnostic you will stare at for the rest of the assignment, because a collapsing entropy curve is how over-confident RL announces itself before the reward curve does.

• tests/adapters.py · run_compute_entropy — (logits: torch.Tensor) -> torch.Tensor
• tests/test_sft.py — test_compute_entropy snapshots the output for a seeded (2, 10, 100) random logit tensor; pure CPU, runs anywhere

Problem (get_response_log_probs): Response log-probs (and entropy) · 2 points

Deliverable: per-token conditional log-probabilities from a causal LM, optionally with per-token entropy.

The workhorse. Run the model, take logits, log-softmax over the vocabulary, and gather the entry corresponding to each label — log p(x_t | x_<t) at every position. Return a dict with key "log_probs", plus "token_entropy" when return_token_entropy=True, both shaped (batch, seq). Two notes. The adapter's type annotation says -> torch.Tensor and is simply wrong — the docstring and the test both want a dict; trust the docstring. And read the second half of the docstring carefully: nothing is masked here. Prompt and padding positions come back with real numbers in them, and it is the training loop's job to mask them. People who helpfully zero out the prompt inside this function fail the snapshot and then break the entropy statistics later.

Use torch.gather on the log-softmax rather than indexing with fancy tensors, and be aware that this is the memory peak of a training step: log-softmax over (batch, seq, 150K) in fp32 is larger than the model. When you need old log-probs for off-policy GRPO you call the same function under torch.inference_mode().

• tests/adapters.py · run_get_response_log_probs — (model, input_ids, labels, return_token_entropy: bool), returns dict[str, torch.Tensor]
• tests/test_sft.py — test_get_response_log_probs calls it with return_token_entropy=True and snapshots both keys; loads the real Qwen model, so cluster-only

Problem (masked_normalize): Masked normalize · 1 point

Deliverable: sum a tensor over a dimension, counting only masked-in elements, then divide by a constant.

Six lines: (tensor * mask).sum(dim) / normalize_constant, with dim=None meaning sum everything. The reason it is a graded problem at all is that it is one half of the length-normalization debate you will run as an experiment in §8 — this is the "sum and divide by a fixed constant" branch, against masked_mean's "divide by the actual token count" branch. Choosing between them changes which tokens get gradient weight, and think_about_length_normalization asks you to predict the difference before measuring it.

Verify the argument order against the adapter, not the PDF. The handout's recommended signature is (tensor, mask, normalize_constant, dim=None); the actual adapter is (tensor, mask, dim=None, normalize_constant=1.0). The handout's later worked example calls it with a keyword constant_normalizer=, which does not exist under either spelling. Keyword-call it as the adapter declares it and the discrepancy is harmless.

• tests/adapters.py · run_masked_normalize — (tensor, mask, dim: int | None = None, normalize_constant: float = 1.0) -> torch.Tensor
• tests/test_sft.py — four snapshots: test_masked_normalize_dim0, _dim1, _dimlast, _dimNone, all with normalize_constant=10.0 and a random boolean mask

Problem (sft_microbatch_train_step): Microbatch train step · 3 points

Deliverable: one forward-and-backward on a microbatch — masked NLL, normalized, scaled for gradient accumulation, backpropagated.

The three points are for getting the scaling right, because everything else is arithmetic. Negate the log-probs, sum them over response tokens only via masked_normalize, divide by gradient_accumulation_steps, call loss.backward() inside the function, and return the scaled scalar plus a metadata dict. Two things people get wrong. They forget the division, which multiplies the effective learning rate by the accumulation count and produces a run that diverges for no visible reason. And they call zero_grad() here — this function must not; accumulation across microbatches is the entire point, and the ten-step test exists precisely to catch it.

That test is worth understanding: test_sft_microbatch_train_step_10_steps calls the step ten times on the same policy_log_probs tensor without ever zeroing, stacks the ten gradients, and snapshots the stack. If your gradients are being cleared, or scaled inconsistently, the stack is wrong even when a single step passes.

• tests/adapters.py · run_sft_microbatch_train_step — (policy_log_probs, response_mask, gradient_accumulation_steps: int, normalize_constant: int | None = 1.0) -> tuple[torch.Tensor, dict[str, torch.Tensor]]
• tests/test_sft.py — test_sft_microbatch_train_step snapshots {loss, policy_log_probs_grad}; test_sft_microbatch_train_step_10_steps snapshots the accumulated stack. Fixtures: gradient_accumulation_steps=2, normalize_constant=10.0

Problem (log_generations): Logging generations · 1 point

Deliverable: a function that samples from the current policy on validation prompts and logs what came out.

The cheapest point in the assignment and the one that saves you the most time later. The handout lists six things to log per example: prompt, generated response, ground truth, the reward breakdown (format / answer / total), the average token entropy of the response, and average response length split three ways — overall, correct-only, incorrect-only. That last split is the diagnostic that catches length hacking, the classic RL-on-reasoning failure where the model discovers that rambling correlates with reward and response length climbs while accuracy does not. Build this before sft_experiment, not after, and call it from every loop you write afterwards.

• No adapter, no test — graded on the code you submit and on it showing up in your experiment logs.
• Log to wandb (already a dependency); §4.3 gives the define_metric incantation for separate train and eval x-axes.

Problem (sft_experiment): Run SFT on the MATH dataset · 2 points · 2 H100 hrs

Deliverable: validation-accuracy curves across five SFT dataset sizes, plus a second run on a correctness-filtered dataset with its size reported.

Part 1: run SFT on {128, 256, 512, 1024} unique examples and on the full set, tuning learning rate and batch size until the full-dataset run clears 15% validation accuracy. Part 2: filter sft.jsonl down to only those examples whose response actually produces the correct answer under r1_zero_reward_fn, train on that, report both the filtered size and the accuracy, and compare. The expected finding is the one STaR is built on — a smaller, verified dataset beats a larger, noisy one — and it is the direct motivation for §5, where you generate the data instead of being given it.

Two GPUs: one holds the policy, one holds a vLLM instance for periodic validation. The handout gives you the plumbing, and you should copy it rather than reinvent it. init_vllm() monkeypatches two things borrowed from TRL — torch.distributed.get_world_size forced to 1 so vLLM will sit on a chosen device, and a memory-profiling assertion disabled because it does not hold in this setup. load_policy_into_vllm_instance() pushes policy.state_dict() straight into llm.llm_engine.model_executor.driver_worker.model_runner.model before every evaluation, so vLLM scores the current weights rather than the checkpoint it booted with. Forgetting that call gives you a validation curve that is flat at the base model's accuracy forever — a genuinely confusing failure, since training loss looks fine. Gradient clipping at 1.0 throughout.

• tests/adapters.py — no adapter; this composes the six functions above
• L15 Alignment: SFT and RLHF — the SFT half of the lecture is this problem's conceptual half

5 · Expert iteration for MATH

The bridge between SFT and RL, and the cheapest way to see that a verifier is worth more than a teacher. Algorithm 2 wraps SFT in an outer loop: freeze a copy of the policy, sample G solutions per question from it, score them, throw away everything with reward 0, finetune on what survives, repeat. It is STaR exactly, and the same idea appears as Cobbe et al. 2021, Language Model Cascades and ReST in the handout's citation list. L16 calls it the RFT baseline in the DeepSeekMath comparison plot, which is the right frame: expert iteration is the strongest thing you can do without a policy gradient, and GRPO has to beat it to justify itself.

Problem (expert_iteration_experiment): Run expert iteration on the MATH dataset · 2 points · 6 H100 hrs

Deliverable: validation curves for at least two rollout counts × two epoch counts, a model above 15% validation accuracy, a two-sentence comparison against SFT and across EI steps, and an entropy-over-training plot.

Run on MATH/train.jsonl with n_ei_steps = 5, sweeping G (rollouts per question) and the number of SFT epochs per EI step. Both knobs trade the same currency. Larger G means more rollouts per question and therefore more surviving correct traces from hard questions — the coverage that makes EI work at all — at linear generation cost. More epochs per step squeezes harder on the traces you already have, at the cost of overfitting the model to its own outputs. The entropy plot is the requested deliverable because it shows the mechanism directly: EI is self-distillation, so entropy falls, and the interesting question is whether it falls faster than accuracy rises. Compare to the entropy curve you get from GRPO later — the handout explicitly asks you to in grpo_off_policy_sweep.

Two operational notes. Pass min_tokens=4 to SamplingParams; without it vLLM will occasionally return an empty string, which becomes an empty response mask, which becomes a division by zero, which becomes a NaN that quietly poisons the run. And keep the stop=["</answer>"] setup from §3 — the handout repeats the instruction here because people drop it when they switch from the eval path to the rollout path. Gradient clipping at 1.0 again.

• STaR: Bootstrapping Reasoning With Reasoning — Zelikman et al. (2022) · the algorithm you are implementing
• Thinking Fast and Slow with Deep Learning and Tree Search — Anthony et al. (2017) · where "expert iteration" gets its name

6 · Primer on policy gradients

No problems in this section — it is four pages of derivation you need before §7 makes sense, based on Spinning Up and Nathan Lambert's RLHF Book. The framing: a causal LM is a categorical stochastic policy, with the text prefix as state and the next token as action; the environment is deterministic (the next state is the old prefix plus the emitted token) and rewards are terminal only — 1 if the finished solution grades correct, 0 otherwise, nothing in between.

From there: the REINFORCE identity ∇J = E[Σ_t ∇ log π(a_t|s_t) · R(τ)], derived in six lines from the log-derivative trick plus the observation that ρ₀, P and R do not depend on θ. Then baselines — subtract any b(s_t) that depends only on state and the estimator stays unbiased, because the expected score function is zero, while the variance drops. Then off-policy: REINFORCE demands fresh rollouts for every single gradient step, which is absurd when generation dominates the cost, so you importance-weight by π_θ(a|s)/π_θold(a|s) and accept the approximation, which is fine as long as the two policies stay close (Degris et al. 2013).

One paragraph in §6.5 is worth memorizing, because it explains why every deliverable in §8 asks for reward curves and not loss curves:

pg_loss is not a loss in the canonical sense—it's not meaningful to report pg_loss on the train or validation set as an evaluation metric, and a good validation pg_loss doesn't indicate that our model is generalizing well. The pg_loss is really just some scalar such that when we call pg_loss.backward(), the gradients we obtain through backprop are the approximate policy gradient.§6.5, "A note on policy gradient 'losses'"

7 · Group Relative Policy Optimization

GRPO's one idea: get your baseline for free by sampling a group. Sample G solutions to the same question, and the mean reward across that group is an estimate of V(s) for that prompt — no learned value network to train, no second set of weights resident on the GPU, no critic that can be wrong. Advantage A⁽ⁱ⁾ = (r⁽ⁱ⁾ − mean(r)) / (std(r) + advantage_eps), per DeepSeekMath and R1. The advantage is constant across every token of a response, which is why it arrives shaped (batch_size, 1) and broadcasts along the sequence.

Algorithm 3 is the loop: for each GRPO step, sample a batch of questions, snapshot π_old ← π, sample G outputs per question, reward them, group-normalize, then take n_train_steps_per_rollout_batch gradient steps on the GRPO-Clip objective. The handout notes in a footnote that this is a special case of DeepSeekMath's GRPO with a verified reward function, no KL term, and no iterative reference or reward model update — the KL term is dropped because in the staff's experiments and in the literature it made no difference to performance while costing a resident reference model's worth of memory.

The clipping is PPO's, and the handout's case analysis is the clearest way to hold it. When A > 0, the objective is min(ratio, 1+ε)·A: you are rewarded for making the token more likely, but the reward saturates once the ratio exceeds 1+ε, so there is no gradient incentive to run away from π_old. When A < 0, symmetric: you push the token down but stop caring below 1−ε. Clipping only matters when π_θ has had a chance to drift from π_old — that is, when you take more than one gradient step per rollout batch.

Problem (compute_group_normalized_rewards): Group normalization · 2 points

Deliverable: score a flat list of rollouts, reshape into groups, normalize within each group, return normalized rewards, raw rewards and metadata.

The bookkeeping is the whole problem. rollout_responses arrives flat, of length rollout_batch_size = n_prompts_per_rollout_batch × group_size, with repeated_ground_truths the same length because each question's answer is repeated group_size times. Reshape to (n_prompts, group_size), subtract the per-group mean, and — if normalize_by_std — divide by the per-group std plus advantage_eps. Flatten back. Get the reshape order wrong (interleaving groups instead of blocking them) and you will normalize across questions instead of within them, which produces a plausible-looking tensor and an advantage signal that is pure noise.

The normalize_by_std flag is the Dr. GRPO ablation baked into the API: Liu et al. 2025 argue that dividing by the group std up-weights questions where the group barely disagreed — the too-easy and too-hard ones, where std is near zero — so they drop the division entirely and use A = r − mean(r). You implement both and test the difference in grpo_group_standard_deviation.

One repo wart to know: the adapter's return annotation reads tuple[torch.Tensor, dict[str, float]] but the docstring and both tests unpack three values — (advantages, raw_rewards, metadata). The annotation is stale; follow the docstring.

• tests/adapters.py · run_compute_group_normalized_rewards — (reward_fn, rollout_responses, repeated_ground_truths, group_size, advantage_eps, normalize_by_std), returns (advantages, raw_rewards, metadata)
• tests/test_grpo.py — test_compute_group_normalized_rewards_normalize_by_std and ..._no_normalize_by_std; 8 rollouts, group_size=4, advantage_eps=1e-6, and a hash-based dummy reward function so the test needs no sympy

Problem (compute_naive_policy_gradient_loss): Naive policy gradient · 1 point

Deliverable: the per-token loss −A_t · log p_θ(o_t | q, o_<t).

A multiply and a negate: -raw_rewards_or_advantages * policy_log_probs, with the (batch, 1) scalar broadcasting over (batch, seq). The single argument name doing double duty is deliberate — pass raw rewards and you get the no_baseline variant, pass group-normalized advantages and you get reinforce_with_baseline. No reduction happens here; you return the full per-token tensor and someone upstream decides how to sum it. The only way to lose the point is to get the sign backwards, and the symptom is unmistakable: reward marches steadily downward.

• tests/adapters.py · run_compute_naive_policy_gradient_loss — (raw_rewards_or_advantages: torch.Tensor, policy_log_probs: torch.Tensor) -> torch.Tensor
• tests/test_grpo.py — test_compute_naive_policy_gradient_loss, one snapshot on seeded (2, 10) tensors

Problem (compute_grpo_clip_loss): GRPO-Clip loss · 2 points

Deliverable: the per-token clipped surrogate, plus metadata recording which tokens clipped.

Form the ratio in log space — ratio = torch.exp(policy_log_probs - old_log_probs), never a division of probabilities — then -torch.min(ratio * advantages, ratio.clamp(1-cliprange, 1+cliprange) * advantages). Note the negation applies to the whole min, and the min is over the two products, not over the ratios: clipping the ratio first and multiplying afterwards gives a different (wrong) answer whenever the advantage is negative, which is the classic PPO implementation bug and the reason there are two tests here at different clip ranges. The metadata dict should record, per token, whether the clipped branch won — that is your clip fraction, the single best diagnostic for how far off-policy you have drifted.

The two tests, test_compute_grpo_clip_loss_large_cliprange and ..._small_cliprange, exist to catch exactly the sign/order confusion above: with a large clip range the clipping is inert and a broken implementation can still match, so the small-range snapshot is the one that actually discriminates.

• tests/adapters.py · run_compute_grpo_clip_loss — (advantages, policy_log_probs, old_log_probs, cliprange: float) -> tuple[torch.Tensor, dict[str, torch.Tensor]]
• tests/test_grpo.py — two snapshots; old_log_probs is the fixture policy_log_probs + randn_like(...), so the ratio genuinely departs from 1
• Proximal Policy Optimization Algorithms — Schulman et al. (2017) · where the clipped surrogate comes from

Problem (compute_policy_gradient_loss): Policy-gradient wrapper · 1 point

Deliverable: a dispatcher over the three loss types with argument validation.

Trivial to write, and the point is really for the assertions. no_baseline requires raw_rewards and routes to the naive loss; reinforce_with_baseline requires advantages and routes to the same function with a different tensor; grpo_clip requires advantages, old_log_probs and cliprange and routes to the clipped loss. Assert the required arguments are non-None for each branch — the handout says so explicitly, and it pays off in §8 when you are switching loss types from a CLI flag at 2am and would otherwise get a silent None broadcast instead of an error. Merge whatever metadata the underlying routine returns into a single dict so the caller sees one shape regardless of branch.

• tests/adapters.py · run_compute_policy_gradient_loss — (policy_log_probs, loss_type: str, raw_rewards, advantages, old_log_probs, cliprange)
• tests/test_grpo.py — test_compute_policy_gradient_loss_no_baseline, ..._reinforce_with_baseline, ..._grpo_clip; all three pass every argument, so a branch that reads the wrong one still runs and produces the wrong snapshot

Problem (masked_mean): Masked mean · 1 point

Deliverable: mean over masked-in elements along a dimension, or over all of them when dim=None.

(tensor * mask).sum(dim) / mask.sum(dim). The counterpart to masked_normalize: this one divides by the actual number of unmasked tokens, so every sequence contributes equally to the batch loss regardless of length, while masked_normalize divides by a fixed constant so longer sequences contribute proportionally more. The handout flags that the sequence-length division is canonical in RL-on-LLM codebases but is not present in the policy-gradient estimator as derived in §6.4 — there is no 1/T factor in equation 21. That observation is the seed of two problems in §8. Beyond loss reduction you will reuse this constantly for statistics: mean response-token entropy, clip fraction over response tokens, mean response length.

• tests/adapters.py · run_masked_mean — (tensor, mask, dim: int | None = None) -> torch.Tensor
• tests/test_grpo.py — test_masked_mean_dim0, _dim1, _dimlast, _dimNone

Problem (grpo_microbatch_train_step): Microbatch train step · 3 points

Deliverable: per-token loss → masked mean over the sequence → mean over the batch → divide by accumulation steps → backward().

The GRPO twin of sft_microbatch_train_step, and the same three points for the same reason. The reduction order is specified and matters: reduce the sequence dimension with masked_mean(..., dim=-1) to get one scalar per rollout, then average over the batch, then scale. Averaging over everything in one shot is not the same thing when responses have different lengths, and it will pass casual inspection while producing a different gradient. Call loss.backward() inside, do not zero gradients, return the scaled loss plus merged metadata.

The docstring mentions a constant_normalize_factor in its Args list "as in Dr. GRPO" which does not appear in the signature. It is vestigial — the signature is the contract. When you get to grpo_length_normalization you will want a variant of this function that reduces with masked_normalize instead; write it as a separate path in your own code rather than changing the adapter's shape, or the snapshots will stop matching.

• tests/adapters.py · run_grpo_microbatch_train_step — (policy_log_probs, response_mask, gradient_accumulation_steps, loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"], raw_rewards=None, advantages=None, old_log_probs=None, cliprange=None)
• tests/test_grpo.py — test_grpo_microbatch_train_step_grpo_clip and ..._10_steps; both use loss_type="grpo_clip", cliprange=0.1, gradient_accumulation_steps=2, and the ten-step version verifies accumulation the same way the SFT one does

Problem (grpo_train_loop): GRPO train loop · 5 points

Deliverable: a complete train loop, a validation-reward-vs-steps plot showing improvement, and a few example rollouts sampled at different points in training.

Assemble Algorithm 3 from the parts. The handout gives starter hyperparameters that are known to work, which turns this from a search into an integration exercise — if it does not improve with these numbers, the bug is in your code, not your configuration.

The derived quantities are worth writing as asserts, exactly as the handout does, because every off-policy configuration in §8 has to keep them consistent: micro_train_batch_size = train_batch_size // gradient_accumulation_steps; n_prompts_per_rollout_batch = rollout_batch_size // group_size; n_microbatches_per_rollout_batch = rollout_batch_size // micro_train_batch_size; and train_batch_size >= group_size. When you change epochs_per_rollout_batch later you must change gradient_accumulation_steps in step or memory usage moves under you and you will misattribute an OOM to the algorithm.

The two tips that separate a run that trains from a run that looks like it trains:

With our implementation of the losses, GRPO-Clip should only be used when off-policy (since it requires the old log-probabilities). … You should not differentiate with respect to the old log-probabilities.§7.2, train-loop tips

Both are quiet failures. On-policy, π_θ and π_old are identical so the ratio is exactly 1, clipping never binds, and GRPO-Clip is just the naive loss with extra steps — but if you compute old_log_probs with gradients attached, the ratio's denominator gets a gradient too, the two cancel, and the update dies without an error. Compute old log-probs once per rollout batch under torch.inference_mode() and reuse them across epochs; recomputing per epoch is pure waste.

Validate every 5–10 steps on at least 1024 validation examples — the handout is explicit that fewer is too noisy to compare hyperparameters against. Log loss, gradient norm, token entropy, clip fraction (off-policy only), and train rewards broken into total / format / answer. L17 gives the ordered suspect list when a run trains but does not improve: rewards all tied inside each group so every advantage is zero; gradients flowing through π_old; or a reward function with an unintended shortcut. Watch reward, not loss.

• No adapter, no test. Graded on the write-up plot and rollout samples.
• L17 Alignment: RL 2 — the mechanics lecture; its compute_deltas / compute_loss demo is this loop in miniature
• DeepSeekMath — Shao et al. (2024) · Algorithm 3's source

8 · GRPO experiments

Nine problems, 19 points, roughly 30 H100-hours. Each run takes 2 GPUs (one vLLM, one policy). The handout gives you permission up front to kill a configuration early if it has clearly diverged before 200 steps — the hour estimates assume you exercise it. Each experiment fixes its winner before the next one starts, so the sequence is a greedy coordinate descent through the design space rather than a grid.

Problem (grpo_learning_rate): Tune the learning rate · 2 points · 6 H100 hrs

Deliverable: validation reward curves across several learning rates, a model above 25% validation accuracy, and two sentences on trends in your other logged metrics.

Start from the defaults and sweep around 1e-5. The 25% bar is the first checkpoint that GRPO is genuinely beating both SFT and EI, which were only asked for 15%. Note divergence explicitly where it happens rather than silently dropping the run — divergence at a given learning rate is a data point. This is where logging gradient norm earns its keep: a run that is about to diverge shows it in the gradient norm well before the reward curve turns over, which is what lets you kill it at step 40 instead of 200. Whatever wins here is fixed for every remaining experiment.

Problem (grpo_baselines): Effect of baselining · 2 points · 2 H100 hrs

Deliverable: validation reward curves for no_baseline versus reinforce_with_baseline, plus two sentences on other metrics.

The cleanest demonstration in the assignment of a piece of theory you derived by hand. §6.5 proved that subtracting a state-dependent baseline leaves the estimator unbiased while cutting variance; here you watch what that buys in practice. Stay on-policy, keep use_std_normalization = True, change only the loss type. With binary rewards the effect is stark: without a baseline, every token of every correct solution gets a positive push and every token of every wrong one gets nothing, so on an easy question where all eight rollouts succeed you take a large gradient step that teaches nothing. With the group baseline, that question contributes exactly zero. Look at gradient norm alongside reward — that is where the variance reduction is visible.

Problem (think_about_length_normalization): Think about length normalization · 1 point

Deliverable: pros and cons of masked_mean versus masked_normalize, reasoned out before running anything, with settings where one looks better.

The handout hands you a worked numerical example from the RLHF Book: a batch of two responses, one 4 tokens and one 7, all ratios 1 and all advantages 2. Under masked_mean both sequences yield 2.0 and the per-token gradients come out 0.25 for the short sequence and 0.1429 for the long one. Under masked_normalize with a constant of 7 they yield 1.1429 and 2.0, and every token in the batch gets an identical 0.1429 gradient. That is the whole trade in one table.

The argument for masked_mean is that it is per-sequence unbiased and keeps a single rambling rollout from dominating the batch. The argument against is that it introduces a 1/T factor the policy-gradient derivation in §6.4 never had, and that this factor systematically discounts each token of a long response — which, when the advantage is negative, means wrong long answers are penalized less per token than wrong short ones. That is the length-bias Dr. GRPO and DAPO both identify as a driver of response-length inflation. The constant-denominator form removes the bias at the cost of letting a single long rollout carry more of the batch gradient.

Problem (grpo_length_normalization): Effect of length normalization · 2 points · 2 H100 hrs

Deliverable: end-to-end runs with each reduction, validation answer-reward curves, commentary on any other metric with a visible trend.

Now measure what you just predicted. The handout's hint points at stability metrics, gradient norm in particular — the constant-denominator form does not damp long sequences, so its gradient norm is both larger and noisier, and whether that costs you depends on your clip value and learning rate. Track mean response length too, split by correct and incorrect: the mechanism you argued about in the previous problem shows up there before it shows up in reward. Fix the winner for everything downstream.

Problem (grpo_group_standard_deviation): Effect of standard deviation normalization · 2 points · 2 H100 hrs

Deliverable: use_std_normalization True versus False, validation reward curves, commentary.

The second Dr. GRPO ablation, and a flag you already implemented in compute_group_normalized_rewards. The claim under test: dividing by the group std up-weights questions where the group agreed, because a small denominator makes small reward differences into large advantages. Those low-variance groups are exactly the questions that are too easy (all eight correct) or too hard (all eight wrong) — precisely the ones with the least to teach — so the standard normalization spends gradient where there is no signal, amplified by advantage_eps being all that stands between you and a division by zero. Watch gradient norm and per-group advantage magnitudes.

Problem (grpo_off_policy): Implement off-policy GRPO

Deliverable: the off-policy training path. No points of its own; it gates the two problems after it.

Three changes to the loop. Take multiple epochs of gradient steps per rollout batch, with rollout_batch_size, epochs_per_rollout_batch and train_batch_size jointly determining how many optimizer updates each batch of rollouts funds. Compute response log-probs from the policy once, after the rollout phase and before the inner gradient loop, and keep them as old_log_probs — under torch.inference_mode(), so no graph is built and no gradient can leak into the ratio's denominator. And switch loss_type to grpo_clip, which is the first point in the assignment where clipping does anything at all. Depending on how you structured grpo_train_loop you may already have all of this.

Problem (grpo_off_policy_sweep): Off-policy GRPO hyperparameter sweep · 4 points · 12 H100 hrs

Deliverable: a broad sweep at <50 steps then a focused sweep at 200, an experiment log justifying the ranges you chose, comparison against the on-policy run in both step-count and wall-clock, and commentary including entropy and response length.

The biggest single compute item outside the leaderboard, and the one where the two-stage protocol matters: cheap broad sweep first to find the shape of the landscape, then a focused sweep at full length. Fix rollout_batch_size = 256 and vary epochs_per_rollout_batch and train_batch_size; the ratio of the two controls how far off-policy you drift before the next rollout. Remember to move gradient_accumulation_steps in lockstep to hold the microbatch size — and therefore memory — constant, or you will be comparing configurations that also differ in how much they fit on the GPU.

The wall-clock plot is the point of the whole exercise. Off-policy will usually look worse per gradient step and better per hour, because rollouts dominate the cost and you are amortizing each batch of them over more updates. That trade is why anyone runs off-policy at all. The handout also asks you to compare the entropy curve here against the one from expert_iteration_experiment — two different mechanisms for the same collapse, worth a paragraph.

Problem (grpo_off_policy_clip_ablation): Off-policy GRPO-Clip ablation · 2 points · 2 H100 hrs

Deliverable: implement an unclipped loss type, run it at your best off-policy settings, compare against GRPO-Clip.

Add a "GRPO-No-Clip" loss that is the bare importance-weighted per-token loss −(π_θ/π_old)·A_t — the objective from equation 27 with no min and no clamp. Everything else held fixed. This tests whether the clipping does real work or is cargo-culted from PPO, and it is the most direct payoff from the case analysis in §7.1. Watch gradient norm above all: without the clip, tokens whose probability has moved a long way from π_old contribute unboundedly large ratios, so the unclipped run should be the noisier one, and how much noise costs you depends on how far off-policy your winning configuration went. Entropy and response length are the other two metrics the handout asks about.

Problem (grpo_prompt_ablation): Prompt ablation · 2 points · 2 H100 hrs

Deliverable: validation answer-reward curves for the R1-Zero prompt versus the question-only prompt, with an explanation of the difference.

Swap r1_zero.prompt for question_only.prompt — which is literally {question} — for both training and validation, and swap the reward function to question_only_reward_fn to match. The second swap is not optional: the R1-Zero reward requires </think> <answer> tags that the question-only prompt never asks for, so keeping the old grader would score a uniform zero. question_only_reward_fn instead pulls the answer out of a \boxed{} via extract_answer.

This is the experiment that cashes the warning from §3. The expected result, from Liu et al. 2025, is that question-only starts far higher because Qwen 2.5 Math was pretrained on data in exactly that shape — the R1-Zero prompt begins by fighting the model's prior and spends its first hundred steps recovering ground the bare prompt had for free. The honest reading is that a large part of "RL improvement" on this setup is the model relearning a format it already knew, and noticing that is the real deliverable. Compare entropy, response length and gradient norm between the two.

9 · Leaderboard: GRPO on MATH

Problem (leaderboard): Leaderboard · 16 points · 16 H100 hrs

Deliverable: the best validation accuracy you can reach within 4 hours of training on 2 H100s, plus a screenshot of accuracy against wall-clock time whose x-axis ends at ≤4 hours.

Sixteen points, a quarter of the assignment, and the only open-ended problem. The clock is the constraint, so this is as much a systems problem as an RL one — which is deliberate, and the handout says so: in the naive two-GPU layout at least one device is idle at all times, either the policy waiting on rollouts or vLLM waiting on a gradient step. It suggests lower precision for rollouts or training, torch.compile, and explicitly frees you from the one-vLLM-one-policy split, encouraging better parallelism. Everything you learned in L10 about inference throughput applies directly.

You are not allowed to use any other data or do SFT on reasoning chains from stronger models, etc. You must report validation accuracy on the entire validation set (all 5K examples), using the sampling hyperparameters given above (temperature 1.0, max tokens 1024). You are allowed to filter the train set, or design a curriculum over the data, as you desire. You must use the R1-Zero prompt for validation, and during validation, you must use exactly the r1_zero_reward_fn reward function provided in the starter code.§9, Dataset

Read that carefully: the algorithm is unconstrained, the evaluation is nailed down, and the training-side freedom is real — you may filter the train set or build a curriculum, and you may use a different reward function during training so long as validation uses the official one. The handout also volunteers that omitting the KL penalty cost nothing in the staff's experiments and in the literature (Liu et al. 2025, GRPO-Zero) while saving a resident reference model, though you are welcome to try KL or other regularizers. The one hard model constraint: Qwen 2.5 Math 1.5B Base, or something you derived from it.

For ideas the handout points at four production RL codebases: veRL, TRL, torchtune and oat. Reading them for systems ideas is encouraged; importing their trainers is not.

Submission is a pull request against assignment5-alignment-leaderboard adding a row to the Markdown table, sorted by decreasing accuracy, with your final whole-validation-set accuracy, a link to a learning curve that clearly shows wall-clock time (upload to ./images or link a public W&B report), and a description of what you did. Top 3 get a prize; the staff reproduce the top 5 runs and re-rank them, so save a snapshot of your best configuration. For calibration, the Spring 2025 board topped out at 68.75%, with the top ten between roughly 62% and 69% — against a zero-shot baseline in the single digits and a 25% bar for a passing GRPO run.

Optional supplement · Instruction tuning and RLHF

We provide, as an entirely optional supplement to the required course materials, an assignment on training language models to follow instructions and aligning language models to pairwise preference judgments.§1, safety/RLHF supplement

A parallel assignment on the same repo: different model (Llama 3.1 8B Base), different benchmarks, different failure modes. Where the required half optimizes a verifiable reward on a narrow task, this half builds a generalist assistant and confronts the fact that "helpful" and "harmless" have no verifier at all — you end up with an LLM judge and a preference dataset instead. This is where L15's second half lands as code, and where L12's material on judges, arenas and length bias becomes concrete.

Version note. The Spring 2025 copy of this supplement is no longer served from the repo. The live file is the Spring 2026 edition on the main branch, which is the version described below; it runs on Modal rather than the Together cluster and quotes B200 rather than H100 hours. Its three test files — test_data.py, test_dpo.py, test_metrics.py — and their five adapters exist unchanged on master alongside the 2025 handout, so you can do the supplement from either branch. The prompts moved: master keeps them in cs336_alignment/prompts/, main keeps them in cs336_alignment/prompts_safety/, separate from the RL prompts. The Alpaca SFT template is byte-identical in both locations, which matters because test_per_instance_dpo_loss asserts an exact numeric value computed against it.

Prompting setup. Every zero-shot evaluation nests two templates: a task prompt formatted with the example, then inserted into the {instruction} slot of zero_shot_system_prompt.prompt. That system prompt frames the interaction as alternating # Query: / # Answer: blocks inside Markdown fences, which gives you a stop string: when the model emits # Query: it has started a new turn and you cut. Generation is greedy throughout — temperature 0.0, top-p 1.0 — for all four benchmarks. The ↪ character in the PDF's prompt blocks marks a visual line continuation, not a newline; copy the prompt files, not the PDF.

Problem (mmlu_baseline): Zero-shot MMLU baseline · 4 points

Deliverable: a parser, an evaluation script, and four short write-ups — parse-failure count, throughput, accuracy, and an error analysis over ten wrong predictions.

The parsing problem is the interesting half. A base model asked a multiple-choice question in free-form text will answer with the letter, or the option's text, or a paraphrase of the option, or a short essay that arrives at the answer. Rather than write a parser for all of that, mmlu_zero_shot.prompt constrains the output format up front — "Respond with a single sentence of the form 'The correct answer is _'" — and your run_parse_mmlu_response extracts the letter, returning None when it cannot. That None path is graded: part (c) asks how many generations fail to parse and what they look like. This is the same evaluation-is-a-design-problem point L12 makes at length, met head-on.

• tests/adapters.py · run_parse_mmlu_response — (mmlu_example: dict[str, Any], model_output: str) -> str | None; the example dict carries subject, question, options, answer
• tests/test_metrics.py — test_parse_mmlu_response and test_parse_mmlu_response_unknown; the second checks you return None rather than guessing
• Measuring Massive Multitask Language Understanding — Hendrycks et al. (2021) · data on main at data/mmlu, 57 subjects, dev/val/test CSVs

Problem (gsm8k_baseline): Zero-shot GSM8K baseline · 4 points

Deliverable: a numeric parser, an evaluation script, and the same four write-ups.

The prompt is two lines — the question and Answer: — and the parsing rule is one line: take the last number in the output. "She sold 15 clips." parses to 15. Simple and deliberately imperfect; it will misfire on a model that restates a figure after its conclusion, and noticing that in the error analysis is part of the point. The supplement flags explicitly that this is a different prompt and a different parser from the RL half's question_only.prompt, which asks for boxed answers — do not reuse question_only_reward_fn here.

• tests/adapters.py · run_parse_gsm8k_response — (model_output: str) -> str | None
• tests/test_metrics.py — test_parse_gsm8k_response, test_parse_gsm8k_response_unknown
• Training Verifiers to Solve Math Word Problems — Cobbe et al. (2021) · data at the repo's data folder and on HuggingFace

Problem (alpaca_eval_baseline): Zero-shot AlpacaEval baseline · 4 points

Deliverable: a generation script emitting AlpacaEval-shaped JSON, a throughput estimate, the winrate and length-controlled winrate, and an error analysis on ten dispreferred examples.

The first benchmark with no ground truth at all. You generate a response to each instruction, and a judge — Llama 3.3 70B Instruct, configured in scripts/alpaca_eval_vllm_llama3_3_70b_fn — decides whether it prefers yours or GPT-4 Turbo's. The output format is load-bearing: a JSON array, one object per instruction, with keys instruction, output, generator (constant across entries) and dataset (carried through from the source). Then alpaca_eval --model_outputs … --reference_outputs data/alpaca_eval/alpaca_eval_gpt4_turbo.json --annotators_config scripts/alpaca_eval_vllm_llama3_3_70b_fn --base-dir .

Report both winrate and length-controlled winrate, and note the gap — L12's point about AlpacaEval's length bias is that judges reliably prefer longer answers, and the LC metric exists to correct for it. Part (d) sends you into annotations_seed0_configs.json to read ten losses and asks whether you agree with the judge, which is the honest version of "is this benchmark measuring what I think".

Problem (sst_baseline): Zero-shot SimpleSafetyTests baseline · 4 points

Deliverable: a generation script emitting JSON-lines with prompts_final and output, throughput, the safe-output proportion, and an error analysis on ten unsafe generations.

A hundred prompts (data/simple_safety_tests/simple_safety_tests.csv) covering critical safety risks, judged by the same 70B model via the provided scripts/evaluate_safety.py. The supplement is candid about why the judge is a model: it is an imperfect proxy for human evaluation, but far less emotionally taxing than asking people to read hundreds of harmful completions. Expect a base model to score poorly — a model with no instruction tuning and no refusal training simply continues the text, and there is nothing in pretraining that makes "how do I evade taxes" a question to decline. That gap is the motivation for everything after it.

• SimpleSafetyTests — Vidgen et al. (2024) · the test suite
• Run: uv run python scripts/evaluate_safety.py --input-path … --model-name-or-path …/Llama-3.3-70B-Instruct --num-gpus 2 --output-path …

Problem (look_at_sft): Inspect instruction tuning data · 4 points

Deliverable: 2–4 sentences on ten random training examples — which classic NLP tasks show up implicitly, and how good the prompts and responses actually are.

Four points for reading ten examples, which tells you how seriously the course takes the habit. The mix is UltraChat-200K plus SafetyTunedLlamas, flattened to single-turn {prompt, response} pairs. You will find QA, summarization, rewriting, classification and open generation all present without ever being labelled as such — instruction tuning is the observation that these are one task in disguise — alongside a certain amount of model-generated filler, since UltraChat is synthetic. Forming an opinion about quality here is what lets you interpret the benchmark movements in §5 rather than just reporting them.

Problem (data_loading): Implement data loading · 3 points

Deliverable: a packed PyTorch Dataset and a batching function over it.

Part (a): format each pair with alpaca_sft.prompt, tokenize, concatenate all documents into one stream with the Llama end-of-text token between them, then chop into consecutive non-overlapping chunks of seq_length, dropping the final short chunk. __getitem__ returns {"input_ids", "labels"}, both (seq_length,), with labels the inputs shifted by one — which means the chunk you slice for a given index needs seq_length + 1 tokens available. Packing rather than padding is a throughput decision: no wasted positions, at the cost of documents bleeding across chunk boundaries. Part (b) is the batcher, with shuffle, one pass = one epoch.

The test compares your output against a stored fixture element by element, so ordering and boundary handling must match exactly, and it separately asserts that shuffle=True produces a different sequence than shuffle=False — shuffle the documents before concatenation, as the docstring says, not the chunks afterward.

• tests/adapters.py · get_packed_sft_dataset — (tokenizer, dataset_path, seq_length: int, shuffle: bool) -> Dataset
• tests/adapters.py · run_iterate_batches — (dataset, batch_size, shuffle)
• tests/test_data.py — test_packed_sft_dataset and test_iterate_batches, both against tests/fixtures/sft_sample.jsonl with seq_length=32 and the checked-in Llama 3 tokenizer fixture — so these two run anywhere, no cluster needed

Problem (sft_script): Training script: instruction tuning · 4 points

Deliverable: a finetuning script for Llama 3.1 8B on the instruction data.

Configurable model and optimizer hyperparameters, gradient accumulation, periodic train and validation logging. You may adapt an earlier training script; you may not use Trainer. Mechanically identical to the SFT section of the required assignment — bfloat16, FlashAttention-2, model(input_ids).logits, cross-entropy, accumulate, step, save_pretrained — at 8B instead of 1.5B, so memory pressure is the difference you will feel.

Problem (sft): Instruction tuning · 6 points · 3 B200 hrs

Deliverable: your training setup, the final validation loss, a learning curve, and the serialized model and tokenizer.

One epoch over the data. Save the tokenizer with the model — DPO reloads both, twice.

Problem (mmlu_sft): Evaluate SFT on MMLU · 4 points

Deliverable: throughput versus baseline, accuracy versus baseline, and an error analysis over ten wrong predictions comparing the two models qualitatively.

One rule governs all four §5 evaluations and is easy to miss: drop the zero-shot system prompt and format everything with the Alpaca template you trained on instead. For MMLU and GSM8K that means nesting twice — format the task prompt with the example, then put that formatted string into the Alpaca {instruction} slot. Evaluate a finetuned model with the wrong wrapper and you measure the mismatch, not the model. The throughput question is not filler either: an instruction-tuned model stops rambling and starts terminating, so examples-per-second usually goes up sharply, and the reason is worth a sentence.

Problem (gsm8k_sft): Evaluate SFT on GSM8K · 4 points

Deliverable: throughput comparison, accuracy comparison, and an error analysis over ten wrong predictions.

Same protocol, same nesting: gsm8k_zero_shot.prompt inside alpaca_sft.prompt. The number to watch is whether reasoning accuracy actually improves. Instruction tuning on general chat data is not math training, and the interesting outcomes are either a modest gain from better answer formatting (the last-number parser now finds a number where it used to find prose) or a small loss — the first hint of the alignment tax that dpo_training part (d) asks about directly.

Problem (alpaca_eval_sft): Evaluate SFT on AlpacaEval · 4 points

Deliverable: throughput, winrate and LC winrate against the baseline, and ten dispreferred examples analysed.

This is where instruction tuning should show its largest single gain, because the benchmark measures exactly what you trained for. Format instructions with alpaca_sft.prompt and run the same annotator config. In the annotations file, entries with "preference" equal to 1.0 are the ones where the judge preferred GPT-4 Turbo — that is your dispreferred set. Compare winrate and LC winrate movement separately: if raw winrate jumps much more than LC winrate, a good part of your gain is that the tuned model writes longer.

Problem (sst_sft): Evaluate SFT on SimpleSafetyTests · 4 points

Deliverable: throughput, safe-output proportion against the baseline, and ten unsafe examples analysed.

The SafetyTunedLlamas portion of the training mix is the reason this number should move, and it is a small fraction of the data — the demonstration is that a modest amount of refusal data goes a long way, which is the empirical claim that paper makes. Re-run scripts/evaluate_safety.py and, as before, judge the judge: part (c) asks where you disagree with it.

Problem (red_teaming): Red-team the instruction-tuned model · 4 points

Deliverable: three misuse categories beyond those already discussed, plus three documented red-teaming attempts on your own model.

Interactive rather than scripted, and the write-up wants method and honesty, not a scoreboard: what you tried, how long you tried it, which strategies you used, whether you succeeded. The framing follows Ganguli et al. 2022 — red-teaming is an evaluation method for understanding failure, not a compliance checkbox. Doing it on a model you finetuned yourself an hour earlier is the point; the hundred-prompt benchmark you just ran gives an aggregate number, and this gives you the texture of how a refusal actually breaks.

Problem (look_at_hh): Inspect HH preference data · 2 points

Deliverable: a loader for the Anthropic HH dataset, plus 2–4 sentences on three "helpful" and three "harmless" examples and whether you agree with the annotators.

Four gzipped JSON-lines files on main at data/hh: harmless-base, helpful-base, helpful-online, helpful-rejection-sampled (training splits only). Each line has a chosen and a rejected conversation sharing a prefix. Three processing rules: drop multi-turn conversations where the human sent more than one message, because those diverge past the shared prompt and are no longer a clean pairwise comparison; split each example into instruction plus chosen/rejected responses; and remember which file each came from, since the write-up compares the two categories. gzip and json are all you need.

The write-up question is the good one. Anthropic deliberately declined to define "helpful" or "harmless", leaving it to annotator judgment, so the labels encode a distribution of opinions rather than a rule. You will find pairs where the difference is obvious and pairs where you would have chosen the other way — and that is the actual epistemic status of the preference signal DPO is about to optimize.

• Training a Helpful and Harmless Assistant with RLHF — Bai et al. (2022) · the HH dataset paper; also on HuggingFace

Problem (dpo_loss): DPO loss · 2 points

Deliverable: the per-instance DPO loss for one (prompt, chosen, rejected) triple under a policy and a frozen reference.

The whole of DPO in one function: −log σ(β·[log π_θ(y_w|x) − log π_ref(y_w|x)] − β·[log π_θ(y_l|x) − log π_ref(y_l|x)]). The derivation in §6.1 is worth following because it explains why this is legitimate — the optimal policy for a reward model can be inverted to express the reward in terms of the policy, and when you plug that into the Bradley-Terry reward-model loss the intractable partition function Z(x) cancels between the two responses. No reward model, no sampling, no RL loop; just four forward passes.

Three implementation details the test enforces. Format prompt and response with the Alpaca template and append the EOS token after each response — forget it and the arithmetic is close but not equal. Sum log-probs over the whole concatenated string rather than conditioning on the prompt; the supplement notes the prompt's contribution cancels in the difference under the same model, so unconditional sums are equivalent and simpler. And return the loss on the same device as the policy — the two models are expected to be on different GPUs.

• tests/adapters.py · run_compute_per_instance_dpo_loss — (lm, lm_ref, tokenizer, beta: float, prompt: str, response_chosen: str, response_rejected: str) -> torch.Tensor
• tests/test_dpo.py — test_per_instance_dpo_loss uses the checked-in tiny-gpt2 / tiny-gpt2-ref fixtures and the gpt2 tokenizer with beta=0.5, and asserts torch.isclose(loss, 0.5785, atol=1e-4) — an exact number, so template and EOS handling must match. Runs on CPU, no cluster needed.

Problem (dpo_training): DPO training · 4 points · 1 B200 hr

Deliverable: a DPO training script and a validation-accuracy screenshot, then AlpacaEval, SimpleSafetyTests, GSM8K and MMLU evaluations of the result.

One epoch over HH starting from your SFT checkpoint. The suggested path trades throughput for simplicity: two GPUs holding two copies of the instruction-tuned model (one frozen as π_ref), no batching at all — one example at a time, because four forward passes through an 8B model is already a lot of memory — with gradient accumulation supplying the effective batch size. RMSprop rather than AdamW, following the original DPO work, because AdamW's optimizer state is unaffordable here without quantization tricks.

Track the implicit reward model's classification accuracy on validation: for each pair, check whether the chosen response has higher log-probability than the rejected one under π_θ. That single number is your training signal made legible, and it is what selects the checkpoint.

Parts (b) through (d) are the payoff. AlpacaEval and SimpleSafetyTests should improve, since HH contains exactly the helpfulness and refusal behaviour they test. GSM8K and MMLU are the control, and the question is whether they fall — the alignment tax, observed in the original HH paper and much of the alignment literature since. Whether you see it, and how large it is at 8B on one epoch, is the actual finding.

• Direct Preference Optimization: Your Language Model is Secretly a Reward Model — Rafailov et al. (2023)
• Training language models to follow instructions with human feedback — Ouyang et al. (2022) · the RLHF pipeline DPO collapses

What you hand in

Gradescope, two files. writeup.pdf with every written question answered and typeset, and code.zip with everything you wrote. The repo ships test_and_make_submission.sh to build the archive: it runs uv run pytest -v ./tests --junitxml=test_results.xml (with || true, so failures do not block packaging), then zips the tree into cs336-spring2025-assignment-5-submission.zip, excluding .venv, caches, ./data, and every large binary extension (.pt, .bin, .safetensors, .pkl). It also excludes *.txt, *.json and *.log — so if any of your generation dumps or configs live in those formats and you want them submitted, move them or adjust the exclusion list.

What the write-up owes. Fifteen of the twenty-six required problems are wholly or partly write-up: one commentary on baseline failure modes, seven experiment reports with reward curves, one pre-registered prediction about length normalization, one train-loop report with sample rollouts, and one leaderboard screenshot with the wall-clock x-axis capped at four hours. Every experiment deliverable asks for validation reward or accuracy curves, never loss curves — §6.5 explains why, and grading follows it.

Leaderboard. Separate from Gradescope: a pull request against assignment5-alignment-leaderboard with your row, your learning curve (uploaded to ./images or a public W&B link), and a description of your approach. Keep a reproducible snapshot of your best configuration — the staff re-run the top five.

Compute. Roughly 52 H100-hours across the graded runs of the required assignment: 2 for SFT, 6 for expert iteration, 6 for the learning-rate sweep, 12 for the off-policy sweep, 2 each for the four remaining ablations, and 16 for the leaderboard. The supplement adds 3 B200-hours for instruction tuning and 1 for DPO. Every GRPO run wants 2 GPUs. Start early; this is not an assignment you can compress into a weekend.

Materials, verified

The assignment itself

• cs336_spring2025_assignment5_alignment.pdf — the Spring 2025 handout, v1.0.0, 38 pages, on master
• cs336_spring2026_assignment5_supplement_safety_rlhf.pdf — the optional instruction-tuning and RLHF supplement, v26.0.0, on main
• tests/adapters.py — all sixteen adapter stubs, eleven required plus five for the supplement
• tests/conftest.py — the snapshot fixtures and every hyperparameter the tests use
• cs336_alignment/drgrpo_grader.py — the MATH answer grader; r1_zero_reward_fn and question_only_reward_fn at the bottom
• README.md · pyproject.toml · CHANGELOG.md — setup, the pinned dependency set, and the release history
• assignment5-alignment-leaderboard — Spring 2025 results; top score 68.75% on the full MATH validation set

Models and data the handout names

• Qwen2.5-Math Technical Report — Yang et al. (2024) · the base model, Qwen2.5-Math-1.5B on HuggingFace
• Measuring Mathematical Problem Solving with the MATH Dataset — Hendrycks et al. (2021) · the benchmark; not publicly redistributable
• Countdown-Tasks-3to4 · GSM8K · Tulu 3 SFT Math · DART-Math — the four open substitutes for MATH
• Math-Verify — HuggingFace · for extracting short ground-truth labels from long-form solutions
• Llama 3.1 8B · Llama 3.3 70B Instruct — the supplement's policy and judge
• UltraChat-200K · SafetyTunedLlamas · Anthropic HH — the supplement's instruction and preference data
• AlpacaEval — Li et al. (2023) · the judged-winrate harness; the assignment pins a fork that forwards kwargs to vLLM

Reasoning and RL papers

• DeepSeekMath: Pushing the Limits of Mathematical Reasoning — Shao et al. (2024) · GRPO's origin; Algorithm 3 is its special case
• DeepSeek-R1 — DeepSeek-AI et al. (2025) · the R1-Zero prompt, the SFT traces, and the group-normalized advantage
• Understanding R1-Zero-Like Training: A Critical Perspective — Liu et al. (2025) · Dr. GRPO; three of the four §8 ablations are this paper
• Proximal Policy Optimization Algorithms — Schulman et al. (2017) · the clipped surrogate
• Policy Gradient Methods for RL with Function Approximation — Sutton et al. (NeurIPS 1999) · the original theorem
• Off-Policy Actor-Critic — Degris et al. (2013) · the approximation behind the importance-weighted gradient
• STaR: Bootstrapping Reasoning With Reasoning — Zelikman et al. (2022) · §5's algorithm
• Thinking Fast and Slow with Deep Learning and Tree Search — Anthony et al. (2017) · expert iteration
• Show Your Work: Scratchpads for Intermediate Computation — Nye et al. (2021) · finetuned chain-of-thought
• Chain-of-Thought Prompting Elicits Reasoning — Wei et al. (2023) · prompted chain-of-thought
• OpenAI o1 System Card — OpenAI et al. (2024) · reasoning RL at scale
• Kimi k1.5: Scaling Reinforcement Learning with LLMs — Kimi Team et al. (2025) · the RL infrastructure this assignment miniaturizes
• SimpleRL-Zoo — Zeng et al. (2025) · zero-RL on open base models
• DAPO: An Open-Source LLM RL System at Scale — Yu et al. (2025) · the other source for the length-normalization argument
• Training Verifiers to Solve Math Word Problems — Cobbe et al. (2021) · GSM8K and verifier-based reranking
• Language Model Cascades — Dohan et al. (2022) · ReST — Gulcehre et al. (2023) · further expert-iteration lineage
• Efficient Memory Management for LLM Serving with PagedAttention — Kwon et al. (2023) · why vLLM is fast enough to RL against

Alignment and safety papers (supplement)

• Direct Preference Optimization — Rafailov et al. (2023) · the loss you implement
• Training language models to follow instructions with human feedback — Ouyang et al. (2022) · the RLHF pipeline, reward model and all
• Training a Helpful and Harmless Assistant with RLHF — Bai et al. (2022) · the HH dataset and the alignment tax
• Measuring Massive Multitask Language Understanding — Hendrycks et al. (2021) · MMLU
• SimpleSafetyTests — Vidgen et al. (2024) · the safety suite
• Red Teaming Language Models to Reduce Harms — Ganguli et al. (2022) · the framing for the red-teaming problem

Reference implementations and further reading

• Spinning Up in Deep RL — Achiam (2018) · §6's presentation follows it closely; the PPO-Clip page is the case analysis §7.1 reuses
• Reinforcement Learning from Human Feedback — Lambert (2024) · source of the length-normalization worked example
• veRL · TRL · torchtune · oat — the four codebases the leaderboard section points at for systems ideas
• sail-sg/understand-r1-zero — where drgrpo_grader.py and the stop-string trick come from
• Open-R1 · TinyZero · GRPO-Zero — small open reproductions of reasoning RL
• vLLM offline inference example (pinned to v0.7.2, the version the assignment requires) · vLLM docs
• torch.inference_mode — how to compute old log-probs without a graph
• CS336 Spring 2025 course site — schedule, slides, all five assignments