Chapter 6: Training reasoning models with reinforcement learning

Introduction

The introduction places reinforcement learning among the ways to improve reasoning models, contrasting training methods with increased inference compute.

• Two routes to better reasoning: Explain that reasoning performance and answer accuracy can improve through either more inference compute or specific model-training methods, and that this chapter focuses on reinforcement learning as a training method.
  Why it matters: This distinction places the chapter’s training approach alongside the inference-time approaches discussed earlier.

Key terms: inference compute budget, model-training methods, reinforcement learning (RL)

6.1 Introduction to RL for LLMs

This section introduces reinforcement learning as a way to improve LLM reasoning during training, contrasts it with inference-time scaling and pretraining, and places it within common LLM training pipelines.

• Two ways to scale reasoning: Distinguish inference-time scaling from training-time scaling by when each spends additional computation, and explain that the two approaches can be combined.
  Why it matters: It clarifies which kind of improvement this chapter pursues without treating earlier inference techniques as alternatives that must be abandoned.
• Placement of reasoning RL: Describe where RL usually fits in LLM training and explain why reasoning-focused RL might instead be applied directly to a pretrained base model.
  Why it matters: The starting model and training stages affect both reasoning accuracy and what you can conclude caused an improvement.
• Pretraining and RL objectives: Explain the conceptual difference between what pretraining optimizes and what RL can optimize for an LLM.
  Why it matters: It explains why adding RL can improve reasoning behavior even though pretraining has already given the model substantial abilities.

Key terms: inference-time scaling, training-time scaling, post-training, pretrained base model, instruction fine-tuning, reasoning training, preference tuning, pretraining, next-token prediction, whole outputs

6.1.1 The original RL pipeline with human feedback (RLHF)

This section explains how RLHF uses human rankings to train a reward model, then uses its scores to fine-tune an LLM’s responses.

• Whole response preference tuning: Distinguish RLHF’s optimization of whole responses according to human preferences from pretraining and supervised fine-tuning based on next-token prediction. Explain why an annotator might prefer the more specific laptop recommendation in the example.
  Why it matters: It clarifies what kind of improvement human feedback is meant to encourage.
• Human ranked reward model: Trace how multiple responses to a prompt are ranked by human annotators, converted from pairwise preferences into training targets, and used to train a reward model that assigns a scalar score to new responses.
  Why it matters: The reward model turns a limited set of human judgments into scores usable throughout training.
• Reward guided model updates: Describe the second RLHF stage: the reward model scores the target LLM’s answers, and those scores are used to fine-tune the LLM. State the shared generate-score-update structure that links RLHF to RLVR, without assuming their reward sources are the same.
  Why it matters: This pipeline is the bridge from human preference tuning to the chapter’s reasoning-focused RL methods.

Key terms: RLHF, preference-tuning, next-token prediction, whole outputs, human preference, human annotators, pairwise preferences, reward model, scalar reward score, target LLM

6.1.2 From human feedback to verifiable rewards

This section explains how reinforcement learning with verifiable rewards (RLVR) replaces a learned reward model with automatic checks, and where that simplification helps or limits training.

• Verifiable reward signals: Explain how RLVR replaces RLHF’s learned reward model with deterministic verification, using the math-answer check as an example.
  Why it matters: This is the key substitution that makes RLVR simpler than the RLHF setup described here.
• Single RLVR training loop: Trace the simplified RLVR loop from a generated response through verification to a model update, and contrast it with RLHF’s two-step pipeline.
  Why it matters: The loop connects automatic verification to training rather than treating the verifier as only an evaluation tool.
• Benefits and limits of verification: Explain why verifiable rewards reduce cost and labeling effort while limiting RLVR to tasks with reliable checks.
  Why it matters: The same requirement that makes rewards cheap and consistent also determines where RLVR can be used.

Key terms: RLVR, verifiable rewards, math verifier, ground truth, RLVR training loop, verifier, rewards, reliable verification signals, deterministic, reproducible

6.2 RLVR using GRPO

The section distinguishes RLVR’s verifiable-reward setup from GRPO’s role in updating the LLM and explains why GRPO is a resource-friendly alternative to PPO.

• RLVR and GRPO roles: Distinguish RLVR’s role in providing automatically verifiable rewards from GRPO’s role as the policy optimization algorithm that uses the learning signal to update the LLM.
  Why it matters: Keeping the reward setup separate from the update algorithm makes the training approach easier to follow.
• GRPO resource advantage: Explain why the section describes GRPO as more resource friendly than PPO and where its learning signal comes from.
  Why it matters: This explains the practical reason for choosing GRPO without treating PPO as incompatible with RLVR.

Key terms: RLVR, automatically verifiable rewards, policy, Group Relative Policy Optimization (GRPO), Proximal Policy Optimization (PPO), GRPO, value model, relative comparisons, sampled responses

6.2.1 High-level GRPO intuition via a chef analogy

A chef analogy introduces how GRPO generates and compares several responses, then updates the model while limiting changes from its original behavior.

• Prompts rollouts and whole-response rewards: Describe how one prompt leads to several rollouts and completions, and explain that each completion receives feedback on its final result rather than on individual generation steps.
  Why it matters: This establishes what is generated and what receives feedback before the model is updated.
• Relative feedback and policy adjustment: Explain how GRPO compares the rewards of responses to the same request and uses that relative feedback alongside how characteristic each response is of the current model to favor better outcomes.
  Why it matters: Comparing responses identifies which choices the adjustment should encourage or discourage.
• Style preservation and final update: Explain why the analogy consults an original cookbook and how its style-preservation penalty combines with the preference-based adjustment before model weights are updated.
  Why it matters: The update aims to improve responses without letting the model drift too far from its starting behavior.

Key terms: prompt, rollouts, completions, reward, relative feedback, logprobs, policy gradient loss, original cookbook, style-preservation penalty, overall loss

6.2.2 The high-level GRPO procedure

This section introduces the GRPO overview and explains why the book begins with a version that omits the KL loss term.

• Optional KL loss term: Explain that the book’s initial GRPO implementation omits the KL loss term, even though it is part of the original formulation, because it is not strictly essential and omitting it simplifies implementation.
  Why it matters: This explains why the implementation that follows differs from the original GRPO formulation.

Key terms: GRPO, KL loss term

6.3 Loading a pretrained model

This section loads a pretrained base model for GRPO rollouts and checks that text generation runs correctly before training.

• Pretrained model for rollouts: Explain that the tokenizer and pretrained base model are loaded so the model can generate rollouts—answers to prompts—needed for GRPO.
  Why it matters: This establishes what model will produce the answers used during GRPO training.
• Generation smoke test: Describe how the section checks the loaded model with a simple math prompt using temperature and top-p sampling, and why an incorrect answer does not invalidate that check.
  Why it matters: It keeps a code-loading check distinct from judging the model’s mathematical accuracy.
• CPU default and GPU option: Explain why the code defaults to CPU and what change the author recommends after completing the chapter.
  Why it matters: It explains the default device choice without mistaking it for a requirement to train on CPU.

Key terms: pretrained model, tokenizer, base model, rollouts, GRPO, temperature, top-p sampler, math prompt, CPU, GPU

6.4 Loading a MATH training subset

This section loads a MATH training subset that excludes evaluation problems and explains which fields are used to train and verify model responses.

• Separate training and evaluation problems: Explain why the MATH training subset excludes every MATH-500 evaluation problem and state the expected size of the loaded training set.
  Why it matters: A clean split makes evaluation a more trustworthy test of the trained model.
• Training example field roles: Describe how the relevant fields of a MATH training example are used: the problem is the prompt, and the answer is the target for verifying the model’s output.
  Why it matters: These field roles determine what the model responds to and how its response is checked.
• Worked solution field tradeoff: Explain why the available worked solution is not used to evaluate intermediate reasoning steps in this training approach.
  Why it matters: A correct answer need not follow the dataset’s particular worked-out path.

Key terms: MATH training set, MATH-500, data leakage, problem, answer, math verifier, solution, intermediate reasoning steps, overfitting

6.5 Sampling rollouts

This section defines rollouts for GRPO and adapts the existing text-generation code to sample responses that can be used in training.

• Rollouts as sampled answers: Explain that a rollout is a complete model-generated answer to a prompt, and that GRPO begins by generating multiple rollouts.
  Why it matters: Rollouts are the responses that the later GRPO steps work with.
• Training-compatible sampling mode: Explain why the earlier generation function’s @inference_mode decorator is unsuitable for this training setup and why sample_response uses @torch.no_grad instead.
  Why it matters: The choice of decorator lets generation remain efficient without making the function incompatible with the later training step.
• Sample response outputs: Describe what sample_response returns and why its generated response includes the end-of-sequence token.
  Why it matters: These outputs make the sampled response usable by later functions, while the ending token protects a useful generation behavior.

Key terms: rollout, prompt, GRPO, @inference_mode, @torch.no_grad, backward pass, gradient tracking, sample_response, token IDs, answer text

6.6 Calculating rewards

This section explains how GRPO assigns rewards to generated rollouts using final-answer correctness and a required answer format, then proposes a partial-credit variation.

• Correctness and format rewards: Explain how the reward function scores each rollout using the math verifier, and why a correct answer must also use the required \boxed{} format to receive a nonzero reward.
  Why it matters: The reward determines which kinds of model responses training encourages.
• Partial credit format shaping: Apply the exercise’s proposed reward change by stating the reward for a correct boxed answer, a correct unboxed answer, and any other answer.
  Why it matters: The exercise shows how changing reward values can distinguish correctness from compliance with the output format.
• Final versus intermediate rewards: State what the section reports about attempts to reward intermediate solution steps and the resulting preference for final-answer correctness rewards.
  Why it matters: It clarifies why the reward discussed here evaluates the final answer rather than each reasoning step.

Key terms: GRPO, rollout, math verifier, reward_rlvr, fallback=None, \boxed{}, format-aware reward shaping, partial credit, process reward models, intermediate solution steps

6.7 Preparing learning signals from rollouts via advantages

This section turns rollout rewards into group-relative advantages and explains how those advantages guide policy updates.

• Group Relative Advantage Calculation: Explain how an advantage compares one rollout’s reward with other rollouts for the same prompt. Describe the calculation using the group’s mean and standard deviation, including the small constant that prevents division by zero.
  Why it matters: This calculation turns individual correctness rewards into a comparison among answers to the same prompt.
• Advantage Guided Policy Updates: Explain how positive, negative, and near-zero advantages affect a rollout’s contribution to a policy update. State what happens to advantages when every rollout receives the same reward.
  Why it matters: The sign and size of an advantage determine whether a rollout encourages a behavior, discourages it, or contributes little to learning.

Key terms: advantage, reward, rollout, mean reward, standard deviation, group relative, gradient, policy update, GRPO loss, zero-advantage cases

6.8 Scoring rollouts with sequence log probabilities

This section computes a differentiable log probability for each complete rollout and explains why GRPO uses that score rather than an average per-token score.

• Whole Rollout Log Probability: Explain how a sequence-level logprob is computed from generated-token logprobs and why GRPO uses it instead of their average.
  Why it matters: The score must describe the same complete response to which GRPO assigns its reward and advantage.
• Differentiable Rollout Scoring: Describe the changes that make the sequence-logprob calculation usable for training, including how PyTorch can propagate gradients through the sum.
  Why it matters: A rollout score can support a model update only if gradients can flow through its calculation.
• Length Effects of Summation: Explain how summing logprobs affects scores for responses of different lengths and the resulting preference between equally good answers.
  Why it matters: Summing rather than averaging changes how response length affects the training signal.

Key terms: token-level logprobs, sequence-level logprobs, rollout, advantage, torch.sum, @torch.inference_mode(), computation graph, torch.gather, SumBackward0, sequence length

6.9 From advantages to policy updates via the GRPO loss

This section combines rollout advantages and sequence-level log probabilities into the GRPO policy gradient loss and explains why the loss is negated and the advantages are detached.

• GRPO policy gradient loss: Describe how the GRPO policy gradient loss is computed: multiply each rollout’s sequence-level log probability by its advantage, average across rollouts, and negate the result. Explain that the sequence-level log probability is the sum of the generated tokens’ log probabilities.
  Why it matters: This calculation turns the rollout scores into the loss used for a policy update.
• Negative loss sign: Explain why the advantage-weighted log probability objective is multiplied by −1 before passing it to a PyTorch optimizer.
  Why it matters: The sign ensures that minimizing the implemented loss pursues the intended policy objective.
• Detached rollout advantages: Explain why the advantages are detached when computing the policy gradient loss.
  Why it matters: Detaching keeps the update focused on model parameters that influence the rollout log probabilities.

Key terms: policy gradient loss, advantage, sequence-level logprob, rollout, maximization, minimization, .detach(), advantages, gradients, logprobs

6.10 Putting everything together in a single GRPO function

This section combines the GRPO stages in one loss function, explains two implementation choices, and interprets a test result with zero loss.

• Combined GRPO loss stages: Explain that compute_grpo_loss combines five conceptual GRPO stages in one function, and that the code’s ordering simplifies its loops without changing the logical operations.
  Why it matters: This distinction helps you read the combined function without mistaking its code order for a different algorithm.
• Simplified policy gradient objective: Explain why pg_loss equals the overall loss in these examples and why the KL loss term is omitted.
  Why it matters: It explains why the reported loss contains no separate KL contribution.
• Zero reward test outcome: Trace why the MATH test example produces zero reward, zero advantages, zero loss, and no parameter update if trained.
  Why it matters: The example shows how a failed correctness check can leave this training example with no update signal.

Key terms: compute_grpo_loss, GRPO stages, pg_loss, policy gradient loss, KL loss term, MATH training dataset, rewards, advantages, loss, \boxed{6}

6.11 Implementing the GRPO training loop

This section puts GRPO into a PyTorch training loop, explains practical settings for running it, and shows how to interpret its reward and loss output.

• GRPO training loop structure: Describe how the RLVR training loop follows a conventional PyTorch loop, while computing a GRPO loss instead of a standard supervised loss.
  Why it matters: It connects the GRPO loss to the familiar steps that actually update the model.
• Training run resource choices: Explain how the section's device, rollout, step, and checkpoint settings balance training practicality against time, memory, and disk use.
  Why it matters: These choices help make a GRPO run feasible on available hardware.
• Reward and loss interpretation: Explain what average reward reports and why fluctuating or near-zero GRPO loss should not be read as a simple measure of progress.
  Why it matters: It prevents a single low loss or noisy step from being mistaken for evidence that reasoning has improved.

Key terms: RLVR, GRPO loss, training loop, checkpoints, cpu, cuda, mps, num_rollouts, max_new_tokens, checkpoint_every

6.12 Loading and evaluating saved model checkpoints

Shows how to load and evaluate saved GRPO checkpoints on MATH-500, then interprets their accuracy, response lengths, and training limitations.

• Checkpoint loading and evaluation: Describe how to load a saved GRPO checkpoint and evaluate it with the existing MATH-500 pipeline.
  Why it matters: This lets you check a trained model’s results without repeating an expensive training run.
• Training and evaluation token limits: Distinguish the training rollout limit from the evaluation generation limit and explain how the training limit affects rewarded answers and response length.
  Why it matters: The distinction prevents you from mistaking a training constraint for the evaluation limit when reading the results.
• GRPO accuracy and training stability: State the reported MATH-500 result and explain what the short training run does, and does not, establish about further training.
  Why it matters: The reported gain is substantial, but it should not be read as evidence that simply training longer will keep helping.
• Response length and compute tradeoff: Explain why the chapter capped training response length despite the potential accuracy benefits of longer responses.
  Why it matters: It explains why the reported experiment balances possible reasoning gains against training cost.

Key terms: checkpoint, PyTorch state dict, .pth, MATH-500 verification pipeline, max tokens, rollout, final boxed answer, average tokens, MATH-500, GRPO

6.13 Summary

The section summarizes how RLHF and RLVR train LLMs, how GRPO turns verified answers into a training loss, and the practical cost and potential benefit of that training.