GRPO Derivation

The Real Problem: PPO Needs a Critic It Cannot Afford

The previous section laid out PPO's machinery: a clipped importance ratio, a value head that supplies the per-token baseline, and a per-token KL anchor to the reference policy. The algorithm works. It also pays a heavy price. A 70B policy demands a 70B value head running in lockstep — the value network is a second model the same size as the policy, with its own forward pass, its own backward pass, and its own slice of AdamW state. At full mixed precision that is roughly $16 N$ bytes per trainable parameter, doubled because there are two such models. For a 70B run this is over a terabyte of memory before a single activation is allocated.

Worse, the value head is dead weight in the most literal sense. Nothing in the policy's output depends on $V(s_t)$; its only job is to subtract a per-state expected return from the actual return so the policy gradient sees a centred signal. The value head is a variance reducer, not a behavioural component. We are paying 70 billion parameters to compute one scalar per token whose only function is to make another loss less noisy.

And the value head is hard to train. It chases a moving target: the policy keeps changing, so the return distribution keeps shifting, so the value head's MSE target keeps drifting. PPO ships with a clipped value loss precisely because uncontrolled value updates routinely blow up. Every engineer who has stood up a from-scratch PPO-RLHF run has felt this pain — value loss oscillating, advantages collapsing, the run wedging on a plateau that has nothing to do with the policy itself.

GRPO was born from a single question: can we get a usable per-token advantage without a value head? The DeepSeek-Math team (2024) answered yes, with a one-line trick that has since propagated through DeepSeek-R1, Qwen-2-Math, Llama-3 math fine-tunes, and most open-source rule-based RL recipes. The trick: for every prompt, sample G rollouts at once, and use the dispersion of their final rewards as the baseline. The group is its own critic.

Intuition: Let the Group Be Its Own Baseline

Imagine you are grading G = 8 students who all wrote the same essay. You do not have an objective rubric; you only have your gut ranking. The natural way to teach the class is relatively: praise the essays that scored above the average of these eight, criticise the ones that scored below. Whether the average is “objectively good” is irrelevant — what matters is which essays stood out from this group.

GRPO is that intuition mechanised. For every prompt, sample G responses. Score each one. Subtract the group mean. Now you have a signed advantage per response — positive means “do more of this”, negative means “do less of this”. There is no value network because the group itself supplies the answer to “what was expected of this prompt?”. The expected reward of the current policy on this prompt is, well, the empirical mean of the G samples we just drew from it.

Why does this work? Because the policy gradient does not need an absolute baseline — it only needs an unbiased one. Any baseline $b(s)$ that does not depend on the action $a$ leaves the gradient unbiased (this is the classical baseline-subtraction theorem from REINFORCE). The group mean of rewards on the same prompt is a perfectly valid such baseline — it depends only on the prompt and on the current policy, not on any individual action. So we are allowed to subtract it, and the gradient stays pointing in the right direction.

One-line mental model: PPO uses the value function $V(s_t)$ as a baseline because $\mathbb{E}_{a}[R \mid s_t] = V(s_t)$. GRPO replaces $V(s_t)$ with the empirical Monte Carlo estimator of that same expectation — the mean reward of G samples drawn at $s_t = \text{prompt}$. The G samples are the value network.

A physical analogy: PPO carries an absolute thermometer that reads “true” expected reward at every state. GRPO carries no thermometer at all; instead it compares the temperature of G samples against each other and trusts that one of them is above average. The first is more informative when calibrated; the second is dramatically cheaper and cannot drift out of calibration because there is nothing to calibrate.

Deriving GRPO from PPO, One Substitution at a Time

We start exactly where § 14.6 ended. PPO's RLHF objective on a single prompt is:

$J^{\text{PPO}}(\theta) = \mathbb{E}_t \left[ \min\left( r_t(\theta) A_t,\; \mathrm{clip}\big(r_t(\theta), 1{-}\varepsilon, 1{+}\varepsilon\big) A_t \right) \right] - \beta\, D_{\text{KL}}\big(\pi_\theta \,\|\, \pi_{\text{ref}}\big)$

where $r_t(\theta) = \pi_\theta(a_t \mid s_t) / \pi_{\theta_{\text{old}}}(a_t \mid s_t)$ is the importance ratio and the advantage $A_t$ is the GAE estimate built on top of a learned value head $V_\phi(s_t)$.

Substitution one — replace GAE with a single trajectory-level reward. For LLM tasks the per-token reward is zero everywhere except at the EOS token, where it is the reward model's score $R(\tau)$. The GAE estimate over such a sparse reward stream is dominated by that one terminal value; in the limit $\lambda \to 1$ it collapses to

$A_t^{\text{MC}} = R(\tau) - V_\phi(s_t)$

which is just the Monte-Carlo return minus the value baseline. Already we have one clean expression for the advantage that does not need the GAE bookkeeping; we keep it.

Substitution two — replace the value baseline with the group mean. Drop the learned $V_\phi(s_t)$ entirely. In its place, sample G rollouts $\tau_1, \ldots, \tau_G$ from the OLD policy on the same prompt and use the empirical mean of their rewards as the baseline:

$\bar R = \frac{1}{G} \sum_{i=1}^{G} R(\tau_i), \qquad A_i = R(\tau_i) - \bar R$

This is unbiased: the group mean does not depend on any individual rollout's actions, so the baseline-subtraction theorem applies. It is also high variance — for a small G the empirical mean is a noisy estimator of the true expected reward. GRPO addresses the variance not by adding parameters back, but by standardising: divide by the group standard deviation.

Substitution three — standardise by the group std. Compute the group standard deviation $\sigma_R$ alongside the mean and scale:

$A_i = \frac{R(\tau_i) - \bar R}{\sigma_R + \epsilon}$

Two things happen at once. First, the scale of the advantage is now per-prompt normalised — an $A_i = +1$ means “one std above the group mean for this prompt”, regardless of whether this prompt's rewards live in $[0, 1]$ or $[-100, 100]$. Second, the loss magnitude across batches becomes comparable, which is what lets a single learning rate and a single KL coefficient work for prompts of wildly different difficulty.

Substitution four — broadcast the scalar advantage to every token. The advantage we just defined is one number per rollout. The policy update needs one number per token. GRPO makes the simplest possible choice: copy the same $A_i$ to every token of rollout $i$:

$A_{i, t} = A_i \quad \text{for all } t \in \text{response tokens of rollout } i$

This is a deliberately coarse credit assignment. A 200-token response gets credit for every token equally — the one tactically correct token gets the same advantage as 199 boilerplate tokens. PPO with GAE would give each token a different advantage based on local TD error; GRPO trades that resolution for the elimination of the value head. In practice, with large enough G and a reward signal that distinguishes responses on overall quality (not on token-level details), this trade pays off.

The Group-Relative Advantage

Collecting the substitutions, the GRPO advantage in its final form is one line:

$A_{i, t} \;=\; \frac{R(\tau_i) - \frac{1}{G} \sum_{j=1}^{G} R(\tau_j)}{\sqrt{\frac{1}{G} \sum_{j=1}^{G} \big(R(\tau_j) - \bar R\big)^2} + \epsilon}$

Read it slowly. The numerator is “this rollout's reward minus the group mean” — the centred signal. The denominator is the group standard deviation (population formula — divide by G, not G−1), regularised by a small $\epsilon$ to prevent division by zero on degenerate groups. The result is broadcast across all tokens of rollout $i$ — every response token sees the same scalar advantage.

Two consequences of this choice are worth naming up front. One: GRPO trades model parameters (the value head) for inference samples (the G rollouts). If your inference is fast and your reward is cheap to compute, GRPO is strictly better. If your inference is slow or your reward call is expensive, GRPO can be worse than PPO in wall-clock terms. Two: the std normalisation makes the gradient scale-free, which means the choice of $\beta_{\text{KL}}$ and the learning rate become more universal — DeepSeek-Math uses the same hyperparameters across math problems of vastly different reward distributions.

Importing PPO's Clip — Why GRPO Keeps the Ratio

With the advantage defined, the rest of GRPO is PPO almost verbatim. Define the same importance ratio:

$r_{i, t}(\theta) = \frac{\pi_\theta(a_{i, t} \mid s_{i, t})}{\pi_{\theta_{\text{old}}}(a_{i, t} \mid s_{i, t})}$

and the same clipped surrogate, replacing PPO's per-token advantage $A_t$ with GRPO's broadcast advantage $A_{i, t}$:

$L_{i, t}^{\text{CLIP}}(\theta) = \min\!\left( r_{i, t}(\theta) A_{i, t},\; \mathrm{clip}\big(r_{i, t}(\theta), 1{-}\varepsilon, 1{+}\varepsilon\big) A_{i, t} \right)$

Why keep the clip at all if the value head is gone? Because the ratio clip protects against the OTHER PPO failure mode — taking K gradient steps on the same rollout (the ‘outer loop’ of PPO/GRPO) and having the policy drift far enough between steps that the importance correction breaks down. The clip prevents a runaway feedback loop in which a high-ratio token gets pushed higher on the next minibatch, raising its ratio further, and compounding. Even without a critic, that feedback loop is real; the clip is the cheap insurance against it.

Crucially, $\varepsilon = 0.2$ remains the right default. The PPO ablations that fixed this value were about ratio dynamics, not about whether a critic was present. Some recent variants (Dr. GRPO, Open-RLHF's “reinforce++”) experiment with $\varepsilon \in [0.1, 0.3]$, but the differences are second-order — get the advantage and the KL right first, then think about $\varepsilon$.

Anchoring to the Reference: The k3 KL Term

Per-step drift is bounded by the clip; cumulative drift over thousands of steps is bounded by a separate per-token KL penalty against the frozen reference (SFT) policy. GRPO keeps PPO's formulation here, including the k3 unbiased estimator:

$\hat D_{\text{KL}}\big(\pi_\theta \,\|\, \pi_{\text{ref}}\big)_{i, t} = \exp\!\big(\log r_{i, t}^{\text{ref}}\big) - \log r_{i, t}^{\text{ref}} - 1$

where $\log r_{i, t}^{\text{ref}} = \log \pi_\theta(a_{i, t} \mid s_{i, t}) - \log \pi_{\text{ref}}(a_{i, t} \mid s_{i, t})$. The k3 estimator is always non-negative, unbiased in expectation, and has lower variance than the naive $\log r$ estimator that early policy- gradient code used. It costs nothing extra to compute — the ingredients are the same $\log \pi_\theta$ and $\log \pi_{\text{ref}}$ we already need for the ratio.

The coefficient $\beta_{\text{KL}}$ tends to be slightly higher in GRPO than in PPO-RLHF — DeepSeek-Math uses $\beta_{\text{KL}} = 0.04$, roughly double the PPO-RLHF default of 0.02. The reason: in PPO the value head provides indirect smoothing of the policy (a noisy value estimate produces noisy advantages which damp the policy gradient); GRPO has no such smoothing, so the KL term has to do more anchoring on its own.

Putting It Together: The Full GRPO Objective

The complete GRPO loss for one prompt's group is the weighted sum of the clipped policy objective and the KL penalty, averaged over tokens within each rollout and then over rollouts within the group:

$\mathcal{L}^{\text{GRPO}}(\theta) = -\frac{1}{G} \sum_{i=1}^{G} \frac{1}{|\tau_i|} \sum_{t \in \tau_i} \left[ L_{i, t}^{\text{CLIP}}(\theta) - \beta_{\text{KL}} \, \hat D_{\text{KL}}\!\left(\pi_\theta \,\|\, \pi_{\text{ref}}\right)_{i, t} \right]$

Read the three layers of averaging from the inside out. The innermost is the per-token loss combining the clipped surrogate and the KL penalty. The middle is $\frac{1}{|\tau_i|}$ — average over the response tokens of rollout $i$. The outermost is $\frac{1}{G}$ — average over rollouts in the group. The leading minus turns the objective into a loss that gradient descent will minimise.

The double-mean (sequence-mean then group-mean) is a deliberate bookkeeping choice that gives every rollout equal weight in the final loss regardless of its length. It also introduces a length-bias artefact that the Dr. GRPO paper analyses in detail (short responses get effectively up-weighted relative to long ones because a small $|\tau_i|$ denominator inflates the contribution of every token in that rollout). The two principal variants in practice are:

For the rest of this section we use vanilla GRPO — the formula above with sequence-mean then group-mean — because it is what DeepSeek-Math, DeepSeek-R1, and Qwen-2-Math ship by default. The engineering subsection below documents the length-bias trade-off you accept by choosing it.

Manual Numerical Walkthrough

We trace one full GRPO update by hand for G = 4 rollouts of T = 3 tokens each, all from the same prompt. Numbers small enough to check on paper; structure identical to a production step.

Step 1: the group of rollouts

Four rollouts have been sampled from the old policy on a single prompt. Each rollout consists of three response tokens; we list the action sequences and the log-probabilities the old policy assigned to them, plus the scalar reward each response received from the reward model (or verifier).

Step 2: the group baseline and standard deviation

Mean of the four rewards: $\bar R = (0.90 + 0.30 - 0.10 + 0.70) / 4 = 0.450$.

Population variance: $\sigma^2_R = \frac{1}{4}\big((0.90 - 0.45)^2 + (0.30 - 0.45)^2 + (-0.10 - 0.45)^2 + (0.70 - 0.45)^2\big)$ $= \frac{1}{4}(0.2025 + 0.0225 + 0.3025 + 0.0625) = 0.1475$.

Standard deviation: $\sigma_R = \sqrt{0.1475} \approx 0.3841$.

Step 3: the per-rollout advantages

Subtract the mean and divide by the std (with $\epsilon = 10^{-8}$ in the denominator):

Two checks worth doing in your head: $\sum_i A_i = +1.171 - 0.390 - 1.431 + 0.651 = +0.001 \approx 0$ (round-off). And the std of the advantages is $\sqrt{(1.171^2 + 0.390^2 + 1.431^2 + 0.651^2) / 4} \approx 1.00$. Both are guaranteed by the standardisation, but worth confirming once to internalise that the advantage scale is per-prompt comparable.

Step 4: broadcast to tokens

Every token of rollout $i$ receives the same scalar advantage $A_i$. The advantage tensor is now of shape $(G, T) = (4, 3)$ with each row constant:

Step 5: the policy ratio at one token

Take token $(i, t) = (1, 1)$ — the first token of rollout 1. The old policy assigned it $\log \pi_{\text{old}} = -1.10$. The live policy (after one gradient step) now assigns it $\log \pi_{\text{new}} = -0.21$ (computed from the new logits via log-softmax — we will see this exact computation in the code below).

Log ratio: $\log r = -0.21 - (-1.10) = +0.89$.

Ratio: $r = e^{0.89} = 2.43$. Way outside the trust region $[0.8, 1.2]$.

Step 6: clipped surrogate at this one token

The two surrogates: $r \cdot A = 2.43 \cdot 1.171 = +2.845$ and $\mathrm{clip}(r, 0.8, 1.2) \cdot A = 1.2 \cdot 1.171 = +1.405$. The min picks the smaller: $+1.405$. The clip fired, the gradient through $r_{1,1}$ at this token is zero — we have already pushed this rewarded token's probability far enough; do not push further.

Compare to rollout 3's first token where $A = -1.431$. Suppose its ratio after one step is $r = 1.6$. Then $r \cdot A = -2.290$ and $\mathrm{clip}(r, 0.8, 1.2) \cdot A = 1.2 \cdot (-1.431) = -1.717$; the min picks $-2.290$. The clip does NOT fire — gradient flows fully because the policy is moving in the rewarded direction (decreasing a bad action's probability), and the asymmetric clip never blocks that. This is exactly the PPO behaviour from § 14.6, carried over into GRPO unchanged.

Step 7: per-token KL at the same first token

The reference model assigned this token $\log \pi_{\text{ref}} = -1.05$. Then $\log r^{\text{ref}} = -0.21 - (-1.05) = +0.84$. The k3 estimator gives $\hat D_{\text{KL}} = e^{0.84} - 0.84 - 1 = 2.317 - 0.84 - 1 = 0.477$. With $\beta_{\text{KL}} = 0.04$, the KL contribution at this token is $0.04 \cdot 0.477 = 0.0191$.

Step 8: per-token loss and aggregation

Per-token loss at $(1, 1)$: $\ell = -L^{\text{CLIP}} + \beta_{\text{KL}} \hat D_{\text{KL}} = -1.405 + 0.0191 = -1.386$. Negative because the clipped surrogate is positive (rewarded action) and dominates the small positive KL.

Repeating this calculation across all $G \cdot T = 12$ tokens (every token of rollouts 1 and 4 yields a negative loss with magnitude bounded by the clip; rollouts 2 and 3 yield positive losses — the model is penalised for sampling them and will reduce their probability), averaging within each rollout, then averaging across rollouts, produces the scalar loss that gets backpropped through the policy.

The walkthrough headline: the entire GRPO update is computable from the four rewards, the four old-policy log-probs, the four reference log-probs, and the new policy's log-probs — no value head, no MSE target, no GAE bookkeeping. The per-token gradient signal is identical within a rollout (same scalar advantage) and differs across rollouts only by the per-token ratio and KL. That is the entire algorithm.

Interactive Visualization: The Group Baseline

Drag the reward sliders below. Each row is one rollout from the same prompt; the live readouts show the group mean, the group std, and the per-rollout advantage. Toggle divide by std (vanilla GRPO) off to see the unscaled baseline-only variant — the colours stay (same sign of advantage) but the magnitudes spread out wildly. Drag a single reward up until it dominates: every other rollout's advantage flips red because the group mean shifts upward by that single change.

Two experiments worth running in the widget. First: set all four rewards equal (say all to 0.5). The std collapses to zero, the $\epsilon$ in the denominator dominates, every advantage is numerically tiny and the gradient signal vanishes. This is the “std collapse” failure mode and is exactly what happens in a real GRPO run when a binary verifier returns all-correct or all-wrong for a group. Recovery: raise sampling temperature so the rollouts diversify. Second: set one reward to +1 and the other three to 0. The +1 rollout gets a strongly positive advantage, the other three get equal small negative advantages, and the gradient will push the policy hard toward the +1 rollout. This is the canonical “one winner in a group of G” pattern that DeepSeek-R1 reasoning training operates in for hours at a time.

Plain Python: One GRPO Step From Scratch

We implement the GRPO loss in pure NumPy on a toy group: G = 4 rollouts, T = 3 tokens, V = 3 vocabulary entries. No autograd, no transformer — just the arithmetic an RL trainer runs every minibatch. Swap the hard-coded logits for a transformer forward pass and the code is unchanged.

G=4, T=3, V=3 → tokens.shape = (4, 3); rewards.shape = (4,)

rollout 1 sampled the sequence [0, 2, 1]; that is the action sequence we will grade

rewards = [+0.90, +0.30, -0.10, +0.70] → rollout 3 was the worst, rollout 1 the best

mean = (0.90 + 0.30 - 0.10 + 0.70) / 4 = 0.450

advantages = [(0.90-0.45)/0.371, (0.30-0.45)/0.371, (-0.10-0.45)/0.371, (0.70-0.45)/0.371] = [+1.213, -0.404, -1.483, +0.674]

advantages_per_rollout.shape = (4,); after broadcasting, advantages.shape = (4, 3) with each row constant

new_logits.shape = (4, 3, 3); after log_softmax + gather → new_logp.shape = (4, 3)

log pi_new[0,0] = log_softmax(new_logits[0,0])[tokens[0,0]] ≈ -0.207; log pi_old[0,0] = -1.10 → log_ratio = +0.893, ratio = 2.443

ratio[0,0] = 2.443, A[0,0] = +1.213 → surr1 = 2.964, surr2 = 1.2 * 1.213 = 1.456; min = 1.456 (clip fires)

log pi_new[0,0] ≈ -0.207, ref_logp[0,0] = -1.05 → log_r_ref = 0.843, kl = exp(0.843) - 0.843 - 1 = 0.481

per_token_loss.shape = (4, 3); per_rollout_loss.shape = (4,); loss = scalar

1"""
2GRPO loss for ONE prompt with G rollouts — pure NumPy, no autograd.
3
4A 'group' is the G full responses sampled from the OLD policy for the same
5prompt. Each response gets a single scalar reward at the end (from a reward
6model, a verifier, or any oracle). GRPO turns that group of scalar rewards
7into per-TOKEN advantages by subtracting the group mean and dividing by the
8group std. No value network, no GAE, no critic.
9
10We pick a tiny vocab (3 tokens) and G = 4 rollouts of length T = 3 each so
11that every intermediate tensor fits on the page.
12"""
13
14import numpy as np
15
16# ---------------------------------------------------------------------------
17# 1. One group — G = 4 rollouts of length T = 3 tokens each from one prompt.
18#    tokens_i is the action sequence the OLD policy actually sampled.
19#    old_logp_i is log pi_old(a_t | s_t) per token (frozen, no gradient).
20#    ref_logp_i is log pi_ref(a_t | s_t) per token (frozen SFT model).
21# ---------------------------------------------------------------------------
22
23G, T, V = 4, 3, 3                    # group size, tokens per rollout, vocab
24
25tokens   = np.array([
26    [0, 2, 1],                       # rollout 1
27    [1, 0, 2],                       # rollout 2
28    [2, 1, 0],                       # rollout 3
29    [0, 1, 2],                       # rollout 4
30])
31
32old_logp = np.array([
33    [-1.10, -0.92, -1.20],
34    [-0.69, -1.20, -1.10],
35    [-1.00, -0.80, -1.05],
36    [-1.05, -0.95, -1.15],
37])
38
39ref_logp = np.array([
40    [-1.05, -0.90, -1.10],
41    [-0.65, -1.18, -1.05],
42    [-0.98, -0.82, -1.02],
43    [-1.00, -0.92, -1.10],
44])
45
46# ---------------------------------------------------------------------------
47# 2. ONE scalar reward per rollout. In real RLHF this is the reward model's
48#    score; in DeepSeek-R1 it is 1.0 if the math answer is correct, else 0.
49# ---------------------------------------------------------------------------
50
51rewards = np.array([+0.90, +0.30, -0.10, +0.70])     # shape (G,)
52
53# ---------------------------------------------------------------------------
54# 3. GROUP-RELATIVE ADVANTAGE — the heart of GRPO.
55#    For every rollout, A_i = (R_i - mean(R)) / (std(R) + eps).
56#    Broadcast the SAME scalar A_i to all T tokens of that rollout.
57# ---------------------------------------------------------------------------
58
59mean = rewards.mean()
60std  = rewards.std()
61eps  = 1e-8
62
63advantages_per_rollout = (rewards - mean) / (std + eps)     # shape (G,)
64advantages = np.broadcast_to(
65    advantages_per_rollout[:, None],                        # (G, 1)
66    (G, T),                                                 # (G, T)
67).copy()
68
69# ---------------------------------------------------------------------------
70# 4. The NEW policy's logits at the SAME tokens (gradient flows here in
71#    real code). Hard-coded "drifted" logits make the walkthrough readable.
72# ---------------------------------------------------------------------------
73
74new_logits = np.array([
75    [[ 1.50, -0.10,  0.30], [ 0.20,  0.30,  1.50], [-0.10,  1.60,  0.20]],
76    [[ 0.10,  1.50,  0.30], [ 1.40, -0.10,  0.30], [ 0.20,  0.30,  1.20]],
77    [[-0.10,  0.30,  1.50], [ 0.30,  1.40, -0.10], [ 1.50, -0.10,  0.30]],
78    [[ 1.40, -0.10,  0.30], [ 0.20,  1.30, -0.10], [-0.10,  0.30,  1.40]],
79])
80
81# Stable log-softmax → log pi_new for the action that was actually taken.
82z = new_logits - new_logits.max(axis=-1, keepdims=True)
83log_probs_all = z - np.log(np.exp(z).sum(axis=-1, keepdims=True))
84new_logp = np.take_along_axis(
85    log_probs_all, tokens[..., None], axis=-1,
86).squeeze(-1)                                               # shape (G, T)
87
88# ---------------------------------------------------------------------------
89# 5. Probability ratio r_t = pi_new / pi_old, computed in log space.
90# ---------------------------------------------------------------------------
91
92log_ratio = new_logp - old_logp
93ratio     = np.exp(log_ratio)                               # shape (G, T)
94
95# ---------------------------------------------------------------------------
96# 6. PPO-style clipped surrogate — but with the GROUP advantage.
97# ---------------------------------------------------------------------------
98
99EPSILON = 0.2
100surr1 = ratio * advantages
101surr2 = np.clip(ratio, 1.0 - EPSILON, 1.0 + EPSILON) * advantages
102clipped_surrogate = np.minimum(surr1, surr2)
103
104# ---------------------------------------------------------------------------
105# 7. k3 KL estimator to the REFERENCE policy (the SFT model), per token.
106#    kl ≈ exp(log_r) - log_r - 1; always positive, unbiased in expectation.
107# ---------------------------------------------------------------------------
108
109log_r_ref  = new_logp - ref_logp
110kl_per_tok = np.exp(log_r_ref) - log_r_ref - 1              # shape (G, T)
111
112# ---------------------------------------------------------------------------
113# 8. Per-token loss, then sequence-mean within each rollout, then group-mean
114#    across rollouts. Two means are deliberate — see "length bias" below.
115# ---------------------------------------------------------------------------
116
117BETA_KL = 0.04
118
119per_token_loss = -clipped_surrogate + BETA_KL * kl_per_tok  # (G, T)
120per_rollout_loss = per_token_loss.mean(axis=-1)             # mean over T → (G,)
121loss = per_rollout_loss.mean()                              # mean over G → ()
122
123print(f"advantages = {advantages_per_rollout.round(3)}")
124print(f"ratios     =\n{ratio.round(3)}")
125print(f"loss       = {loss:.4f}")

PyTorch: A Production-Shaped GRPO Step

Now the full thing — clipped surrogate, per-token KL to a reference model, gradient clipping, sequence-mean then batch-mean, and the four diagnostics that tell you whether the run is healthy. This is the inner loop that lives inside OpenRLHF's GRPO trainer, TRL's GRPOTrainer, and the DeepSeek-R1 reference code. Drop in any HF causal LM for policy and any frozen copy of the SFT model for ref_policy and this trains a real GRPO run.

If batch has N=4 prompts × G=16 rollouts → B=64 sequences; input_ids.shape = (64, S)

B=64, GROUP_G=16 → N=4; r.shape = (4, 16); r_mean.shape = (4, 1); r_std.shape = (4, 1)

adv_per_seq.shape = (4, 16) → reshape to (64, 1) → expand to (64, S). On row 0 every column is the same scalar.

Typical first GRPO step: log_ratio ≈ 0 ± 0.05 → ratio ≈ 1.0 ± 0.05

Healthy step: approx_kl ≈ 0.02, clipfrac ≈ 0.18, kl_ref ≈ 12 nats, adv_var ≈ 0.5

1"""
2PyTorch GRPO step — the inner loop of DeepSeek-R1, Qwen-2 math, Llama-3 RLHF
3when run in critic-free mode. Three models live in memory (policy, ref,
4no critic, no reward model on the training cluster).
5
6A 'rollout batch' is a list of N prompts, each rolled out to G full responses
7by the OLD policy (frozen snapshot). For each (prompt, response) token we have:
8   - old_logp      : log pi_old(a_t | s_t)            (no gradient)
9   - ref_logp      : log pi_ref(a_t | s_t)            (no gradient)
10   - reward_per_seq : scalar reward per FULL response (no gradient)
11
12K epochs of mini-batched updates on this same rollout follow. Each minibatch
13re-forwards the rollout tokens through the LIVE policy and computes the loss
14below.
15"""
16
17import torch
18import torch.nn.functional as F
19
20EPSILON  = 0.2          # PPO-style ratio clip (same default as PPO)
21BETA_KL  = 0.04         # per-token KL to reference; DeepSeek-Math uses 0.04
22MAX_GRAD = 1.0          # global grad-norm clip
23GROUP_G  = 16           # G rollouts per prompt
24
25def grpo_step(batch, policy, ref_policy, optimizer):
26    """
27    batch is one minibatch from the rollout, already on device. Tensors are
28    grouped: B = N_prompts * G, then reshape (N, G, S) to compute group stats.
29
30      input_ids       : (B, S)        prompt + response tokens
31      attn_mask       : (B, S)
32      action_mask     : (B, S)        1 on RESPONSE tokens, 0 on prompt/pad
33      old_logp        : (B, S)        log pi_old at each token (frozen)
34      ref_logp        : (B, S)        log pi_ref at each token (frozen)
35      reward_per_seq  : (B,)          one scalar per (prompt, response)
36      prompt_id       : (B,)          which prompt each sequence belongs to
37    """
38    input_ids      = batch["input_ids"]
39    attn_mask      = batch["attn_mask"]
40    action_mask    = batch["action_mask"]
41    old_logp       = batch["old_logp"]
42    ref_logp       = batch["ref_logp"]
43    reward_per_seq = batch["reward_per_seq"]
44    B, S = input_ids.shape
45
46    # --- Group-relative advantage (the GRPO core) ------------------------
47    # Reshape (B,) -> (N, G) so each row is one prompt's G rollouts.
48    assert B % GROUP_G == 0, "batch must be a whole number of groups"
49    N = B // GROUP_G
50    r = reward_per_seq.view(N, GROUP_G)                          # (N, G)
51    r_mean = r.mean(dim=1, keepdim=True)                         # (N, 1)
52    r_std  = r.std(dim=1, keepdim=True, unbiased=False)          # (N, 1)
53    adv_per_seq = (r - r_mean) / (r_std + 1e-8)                  # (N, G)
54    advantages  = adv_per_seq.view(B, 1).expand(B, S)            # (B, S)
55
56    # --- Forward the LIVE policy on the rollout --------------------------
57    out = policy(input_ids=input_ids, attention_mask=attn_mask, use_cache=False)
58    logits = out.logits[:, :-1, :]                               # predict t+1 from <=t
59    targets = input_ids[:, 1:]
60    am      = action_mask[:, 1:].float()
61
62    log_probs_all = F.log_softmax(logits, dim=-1)
63    new_logp = log_probs_all.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
64
65    # --- PPO-style clipped surrogate, group advantage --------------------
66    log_ratio = new_logp - old_logp[:, 1:]
67    ratio     = log_ratio.exp()
68
69    surr1 = ratio * advantages[:, 1:]
70    surr2 = torch.clamp(ratio, 1 - EPSILON, 1 + EPSILON) * advantages[:, 1:]
71    policy_objective = torch.min(surr1, surr2)
72
73    # --- k3 KL to reference, per token -----------------------------------
74    log_r_ref   = new_logp - ref_logp[:, 1:]
75    kl_per_tok  = (log_r_ref.exp() - 1) - log_r_ref              # ≥ 0
76    per_token   = -policy_objective + BETA_KL * kl_per_tok       # (B, S-1)
77
78    # --- Two means: sequence-mean THEN group/batch-mean ------------------
79    # Masked mean over tokens within each sequence.
80    seq_len     = am.sum(dim=-1).clamp(min=1.0)                  # (B,)
81    seq_loss    = (per_token * am).sum(dim=-1) / seq_len         # (B,)
82    loss        = seq_loss.mean()                                # scalar
83
84    # --- One gradient step ------------------------------------------------
85    loss.backward()
86    gn = torch.nn.utils.clip_grad_norm_(policy.parameters(), MAX_GRAD)
87    optimizer.step()
88    optimizer.zero_grad(set_to_none=True)
89
90    # --- Diagnostics — the four numbers an RL engineer watches -----------
91    with torch.no_grad():
92        approx_kl = ((old_logp[:, 1:] - new_logp) * am).sum() / am.sum().clamp(min=1.0)
93        clipfrac  = (((ratio - 1.0).abs() > EPSILON).float() * am).sum() / am.sum().clamp(min=1.0)
94        kl_ref    = (kl_per_tok * am).sum() / am.sum().clamp(min=1.0)
95        adv_var   = adv_per_seq.var(dim=1).mean()                # within-group spread
96
97    return {
98        "loss":       loss.detach(),
99        "kl_ref":     kl_ref.detach(),
100        "approx_kl":  approx_kl.detach(),
101        "clipfrac":   clipfrac.detach(),
102        "grad_norm":  gn.detach(),
103        "adv_var":    adv_var.detach(),
104        "reward":     reward_per_seq.mean().detach(),
105    }

At Massive Scale: What Happens at 70B+

The same code scales to a 70B GRPO run, but three forces change in spirit.

Memory: the value head was half the cost

At 70B, each model is ~140 GB in bf16. PPO needs four such models (policy, value, ref, reward) plus optimiser state for the two trainable ones — about 2 TB before activations. GRPO needs two (policy, ref) plus optimiser state for just the policy — about 980 GB. Cutting the trainable-model count from 2 to 1 is the savings that matters: AdamW state is ~12 N bytes per trainable parameter (fp32 master copy plus fp32 m and v), so eliminating the value head's optimiser state alone saves ~840 GB. This is what makes 70B GRPO runs feasible on the same hardware that struggles with 70B PPO-RLHF.

The rollout-vs-train balance shifts further toward rollout

PPO already spent ~85% of wall-clock time on rollout (§ 14.6). GRPO multiplies the rollout cost by G — for G = 16 a single prompt now demands 16 full responses. The optimisation phase, on the other hand, is cheaper because there is no value-head gradient and no value-loss computation. Production GRPO systems push the rollout fraction to 90%+ and architect around that: dedicated vLLM rollout clusters, async weight syncs, and batched-prompt servers that emit all G rollouts of a prompt in one forward call. The architectural choice that PPO made under pressure (separate rollout and training clusters) becomes the only viable choice in GRPO.

Distributed group statistics: an all-reduce per group

If the G rollouts of a single prompt are split across FSDP ranks, the group mean and std must be computed with an all-reduce within each group's rank set, not over the global batch. In practice every production GRPO collator keeps each prompt's G rollouts on a single rank so the group statistics are a local operation — the “contiguous-group” layout rule. Get this wrong and you compute statistics over a mixture of prompts, which is mathematically a different and much worse algorithm. The first thing to assert in any custom GRPO collator is prompt_id[i*G:(i+1)*G].unique().numel() == 1 for every group.

Reward bottleneck disappears when the reward is a verifier

For DeepSeek-R1 the reward is not a learned model — it is a deterministic verifier (math answer checker, code executor, JSON validator). That verifier runs on CPU in microseconds, with no GPU memory at all. The “reward model bottleneck” that defines PPO-RLHF (the fourth 70B model that must be queried per rollout) simply vanishes. This is why GRPO + verifier is the dominant recipe for reasoning RL in 2024-2025: the cost structure is fundamentally simpler than PPO-RLHF, and the per- example signal-to-noise is higher because verifiers are not subject to reward hacking.

Sampling temperature becomes a load-bearing hyperparameter

In PPO, sampling temperature affects exploration but not the algorithm's stability. In GRPO it affects both. If temperature is too low, every rollout in a group is nearly identical, the reward variance collapses, the std $\sigma_R \to 0$, and the gradient signal vanishes. If temperature is too high, rollouts become garbage, rewards are uniformly low, and the policy gets pushed toward those uniform-low responses. The DeepSeek-Math sweet spot is $T = 0.6$ to $T = 1.0$ with a small top-p cutoff; every production GRPO trainer logs adv_var per step and is prepared to bump temperature if it collapses.

Engineering Reality: Length Bias, Std Collapse, and Other Knobs

After a season of GRPO runs every team converges on the same set of failure modes and the same set of fixes. The list is shorter than PPO's because the algorithm has fewer moving parts.

• Std collapse. Symptom: training stalls and adv_var drops to ~1e-6 in the diagnostics. Cause: every rollout in a group received the same reward (binary verifier returns all-correct or all-wrong), so $\sigma_R \to 0$ and the advantage collapses to a tiny ratio governed by the $\epsilon$ floor. Fix: raise sampling temperature so the rollouts diversify; or use a richer reward (partial credit instead of binary) so equal scores are rarer; or filter out collapsed groups entirely from the loss (most production trainers do this).
• Length bias. Symptom: the policy converges to short, terse responses regardless of prompt. Cause: vanilla GRPO's sequence-mean normalisation upweights short rollouts in the loss; a short rewarded rollout contributes more per-token gradient than a long rewarded one, so the policy learns to prefer brevity. Fix: switch to Dr. GRPO (sum over tokens divided by a fixed max-length constant); or add an explicit length-target reward.
• KL exploding. Symptom: kl_ref grows above 100 nats per response and never recovers, eval quality collapses. Cause: $\beta_{\text{KL}}$ is too small for the chosen reward signal. Fix: raise to 0.05 – 0.1 and restart from a pre-collapse checkpoint. Many teams use an adaptive controller that adjusts $\beta_{\text{KL}}$ to hold realised KL at a target value (e.g. 20 nats).
• KL stuck near zero. Symptom: reward barely moves, kl_ref stays under 1 nat. Cause: policy is not changing — either LR is too low, or grad_norm is being clipped to near zero, or there is a reference-vs-policy aliasing bug (both tensors point to the same weights). Fix: assert id(ref_policy.state_dict()) != id(policy.state_dict()) and that the reference's parameters have requires_grad=False.
• Group layout wrong. Symptom: reward seems to go up but eval gets worse. Cause: the batch is not laid out contiguously by prompt, so the group statistics are computed over a mixture of prompts. Fix: explicitly enforce contiguous- group layout in the collator and assert that prompt_id[i*G:(i+1)*G] is a single value for every $i$.
• Clipfrac too high. Symptom: clipfrac > 0.5 per step. Cause: the policy is moving too fast — usually too many GRPO epochs per rollout (K > 2) or LR too high. Fix: drop K to 1 (DeepSeek-R1 uses a single epoch) or halve LR. The clip is doing its job, but you are wasting compute on gradients that are being clipped to zero.
• Rollout decoding mismatch. Symptom: ratio distribution centred far from 1.0 on the very first step before any gradient has been taken. Cause: the rollout used a sampler that the trainer cannot reproduce — top-p, temperature, or stop-token logic differs between vLLM and the live forward pass. Fix: greedy reconstruction is fine for log-prob purposes; just ensure the trainer's logit transform matches what the rollout engine emitted. trl, OpenRLHF, and DeepSpeed-Chat all have explicit code for this.
• Memory blowup on the reference model. Symptom: OOM after enabling reference KL. Cause: holding two full bf16 copies of a 70B policy on the same GPUs. Fix: shard the reference with FSDP or place it on CPU and run reference forwards once per rollout (it is not in the gradient path, so the latency hit is tolerable). LoRA tricks that share weights between the policy and the reference (with a different adapter for each) are increasingly common.

The mental model that unifies the section: GRPO is PPO with one substitution. Replace the value head with the group mean. Replace the value-head's normalisation with the group std. Everything else — the importance ratio, the clipped surrogate, the k3 KL to a reference — survives unchanged. Two models live in memory instead of four. The cost moves from per-step trainable parameters to per-prompt inference samples. For LLM tasks where the per-prompt rollout is cheap (vLLM, async batching) and the reward is a verifier (no learned reward model), GRPO is strictly better than PPO. For everything else it is a one-line decision: how much variance can you tolerate in exchange for a free 70B model?