1️⃣ RLHF via GRPO

Learning Objectives

Understand how the RL agent / action / environment paradigm works in the context of autoregressive language models

Understand GRPO: group-relative advantages, the clipped surrogate objective, and the KL penalty to a reference model

Understand why we estimate the KL penalty from sampled tokens

Write a full RLHF training loop, and train against a toy reward function.

Observe and understand reward hacking and mode collapse

Experiment with a more complex reward function (a sentiment classifier)

The language model "environment"

We'll start by discussing how we apply the reinforcement learning framework of states/actions/rewards to the setting of autoregressive language modelling.

States, actions and episodes

Our policy $\pi_\theta$ is an autoregressive language model. The actions $a_t$ are the tokens generated by the model (i.e. the action space, $\mathcal{A}$, is the model's vocabulary). The initial state $s_0$ is the prompt to the model, and the states $s_t$ are the entire sequence up to that point (not just the most recent token). In other words, given a state $s_t := s_0 a_0 \ldots a_{t-1}$ (sequence) and sampled action $a_t \sim \pi_\theta(\cdot | s_t)$ (token generation), the dynamics of the environment is just concatenation: $s_{t+1} := s_t \; a_t$. For every timestep before the end of the episode, the reward is zero, and for the final timestep, the reward $r_T$ is given by the reward function applied to the entire sequence, so we get $r_T = R(s_T) \equiv R(s_0 a_0 \ldots a_{T-1})$. The environment is obviously deterministic and the dynamics aren't interesting, all of the complexity of this problem is finding a sequence of tokens $a_0 \ldots a_{T-1}$ that the reward function $R$ finds pleasing.

In the strictest of senses this environment is "Markovian", but only because we have a state for every possible sequence of tokens, so it's unlikely we'll ever visit the same state twice. It's better to think of this as a non-Markovian environment where the reward function depends on the entire history of what the model generated.

Each episode is a fixed length (i.e. all our sampled outputs will have the same number of tokens generated from them), and starts from the same prompt, chosen before training. Since the only reward comes at the end of an episode of fixed length, discounting would only scale the reward by a fixed constant, so we don't discount. One may wish to use a discount if the generation length was variable, which would all-else-equal encourage the model to give brief responses that still score high reward.

Why would discounting only scale the final reward by a fixed constant?

The expected return for the initial state is

\mathbb{E}[G_0] = \mathbb{E}[r_1 + \gamma r_2 + \ldots + \gamma^{T-1} r_T] = \gamma^{T-1} \mathbb{E}[r_T]

as all other rewards $r_1, \ldots, r_{T-1}$ are zero. Since $T$ is fixed, $\gamma^{T-1}$ is a constant, so we may as well drop it.

Note about variable generation length

Note that this isn't true if the generation length is variable: We might give the agent a budget of $N$ tokens to generate, and at any point it can sample the end-of-text token early, in which case we stop generation and then query the reward model. If the reward function is a hard 0 or 1 reward for "did the model complete the task?" then discounting would help coax the model towards brevity. It might also provide the agent some softer feedback towards brevity rather than a hard cut-off at maximum generation length.

Note - a key part of RLHF is the actual gathering of and learning from human feedback, in order to train the reward function. We're not going to be doing that here, instead we'll be working with a fixed reward function (see the bonus section for training your own). So today is really "RL fine-tuning of language models", which is also how models are trained on tasks with automatically checkable answers like maths and code (often called RLVR, RL from verifiable rewards).

Policy gradients and baselines

Our goal is to maximize the expected reward

J(\theta) = \mathbb{E}_{o \sim \pi_\theta(\cdot \mid q)}[R(q,o)]

where $q$ is a prompt and $o = a_0 \ldots a_{T-1}$ is the generated completion (we'll use the GRPO paper's notation $q$ for "question" and $o$ for "output" from here on). The policy gradient theorem tells us that

\nabla_\theta J(\theta) = \mathbb{E}_{o \sim \pi_\theta}\left[ \sum_{t=0}^{T-1} A_t \, \nabla_\theta \log \pi_\theta(o_t \mid q, o_{\lt t}) \right]

where the advantage $A_t$ measures how much better the outcome was than expected. The simplest choice is $A_t = R(q, o)$, i.e. push up the log-probability of every token in a sequence in proportion to its reward. (Recall that we don't discount, and that we only score the completion at the end of generation, so this single reward is the same as the return $G_t$ for every timestep $t$.)

But this has a problem: if all the rewards are positive, we push every sampled sequence up, and rely on the good ones being pushed up more. We instead want to push a sequence up or down, depending how good it was compared to what else we could have done. We can subtract any baseline $b(q)$ which doesn't depend on the action without changing the expected gradient, and a good baseline (close to the expected reward) reduces the variance a lot.

PPO solved this by using the value $V(s_t)$ as the baseline, and trained a critic to guess $V(s_t)$. This can be achieved by grafting a value head that reads activations from the residual stream late in the model (which already contain a rich representation of the features in the text generated), and learns to map this via a MLP to a single scalar, that is then trained to predict the empirically obtained returns).

But the critic increases the parameter count, it starts randomly initialized, and estimating the value of a half-completed sentence is hard. GRPO does something much simpler: for each prompt $q$, it samples a group of $G$ completions $o_1, \ldots, o_G$, and uses the group's mean reward as the baseline:

\hat{A}_{i} = \frac{r_i - \operatorname{mean}(r_1, \ldots, r_G)}{\operatorname{std}(r_1, \ldots, r_G)}

We also divide by the standard deviation of the rewards to adjust for scale (the learning dynamics shouldn't change if we scaled up all the rewards by a factor of 10. The absolute value of the reward doesn't matter, only that we are getting the most reward that we can, compared to what else we could be doing). Note that every token of the completion $o_i$ gets the same advantage $\hat{A}_i$. (This is fine for now, but more complex methods try to attribute credit to a particular token/sentence of the rollout, to try and reward the critical or most important steps of, say, a mathematical proof.) A completion that did better than the average of all completions gets its tokens reinforced, one that did worse gets discouraged.

This approach does have a downside, namely that it must be the case that the rollout of groups is diverse. If every rollout is scored the same (suppose the reward is 0/1 for incorrect/correct), then if we get a set of rollouts where if every rollout is wrong, or if every rollout is correct, the entire group has no effect during training, and is effectively skipped. The agent needs a sense of better v.s. worse responses to make any progress.

The GRPO objective

Here is the GRPO objective function (to be maximized), from section 4.1 of the DeepSeekMath paper:

J_{\text{GRPO}}(\theta) = \mathbb{E}\left[ \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \left( \min\Big[ \rho_{i,t} \, \hat{A}_{i}, \; \text{clip}(\rho_{i,t}, 1-\epsilon, 1+\epsilon)\,\hat{A}_{i} \Big] \; - \beta \, D_{\mathrm{KL}}\!\left[\pi_\theta \,\|\, \pi_{\text{ref}}\right] \right) \right], \qquad \rho_{i,t} = \frac{\pi_\theta(o_{i,t} \mid q, o_{i,\lt t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,\lt t})}

There are three ingredients, each of which you'll implement:

Group-relative advantages $\hat{A}_i$, as described above.
The clipped surrogate objective $\min[\rho \hat{A}, \text{clip}(\rho) \hat{A}]$. Generating samples is expensive, so we want to take several gradient steps on each batch of samples. But after the first step, the policy $\pi_\theta$ we're updating is no longer the policy $\pi_{\theta_\text{old}}$ that generated the samples. The importance ratio $\rho$ corrects for this, but it can explode or collapse as $\pi_\theta$ drifts away from $\pi_{\theta_\text{old}}$. Clipping stops us from moving too far from $\pi_{\theta_\text{old}}$ on data generated from it (the further we move from $\pi_{\theta_\text{old}}$, the less reinforcing the rollout makes sense, as it is now a rollout that the current model might not generate anyway).
The KL penalty $\beta D_{KL}[\pi_\theta \| \pi_\text{ref}]$, which stops the policy from drifting too far from the original pretrained model $\pi_\text{ref}$. We want to maximize reward, but not in a way which causes the model to generate out-of-distribution garbage. This is especially important if the reward model is simpler than the policy, and is easily tricked into giving high rewards to incoherent outputs.

The training loop alternates between two phases, just like PPO:

Rollout phase: sample a group of completions of the prompt from the current policy, compute their rewards and advantages, and record the log-probabilities of the sampled tokens under both the current policy (which becomes $\pi_{\theta_\text{old}}$) and under the original model $\pi_{\text{ref}}$.
Learning phase: split these samples into minibatches, and take a gradient ascent step on the objective above for each.

A note on how this maps onto our setup. In the paper, each training batch contains groups for many different prompts. To keep things simple, we'll use a single prompt for all our rollouts, so the whole batch is one group (the period-counting reward doesn't care much what the prompt is anyway). And all our completions have exactly gen_len tokens, so the $\frac{1}{G} \sum_i \frac{1}{|o_i|} \sum_t$ is just the mean over all generated tokens.

Reverse KL Penalty

Note that we compute $D_{KL}(\pi_\theta \| \pi_\text{ref})$, not the other way around. This is because we want to penalize our new model for generating outputs which would be extremely unlikely under the old model, i.e. when $\pi_\theta$ is high and $\pi_\text{ref}$ is low. We generally want to focus our model's output into a more concentrated version of the distribution it already has. For example in RLHF, we want to keep a low probability on completely incoherent behaviour which the original model would never have generated. But on the other hand, it's clearly fine for there to be some behaviours (e.g. violent or sexual content) which have a nontrivial probability in our base model but near-zero probability in our new model, which is exactly how we want the new model to act! For more on the intuition behind this orientation of the distributions in KL divergence, see this post.

KL divergence v.s. reverse KL divergence

Assume $P$ is the true distribution, and $Q$ is the distribution we're trying to fit to $P$.

$D_{KL}(P || Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}$ blows up when $Q(x)$ is zero and $P(x)$ is positive, so we would expect that $Q$ tries to "cover" $P$ anywhere where $P(x)$ is positive. This means that minimizing $D_{KL}(\pi_\text{ref} || \pi_\theta)$ will cause our model to be able to do everything the base model can do, plus it can also do things out-of-distribution for the base model, which is undesirable.
$D_{KL}(Q || P) = \sum_x Q(x) \log \frac{Q(x)}{P(x)}$ blows up when $P(x)$ is zero and $Q(x)$ is positive, so $Q$ should never assign any probability mass to something that $P$ doesn't ($P$ "covers" $Q$), but $Q$ will instead try to cover a subset of $P$ that it fits the best.

This can be illustrated with an example. Let $P$ be a mixture of two Gaussians, and $Q \sim \mathcal{N}(\mu, \sigma^2)$ be a unimodal Gaussian (blue) parameterized by $\mu$ and $\sigma^2$. We learn parameters $\mu,\sigma^2$ that minimize both $D_{KL}(P || Q)$ and $D_{KL}(Q || P)$, and draw the resulting distribution $Q$ (here in blue), showing the expected behaviour: the reverse KL divergence tries to lock onto a subset of the true distribution $P$ that it fits the best. This is the target for RLHF: the tuned-model focuses on a subset of the behaviours of the base model that maximise the return according to what the reward function values.

GRPO training args

Here are the arguments we'll use for training. A few of them deserve comment:

prompt and batch_size: each rollout phase samples batch_size completions of prompt. These form a single GRPO group, so batch_size is the group size $G$.
total_phases controls how long training lasts for. Each phase is one rollout phase followed by one learning phase, in which we take batches_per_learning_phase * num_minibatches gradient steps.
gen_len is the number of tokens we generate for each completion, and temperature is the sampling temperature.
kl_coef is the coefficient $\beta$ of the KL penalty, and clip_coef is the clipping parameter $\epsilon$.

@dataclass
class GRPO_Args:
    # Basic / global
    seed: int = 1

    # Wandb / logging
    use_wandb: bool = False
    wandb_project_name: str = "RLHF"
    wandb_entity: str | None = None

    # Duration of different phases
    total_phases: int = 100
    prompt: str = "This is"
    batch_size: int = 128
    num_minibatches: int = 4
    batches_per_learning_phase: int = 2

    # Optimization hyperparameters
    lr: float = 1e-5
    max_grad_norm: float = 1.0
    warmup_steps: int = 10
    final_scale: float = 0.1

    # GRPO objective
    clip_coef: float = 0.2
    kl_coef: float = 0.3

    # Base model & sampling arguments
    base_model: str = BASE_MODEL
    gen_len: int = 30
    temperature: float = 1.0

    # Reward function, mapping a list of samples (prompt + completion) to a tensor of rewards
    reward_fn: Callable = lambda samples: t.zeros(len(samples), device=device)

    def __post_init__(self):
        assert self.total_phases > self.warmup_steps, "total_phases must be greater than warmup_steps"
        assert self.batch_size % self.num_minibatches == 0, "batch_size should be divisible by num_minibatches"
        self.minibatch_size = self.batch_size // self.num_minibatches

Sampling from a language model

If you didn't go through the sampling exercises in section [1.1] Transformer from Scratch, you might want to go back to them and work through the first few of them (this is not essential). Otherwise, here's a quick refresher:

Sampling methods

The simplest form of sampling is greedy sampling, where we autoregressively generate text by always choosing the most likely token at each step (i.e. argmaxing over logits), appending this to our sequence, and continuing.
Most other forms of sampling are non-deterministic, i.e. they involve randomness. The most basic form of random sampling is choosing the next token according to the model's logit distribution.
Other common refinements of this basic method are:
- Top-k sampling, where we only consider the top-k most likely tokens at each step, and choose from these according to the model's logit distribution.
- Top-p sampling (also called nucleus sampling), where we only consider the most likely tokens that have cumulative probability at least $p$ at each step, and choose from these according to the model's logit distribution.

We've provided the sampling code for you below. Make sure you understand it before moving on, since there are a few things specific to our situation:

We use Hugging Face's model.generate, and pass top_k=0 and top_p=1.0 to turn off truncated sampling. With temperature=1 (the default), the tokens then come from exactly $\pi_\theta$, which is what the policy gradient assumes. Truncated sampling methods like top-k would sample from a different distribution to the one whose log-probabilities we train on.
We always generate exactly gen_len new tokens: passing eos_token_id=None stops generate from ending a sequence early at the end-of-text token, so every episode has the same length.
We start every prompt with the end-of-text token <|endoftext|>. For a base model this marks the start of a new document, just as it did for GPT-2.
Every sequence in the batch starts from the same prompt, so they're all the same length and need no padding. The generated tokens of every sequence are its last gen_len tokens.

@t.no_grad()
def get_samples(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    prompt: str,
    batch_size: int,
    gen_len: int,
    temperature: float = 1.0,
) -> tuple[Int[Tensor, "batch seq"], list[str]]:
    """
    Samples `batch_size` completions of `gen_len` tokens, all from the same prompt.

    Args:
        model: the language model to sample from
        tokenizer: the model's tokenizer
        prompt: the prompt to complete
        batch_size: the number of completions to sample
        gen_len: the number of new tokens to generate for each completion
        temperature: the sampling temperature (1.0 means sampling from the model's distribution)

    Returns:
        sample_ids: the token ids of the samples (prompt + completion). The completions are the last `gen_len` tokens.
        samples: the samples as strings (prompt + completion, without special tokens), used for the reward
    """
    prompt_ids = tokenizer(tokenizer.eos_token + prompt, return_tensors="pt").input_ids.to(device).repeat(batch_size, 1)
    sample_ids = model.generate(
        prompt_ids,
        # Without an explicit attention mask, `generate` treats every pad token in the prompt as padding. Our pad token
        # is the end-of-text token that starts the prompt, so it would mask that out and sample from a different
        # context to the one we compute log-probabilities on.
        attention_mask=t.ones_like(prompt_ids),
        do_sample=True,
        temperature=temperature,
        top_k=0,  # no truncation: sample from the model's full distribution
        top_p=1.0,
        max_new_tokens=gen_len,
        eos_token_id=None,  # don't stop at the end-of-text token: always generate exactly gen_len tokens
        pad_token_id=tokenizer.eos_token_id,
    )

    completions = tokenizer.batch_decode(sample_ids[:, -gen_len:], skip_special_tokens=True)
    samples = []
    for completion in completions:
        samples.append(prompt + completion)
    return sample_ids, samples

Here's some example use of this function.

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, dtype=t.float32).to(device)

sample_ids, samples = get_samples(model, tokenizer, prompt="So long, and thanks for all the", batch_size=5, gen_len=15)

table = [[str(ids.tolist()), repr(sample)] for ids, sample in zip(sample_ids, samples)]
print(tabulate(table, ["Token ids", "Sample"], tablefmt="simple_grid", maxcolwidths=60))

Exercise - implement `reward_fn_char_count`

Difficulty: 🔴⚪⚪⚪⚪

Importance: 🔵🔵⚪⚪⚪

You should spend 5-10 minutes on this exercise.

We'll start with a very basic reward function: counting the total number of a given character (by default, periods) in the sequence.

An interesting thing to note about this reward function - it counts over all characters, but the episode length is defined in terms of tokens. This means that theoretically our model could reward hack by outputting tokens with more than one . character. The GPT-2 tokenizer vocabulary includes tokens of up to 64 periods:

>>> tokenizer.convert_tokens_to_ids("." * 64)
23193

so rewards would be through the roof if these were generated! However, remember that RL is about performing actions, getting feedback on those actions and using that feedback to influence your policy. These tokens are so unlikely to be generated after an ordinary prompt that they'll probably never be sampled and thus never positively reinforced... at least at first. Keep an eye out for them in your training runs. As the model learns to emit periods, it will probably stumble onto the two period token, which internally makes the model more likely to emit other similar period tokens.

Why would the tokenizer have a token specifically for 64 periods?

Lots of the training corpus had long strings of periods as separators in the text (think tables of contents), and so the byte-pair encoding algorithm learned tokens for strings of periods of various lengths, so it can construct a string of periods of any length with few tokens.

If we were worried about this, we could instead have the reward be the number of tokens that contain at least one ., or normalize over character count, or something similar. For now, this is just an easy reward function that we can use to quickly verify that our RLHF trainer is working.

def reward_fn_char_count(generated_sample: list[str], char: str = ".") -> Float[Tensor, " batch"]:
    """
    Reward function counting the number of instances of a particular character.

    Args:
        generated_sample: the generated samples to evaluate
        char: the character to count

    Returns:
        A float tensor of shape (batch,) giving the count of `char` in each sample.
    """
    raise NotImplementedError()


# Test your reward function
A = "This is a test."
B = "......"
C = "Whatever"

t.testing.assert_close(reward_fn_char_count([A]), t.tensor([1.0], device=device))
t.testing.assert_close(reward_fn_char_count([A, B, C]), t.tensor([1.0, 6.0, 0.0], device=device))
t.testing.assert_close(reward_fn_char_count([A], " "), t.tensor([3.0], device=device))
print("All tests for `reward_fn_char_count` passed!")

Solution

def reward_fn_char_count(generated_sample: list[str], char: str = ".") -> Float[Tensor, " batch"]:
    """
    Reward function counting the number of instances of a particular character.

    Args:
        generated_sample: the generated samples to evaluate
        char: the character to count

    Returns:
        A float tensor of shape (batch,) giving the count of `char` in each sample.
    """
    return t.tensor([item.count(char) for item in generated_sample], device=device, dtype=t.float)

Exercise - brainstorm your reward function

Difficulty: 🔴🔴⚪⚪⚪

Importance: 🔵🔵🔵⚪⚪

You should spend ~5 minutes on this exercise.

Take 5 minutes (on your own or with a partner) to brainstorm how the model might be able to maximize the output of periods in ways which don't produce incoherent output (e.g. collapsing into only outputting periods). Remember we have a KL penalty with the reference model, meaning the model is penalized for producing outputs which would be very unlikely under the original model. What ideas can you come up with? When you train your model and observe the output, you should come back here and see how many of the period-maximizing behaviours you predicted actually occur.

This exercise is a great way to start thinking about the effects of different reward functions - although it's only a toy example, it still illustrates the important alignment concept that the behaviour induced by certain reward functions might not always be what you expect!

Spoiler - which behaviours will your model pick up?

The strategies adopted by the model vary a lot depending on the prompt, and thanks to mode collapse it will often find one of these behaviours and entirely ignore the others.

Some common strategies include:

Shorter sentences
Repeating U.S. or U.S.A.
Library versions e.g. Python 2.7.12 or the 2.6.0.2 release
Names with initials e.g. C. S. Lewis or titles e.g. Dr. and PhD.
Abbreviations e.g. Data-R.A.R. series or "L.A. Times"
Decimals in numbers e.g. 9.5cm x 7.5 cm
Ellipses e.g. the man . . . the woman . . .

Exercise - implement `compute_grpo_advantages`

Difficulty: 🔴⚪⚪⚪⚪

Importance: 🔵🔵🔵🔵⚪

You should spend up to 10-15 minutes on this exercise.

Implement the group-relative advantages from the GRPO paper:

\hat{A}_{i} = \frac{r_i - \operatorname{mean}(r_1, \ldots, r_G)}{\operatorname{std}(r_1, \ldots, r_G) + \epsilon}

where the mean and standard deviation are taken over the group. Since all our rollouts share a single prompt, the group is the whole batch. This gives one advantage per sequence; we'll apply it to every token of the sequence later.

What happens to the advantages if every sample got the same reward? Make sure your function returns zeros rather than NaNs in this case.

Answer

If every sample gets the same reward, they all get advantage zero: there's no signal about which completions were better, so the batch contributes nothing to the gradient. (This is why the $\epsilon$ is needed: it avoids $0 / 0$.) In practice, RL on hard tasks wastes a lot of samples this way, on prompts the model always fails or always succeeds at.

What about multiple prompts?

In practice we'd train on many prompts at once, sampling a group of completions for each, and normalize within each group: the mean reward for a prompt is our estimate of how well the policy does on that prompt. If we normalized over the whole batch instead, every sample from an easy prompt would get a positive advantage (and every sample from a hard prompt a negative one), so we'd mostly be learning "easy prompts are good", which isn't something the policy can act on. You'll see this in the RLVR material, which trains on many different problems at once.

def compute_grpo_advantages(rewards: Float[Tensor, " batch"], eps: float = 1e-6) -> Float[Tensor, " batch"]:
    """
    Computes GRPO advantages: each reward, normalized by the mean and std of the rewards in the group (the whole batch,
    since all our rollouts share one prompt).

    Args:
        rewards: the reward for each sequence
        eps: small constant added to the std dev for numerical stability

    Returns:
        The GRPO advantage of each sequence.
    """
    raise NotImplementedError()


tests.test_compute_grpo_advantages(compute_grpo_advantages)

Aside - should we divide by the standard deviation?

Dividing by the group's standard deviation is part of the original GRPO algorithm, but it has a side effect: groups where the rewards are nearly all equal (so the std is tiny) get their advantages scaled up. For a 0/1 reward, that means prompts which are almost always solved (or almost always failed) are weighted most heavily. Dr. GRPO (Liu et al., 2025) argues that this (along with the $\frac{1}{|o_i|}$ length normalization) biases training, and proposes dropping both.

Solution

def compute_grpo_advantages(rewards: Float[Tensor, " batch"], eps: float = 1e-6) -> Float[Tensor, " batch"]:
    """
    Computes GRPO advantages: each reward, normalized by the mean and std of the rewards in the group (the whole batch,
    since all our rollouts share one prompt).

    Args:
        rewards: the reward for each sequence
        eps: small constant added to the std dev for numerical stability

    Returns:
        The GRPO advantage of each sequence.
    """
    return (rewards - rewards.mean()) / (rewards.std() + eps)

Exercise - implement `get_logprobs`

Difficulty: 🔴🔴⚪⚪⚪

Importance: 🔵🔵🔵⚪⚪

You should spend up to 10-15 minutes on this exercise.

Everything in the GRPO objective is written in terms of $\log \pi(o_t \mid q, o_{\lt t})$, the log-probability that a model assigns to each generated token (we don't care about the prompt tokens, as the model has no effect on them). You should implement get_logprobs, which takes the logits of a model, and the token ids, and returns these log-probabilities for the last gen_len tokens.

Be careful with the indexing: the logits at position $s$ are the model's prediction for the token at position $s+1$. So the log-probability of the token at position $s$ comes from the logits at position $s-1$.

def get_logprobs(
    logits: Float[Tensor, "batch seq vocab"],
    tokens: Int[Tensor, "batch seq"],
    gen_len: int,
) -> Float[Tensor, "batch gen_len"]:
    """
    Returns the log-probabilities of the last `gen_len` tokens (the generated ones), under the model with these logits.

    Args:
        logits: the model's output logits. These only need to cover the last `gen_len + 1` positions of `tokens`, but
            may cover more.
        tokens: the token ids, of which the last `gen_len` are the generated tokens
        gen_len: the number of generated tokens

    Returns:
        The logprobs on the tokens sampled.
    """
    raise NotImplementedError()


tests.test_get_logprobs(get_logprobs)

Solution

def get_logprobs(
    logits: Float[Tensor, "batch seq vocab"],
    tokens: Int[Tensor, "batch seq"],
    gen_len: int,
) -> Float[Tensor, "batch gen_len"]:
    """
    Returns the log-probabilities of the last `gen_len` tokens (the generated ones), under the model with these logits.

    Args:
        logits: the model's output logits. These only need to cover the last `gen_len + 1` positions of `tokens`, but
            may cover more.
        tokens: the token ids, of which the last `gen_len` are the generated tokens
        gen_len: the number of generated tokens

    Returns:
        The logprobs on the tokens sampled.
    """
    logits_on_gen = logits[:, -gen_len - 1 :]
    tokens_on_gen = tokens[:, -gen_len - 1 :]
    logprobs_on_gen = logits_on_gen.log_softmax(dim=-1)
    return eindex(logprobs_on_gen, tokens_on_gen, "b s [b s+1]")

    # alternatively with gather (the logits at position s predict the token at s+1)
    # return logprobs_on_gen[:, :-1].gather(dim=-1, index=tokens_on_gen[:, 1:, None]).squeeze(-1)

We've also given you two helper functions which use get_logprobs:

get_logits runs the model on a batch of samples. It passes the logits_to_keep argument so that the model only computes logits for the last gen_len + 1 positions, which are all we need. This saves a lot of memory: the logits are a (batch, seq, d_vocab) tensor of floats, and d_vocab is 50k for GPT-2 and 152k for Qwen. This is done internally by slicing out the residual stream along the seq axis before unembedding, and it saves computing the logits for the prompt tokens (which we would discard anyway).
get_logprobs_batched computes the log-probabilities for a whole batch of samples, in chunks (to save memory) and without gradients. We use it in the rollout phase.

def get_logits(
    model: PreTrainedModel,
    sample_ids: Int[Tensor, "batch seq"],
    gen_len: int,
) -> Float[Tensor, "batch gen_len_plus_1 vocab"]:
    """Runs the model on a batch, returning float32 logits for the last `gen_len + 1` positions only."""
    return model(input_ids=sample_ids, logits_to_keep=gen_len + 1).logits.float()


@t.no_grad()
def get_logprobs_batched(
    model: PreTrainedModel,
    sample_ids: Int[Tensor, "batch seq"],
    gen_len: int,
    chunk_size: int,
) -> Float[Tensor, "batch gen_len"]:
    """Log-probabilities of the generated tokens under `model`, computed in chunks of `chunk_size` sequences."""
    logprobs = []
    for chunk in sample_ids.split(chunk_size):
        logits = get_logits(model, chunk, gen_len)
        logprobs.append(get_logprobs(logits, chunk, gen_len))
    return t.cat(logprobs)

The KL penalty

Now, we'd like to use the KL penalty to ensure that the new model doesn't diverge too much from the original reference model.

The formula for KL divergence of two distributions, $D_{\text{KL}}(P || Q)$, is $\sum_i P_i \log (P_i / Q_i)$. Recall that we want our new policy to be $P$ and the reference policy to be $Q$ (because this penalizes our new model for generating outputs which would be very unlikely under the original reference model), so in the form that we want, the KL penalty looks like:

\frac{1}{T} \sum_{t=1}^T D_{KL}(\pi_\theta(\cdot \mid s_t) \,\|\, \pi_\text{ref}(\cdot \mid s_t)) = \frac{1}{T} \sum_{t=1}^T \sum_a \pi_\theta(a \mid s_t) \log \frac{\pi_\theta(a \mid s_t)}{\pi_\text{ref}(a \mid s_t)}

Note that this formula performs a sum over the vocabulary dimension, as that is required to compute the KL divergence between the two distributions, but we mean over the position dimension, as we want the average KL divergence per token. If we are also computing this over a batch of sequences, we compute this for each sequence in the batch, and mean over that dimension too.

Below is what a naive implementation of this formula would look like, written in a numerically stable way. Note that we take logprobs as input, so the log-ratio inside the sum is just logprobs - ref_logprobs, and the weight $\pi_\theta(a \mid s_t)$ is recovered as logprobs.exp():

\log \frac{\pi_\theta(a \mid s_t)}{\pi_\text{ref}(a \mid s_t)} = \log \pi_\theta(a \mid s_t) - \log \pi_\text{ref}(a \mid s_t)

def calc_kl_penalty_naive(
    logprobs: Float[Tensor, "minibatch_size gen_len d_vocab"],
    ref_logprobs: Float[Tensor, "minibatch_size gen_len d_vocab"],
    kl_coef: float,
) -> Float[Tensor, ""]:
    probs = logprobs.exp()
    kl_div = (probs * (logprobs - ref_logprobs)).sum(-1)  # sum d_vocab

    return kl_coef * kl_div.mean()  # mean batch, seq


tests.test_calc_kl_penalty_naive(calc_kl_penalty_naive)
tests.test_calc_kl_penalty_naive_stability(calc_kl_penalty_naive)

This would work fine, but there are two problems

This would require us to cache the reference model's full logprobs in the replay buffer for every sequence generated. The logprobs are of shape (batch_size, gen_len, d_vocab), and while batch_size is a modest 128 and gen_len is 30, d_vocab ≈ 50k for GPT-2, and even larger for larger models: for the Llama-3 family d_vocab ≈ 128k, for Qwen d_vocab ≈ 152k, and as high as d_vocab ≈ 262k for the Gemma 3 family.

Even with our modestly small setup, 128 * 30 * 50k * 4 bytes/float ≈ 740 MiB of memory would be required just for the reference model's logprobs. That's half the size of GPT-2 medium itself!
Computing the KL divergence is both expensive and wasteful: requiring $O(\text{batch\_size} \times \text{gen\_len} \times \text{d\_vocab})$ operations, most of which don't meaningfully contribute. Most of the probability mass is concentrated on a few tokens, so the exact KL sum, which weights each token's log-ratio by its probability $p(x)$, is dominated by just a few terms. But even caching the logprobs for the top-k tokens can be avoided: as we'll see next, we can estimate the KL keeping only the single token actually sampled, with no extra logprobs at all.

@t.no_grad()
def kl_diagnostics(model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, prompt: str, k: int = 50):
    """Plot the model's top-k next-token probabilities after `prompt`, to show how concentrated the
    distribution is: almost all the mass sits on a handful of tokens."""
    input_ids = tokenizer(tokenizer.eos_token + prompt, return_tensors="pt").input_ids.to(device)
    logprobs = model(input_ids).logits[0, -1].float().log_softmax(-1)
    top = logprobs.argsort(descending=True)[:k]
    probs = logprobs[top].exp().numpy(force=True)
    labels = [repr(tokenizer.decode([i]))[1:-1] for i in top]

    fig, ax = plt.subplots(figsize=(9, 3.5), dpi=200)
    ax.bar(np.arange(k), probs, color="#378ADD")
    ax.set(ylabel="probability", title=f'top-{k} next-token probs after "{prompt}"')
    ax.set_xticks(np.arange(k), labels, rotation=90, fontsize=9)
    fig.tight_layout()
    return fig

kl_diagnostics(model, tokenizer, "The cat sat on the")
plt.show()

A cheaper KL

TL;DR We instead use an estimator of $\mathbb{E}_{x \sim P}[\log \frac{P(x)}{Q(x)}]$ (with $P = \pi_\theta$ and $Q = \pi_\text{ref}$), computed only from the token $x$ actually sampled, taking the mean over other dimensions. Writing $\rho^\text{ref} = \frac{P(x)}{Q(x)}$, we use the k2 estimator $\tfrac{1}{2}(\log \rho^\text{ref})^2$ as our penalty: its value is a slightly biased estimate of the KL, but its gradient is exactly the gradient of $D_{KL}(\pi_\theta \| \pi_\text{ref})$, which is what we actually need. If you care for the math, read on.

Math behind the cheap KL estimator

calc_kl_penalty_naive is exact, but expensive. We'd now like to construct some estimators of

D_{KL}(\pi_\theta || \pi_\text{ref}) = \mathbb{E}_{x \sim \pi_\theta} \left[ \log \frac{\pi_\theta(x)}{\pi_\text{ref}(x)} \right]

using only the logprobs from the tokens actually sampled by the policy. Define the per-token ratio to the reference policy on the sampled token $a_t$ as $\rho^\text{ref}_t = \dfrac{\pi_\theta(a_t \mid s_t)}{\pi_\text{ref}(a_t \mid s_t)}$, so that $\log \rho^\text{ref}_t = \log \pi_\theta(a_t \mid s_t) - \log \pi_\text{ref}(a_t \mid s_t)$. (The superscript distinguishes it from the importance ratio $\rho_t = \pi_\theta / \pi_{\theta_\text{old}}$ in the clipped objective: both compare the current policy to another policy, with a different policy in the denominator.) We have the following three potential estimators of the KL divergence, details about where they come from can be read in John Schulman's blog.

k1 $= \log \rho^\text{ref}_t$. The KL divergence is defined as an expectation, and this is just a sample of that expectation $D_{KL}(\pi_\theta || \pi_\text{ref}) = \mathbb{E}_{a \sim \pi_\theta} \left[ \log \frac{\pi_\theta(a \mid s_t)}{\pi_\text{ref}(a \mid s_t)} \right] \approx \log \frac{\pi_\theta(a_t \mid s_t)}{\pi_\text{ref}(a_t \mid s_t)} = \log \rho^\text{ref}_t$ Unbiased, but high variance: it's negative on about half the samples even though KL $\ge 0$.
k2 $= \tfrac{1}{2}(\log \rho^\text{ref}_t)^2$, the square of k1. Always $\ge 0$ and far lower variance; biased, but the bias is tiny when the two distributions are close. This works because its expectation is an f-divergence (with $f(x) = \tfrac{1}{2} x (\log x)^2$), and any f-divergence with $f''(1) = 1$ agrees with the KL divergence up to second order when $p$ is close to $q$ (Theorem 7.20, Polyanskiy & Wu, 2023). This is exactly the KL term used by prime-RL (with the 1/2 rolled into the kl_coef).
k3 $= (\rho^\text{ref}_t)^{-1} - 1 + \log \rho^\text{ref}_t$ (i.e. $e^{-\log \rho^\text{ref}_t} - 1 + \log \rho^\text{ref}_t$), this is k1 plus the zero-mean control variate $((\rho^\text{ref}_t)^{-1} - 1)$. It is unbiased, $\ge 0$, and low variance: it's Schulman's $(r-1) - \log r$ with $r = (\rho^\text{ref}_t)^{-1}$. It's his recommended estimator for measuring the KL, and it's the per-token KL penalty used in the GRPO paper (and by TRL's GRPOTrainer). But as we'll see below, it's the wrong choice for a loss.

Justification for the addition of the 1/rho_ref - 1 term in k3

Adding the term $(\rho^\text{ref}_t)^{-1} - 1$ doesn't bias the estimator, as the expectation of it under $\pi_\theta$ is 0:

\mathbb{E}_{a \sim \pi_\theta}\!\left[(\rho^\text{ref}_t)^{-1} \mid s\right] = \sum_a \pi_\theta(a | s)\,\frac{\pi_\text{ref}(a | s)}{\pi_\theta(a | s)} = \sum_a \pi_\text{ref}(a | s) = 1 \quad\Longrightarrow\quad \mathbb{E}_{a \sim \pi_\theta}\!\left[(\rho^\text{ref}_t)^{-1} - 1 \mid s\right] = 0.

So, $\lambda ((\rho^\text{ref}_t)^{-1} - 1) + \log \rho^\text{ref}_t$ is an unbiased estimator of the KL divergence for any $\lambda$. The sneaky trick is then recognizing that $ \log(x) \leq x - 1$, so if $\lambda=1$, then while $\log \rho^\text{ref}_t$ may not be positive,

k_3 = ((\rho^\text{ref}_t)^{-1} - 1) + \log \rho^\text{ref}_t = (\rho^\text{ref}_t)^{-1} - 1 - \log (\rho^\text{ref}_t)^{-1} \geq (\rho^\text{ref}_t)^{-1} - 1 - \left( (\rho^\text{ref}_t)^{-1} - 1 \right) = 0

will be non-negative on every sample, unlike **k1**.

Estimating the KL vs. differentiating it: why we use k2 as the loss

So far we've compared estimators by how well their values estimate the KL. But we don't use the penalty's value for anything; we backpropagate through it, so what matters is its gradient. The sampled token $a_t$ is held fixed when we differentiate (we don't differentiate through the sampling), so for any per-token penalty $k(\log \rho^\text{ref}_t)$ the expected gradient is $\mathbb{E}_{a \sim \pi_\theta}\big[k'(\log \rho^\text{ref})\, \nabla_\theta \log \pi_\theta(a \mid s)\big]$. Using $\sum_a \pi_\theta \nabla_\theta \log \pi_\theta = \sum_a \nabla_\theta \pi_\theta = 0$, the three estimators give:

k1 ($k' = 1$): $\;\sum_a \pi_\theta \nabla_\theta \log \pi_\theta = 0$. Zero expected gradient: as a loss, it's pure noise.
k2 ($k' = \log \rho^\text{ref}$): $\;\sum_a \pi_\theta \log \frac{\pi_\theta}{\pi_\text{ref}} \nabla_\theta \log \pi_\theta = \nabla_\theta D_{KL}(\pi_\theta | \pi_\text{ref})$. Exactly the gradient we want.
k3 ($k' = 1 - (\rho^\text{ref})^{-1}$): $\;\sum_a (\pi_\theta - \pi_\text{ref}) \nabla_\theta \log \pi_\theta = -\sum_a \pi_\text{ref} \nabla_\theta \log \pi_\theta = \nabla_\theta D_{KL}(\pi_\text{ref} | \pi_\theta)$. This is the gradient of the forward, mode-covering KL: the direction we argued against earlier!

So k3 is the best of the three at measuring $D_{KL}(\pi_\theta \| \pi_\text{ref})$ while the policy is close to the reference, but differentiating it optimizes the other KL. k2 is biased as a measurement, but differentiating it gives exactly the reverse-KL gradient, so that's what we'll use in the loss. For more on this pitfall, see Tang & Munos (2025), On a few pitfalls in KL divergence gradient estimation for RL.

(Note that the k2 gradient $\log \rho^\text{ref}_t \, \nabla_\theta \log \pi_\theta(a_t \mid s_t)$ is exactly the REINFORCE gradient you'd get by putting k1 into the reward instead, i.e. treating $-\beta \log \rho^\text{ref}_t$ as a per-token reward, which is how the original RLHF papers applied their KL penalty. k1 is fine in the reward; it's only useless as a directly-differentiated loss.)

Can't we just swap the arguments of k3? Swapping $\pi_\theta$ and $\pi_\text{ref}$ in k3 gives $\rho^\text{ref}_t - 1 - \log \rho^\text{ref}_t$, with $k' = \rho^\text{ref} - 1$, so its expected gradient is $\sum_a (\pi_\theta \rho^\text{ref} - \pi_\theta) \nabla_\theta \log \pi_\theta = \sum_a \frac{\pi_\theta^2}{\pi_\text{ref}} \nabla_\theta \log \pi_\theta = \tfrac{1}{2} \nabla_\theta \sum_a \frac{\pi_\theta^2}{\pi_\text{ref}}$, the gradient of the $\chi^2$ divergence $\chi^2(\pi_\theta \| \pi_\text{ref})$ (up to a factor of $\frac{1}{2}$). Its value doesn't estimate the KL either, since $\mathbb{E}_{\pi_\theta}[\rho^\text{ref}] \ne 1$.

We're also being a little sloppy here. The derivations above average over tokens sampled from $\pi_\theta$, but our tokens were actually sampled from $\pi_{old}$, the policy as it was at rollout time. Note that the logprobs we feed into the penalty are the fresh ones from re-running the current model $\pi_\theta$ (so that gradients flow through the penalty), compared against the frozen $\pi_{ref}$: it's only the sampling distribution of the tokens that is stale. Now, $\pi_\theta$ and $\pi_{old}$ are not the same distribution, but they are close, especially on the tokens with high probability mass. So a token sampled from $\pi_{old}$ was probably also going to be picked by $\pi_\theta$, which is why we can get away with using these samples to estimate the expectation under $\pi_\theta$.

Exercise - implement `calc_kl_penalty`

Difficulty: 🔴🔴⚪⚪⚪

Importance: 🔵🔵🔵⚪⚪

You should spend up to 10-15 minutes on this exercise.

Implement calc_kl_penalty below using the k2 estimator $\tfrac{1}{2}(\log \rho^\text{ref}_t)^2$, where $\rho^\text{ref}_t = \frac{\pi_\theta(a_t | s_t)}{\pi_{\text{ref}}(a_t | s_t)}$ is the ratio of the current and reference policies on the sampled token. Make sure to do this in a numerically stable way: compute $\log \rho^\text{ref}_t$ directly as a difference of logprobs, rather than exponentiating the logprobs and dividing. We compute k2 for every sequence position, and take the mean both over sequence and batch dimensions.

def calc_kl_penalty(
    logprobs: Float[Tensor, "minibatch_size gen_len"],
    ref_logprobs: Float[Tensor, "minibatch_size gen_len"],
    kl_coef: float,
) -> Float[Tensor, ""]:
    """
    Computes the KL penalty `kl_coef * D_KL(pi_theta || pi_ref)` from only the generated tokens' logprobs, using the k2
    estimator 0.5 * (log pi_theta - log pi_ref)^2.
    Args:
        logprobs:
            log pi_theta(a_t | s_t) for each generated token (under the new model).
        ref_logprobs:
            log pi_ref(a_t | s_t) for each generated token (under the reference model).
        kl_coef:
            The coefficient of the KL penalty.

    Output:
        The (estimated) KL divergence between the new and reference policies, scaled by kl_coef.
    """
    assert logprobs.shape == ref_logprobs.shape, "logprobs and ref_logprobs must have the same shape"

    raise NotImplementedError()


tests.test_calc_kl_penalty(calc_kl_penalty)
tests.test_calc_kl_penalty_stability(calc_kl_penalty)

Bias/variance tradeoff for KL estimators

A tuned policy should never put mass on tokens the base model wouldn't: it just concentrates the base's distribution. For the moment, we will use as a stand-in for the tuned model simply the base model with a lower temperature (as if the model was trained to have a sharper distribution, reinforcing what the model would already do). A temperature close to 1 keeps the KL small, which is the regime real RLHF actually trains in (the KL penalty keeps the policy near the reference).

We sweep over text from the Pile: at every token position we get the two next-token distributions, compute the exact mean per-token $D_{\text{KL}}(\pi_\theta \,|\, \pi_\text{ref})$, then see how well each estimator recovers it from sampled tokens alone, using the three estimators from the dropdown above: k1 $= \log \rho^\text{ref}$, k2 $= \tfrac{1}{2}(\log \rho^\text{ref})^2$ and k3 $= (\rho^\text{ref})^{-1} - 1 + \log \rho^\text{ref}$.

@t.no_grad()
def kl_estimator_comparison(model, tokenizer, texts, temperature=0.9, seq_len=128, n_seqs=128, batch_size=8, seed=0):
    """Compare the k1/k2/k3 KL estimators' bias and variance to the exact D_KL(pi_theta || pi_ref), where pi_ref is
    `model` and pi_theta is the same model with its logits sharpened to `temperature` < 1 (a small stand-in for an
    RL-tuned policy). Chunks `texts` into fixed-length sequences and estimates the KL from one sampled token each.
    The token sampling uses a seeded generator so the printed table is reproducible (the estimators are stochastic)."""
    ids = t.cat([t.tensor(tokenizer(txt).input_ids) for txt in texts])
    n = min(n_seqs, len(ids) // seq_len)
    seqs = ids[: n * seq_len].reshape(n, seq_len).to(device)
    rng = t.Generator(device=device).manual_seed(seed)

    true_kls, log_ratios = [], []
    for batch in tqdm(seqs.split(batch_size), desc="batches"):
        logits = model(batch).logits[:, :-1].float()
        logp = (logits / temperature).log_softmax(-1)  # pi_theta (sharpened policy)
        logq = logits.log_softmax(-1)  # pi_ref (temperature 1)
        true_kls.append((logp.exp() * (logp - logq)).sum(-1).flatten())  # exact per-token KL
        # one token sampled from the policy at each position (multinomial takes a generator; Categorical doesn't)
        x = t.multinomial(logp.exp().reshape(-1, logp.shape[-1]), 1, generator=rng).reshape(*logp.shape[:-1], 1)
        log_ratios.append((logp.gather(-1, x) - logq.gather(-1, x)).flatten())  # log pi_theta(x) - log pi_ref(x)

    true_kl = t.cat(true_kls).mean().item()
    log_ratio = t.cat(log_ratios)
    estimators = {
        "k1 = log_ratio": log_ratio,
        "k2 = 0.5*log_ratio**2": 0.5 * log_ratio**2,
        "k3 = 1/ratio-1+log_ratio": (-log_ratio).exp() - 1 + log_ratio,
    }
    print(f"mean per-token true KL = {true_kl:.4f}   ({log_ratio.numel()} tokens)")
    print(f"{'estimator':24} {'mean':>9} {'bias':>8} {'stdev/true':>11}")
    for name, est in estimators.items():
        m, s = est.mean().item(), est.std().item()
        print(f"{name:24} {m:>9.4f} {(m - true_kl) / true_kl:>+8.1%} {s / true_kl:>11.2f}")

If you run this for yourself:

from datasets import load_dataset

pile = load_dataset("NeelNanda/pile-10k", split="train[:64]")
kl_estimator_comparison(model, tokenizer, pile["text"], temperature=0.9)  # T=1 ref vs sharpened T=0.9 policy

You should see something like:

mean per-token true KL = 0.0301   (16256 tokens)
estimator                     mean     bias  stdev/true
k1 = log_ratio              0.0288    -4.4%        7.50
k2 = 0.5*log_ratio**2       0.0259   -14.0%        1.92
k3 = 1/ratio-1+log_ratio    0.0289    -4.0%        2.67

In this realistic small-KL regime the lesson is clean: * k1 is unbiased but has huge variance (several times the true KL per sample, which is functionally useless); * k2 has low variance but is biased (it underestimates the true KL here); * k3 is unbiased and low-variance: the best of both worlds for measuring the KL.

So if you want to log how far your policy has drifted from the reference while it's still close, k3 is the one to use. But k3 has a weakness of its own: once training has moved the policy further away from the reference, k3 becomes heavy-tailed (a single sampled token that the policy has made very unlikely can make $(\rho^\text{ref})^{-1}$ enormous), and plain k1 averaged over the batch is the steadier measurement. That's what our trainer logs, since with kl_coef=0 the policy is free to drift a long way. But remember from the "Estimating the KL vs. differentiating it" dropdown above that this table only compares the estimators' values: as a loss, k2's bias in value doesn't matter, because its gradient is exactly the reverse-KL gradient, while k3's gradient is that of the forward KL.

Solution

def calc_kl_penalty(
    logprobs: Float[Tensor, "minibatch_size gen_len"],
    ref_logprobs: Float[Tensor, "minibatch_size gen_len"],
    kl_coef: float,
) -> Float[Tensor, ""]:
    """
    Computes the KL penalty `kl_coef * D_KL(pi_theta || pi_ref)` from only the generated tokens' logprobs, using the k2
    estimator 0.5 * (log pi_theta - log pi_ref)^2.
    Args:
        logprobs:
            log pi_theta(a_t | s_t) for each generated token (under the new model).
        ref_logprobs:
            log pi_ref(a_t | s_t) for each generated token (under the reference model).
        kl_coef:
            The coefficient of the KL penalty.

    Output:
        The (estimated) KL divergence between the new and reference policies, scaled by kl_coef.
    """
    assert logprobs.shape == ref_logprobs.shape, "logprobs and ref_logprobs must have the same shape"

    log_ratio = logprobs - ref_logprobs
    k2 = 0.5 * log_ratio**2
    return kl_coef * k2.mean()

Skipping the entropy bonus

If you've done the PPO material, you might be wondering where the entropy bonus went. In many setups for RL for language models, the entropy bonus can be safely skipped, as the KL penalty already contains a $-H(\pi_\theta)$ term, since

D_{KL}(\pi_\theta \| \pi_\text{ref}) = -H(\pi_\theta) - \mathbb{E}_{\pi_\theta}[\log \pi_\text{ref}] = -H(\pi_\theta) + \text{CE}(\pi_\theta \| \pi_\text{ref})

(where $\text{CE}$ is the standard cross-entropy loss). Minimizing the KL pulls $\pi_\theta$ towards $\pi_\text{ref}$, and a pretrained base model spreads its probability over many plausible continuations. So the KL penalty already rewards keeping the entropy of the policy up, which is usually enough to prevent collapse without a separate explicit entropy bonus.

The clipped surrogate objective

Exercise - implement `calc_clipped_surrogate_objective`

Difficulty: 🔴🔴🔴⚪⚪

Importance: 🔵🔵🔵🔵🔵

You should spend up to 20-30 minutes on this exercise.

This is the heart of the algorithm. Let's build it up in steps.

Step 1: the importance ratio. For a single gradient step on fresh samples from $\pi_\theta$, the policy gradient from earlier says to maximize $\sum_t \hat{A}_t \log \pi_\theta(o_t)$ (holding the samples fixed). But we want to take several gradient steps per batch of samples, so after the first step our samples come from an older policy $\pi_{\theta_\text{old}}$. Importance sampling fixes this: for any function $f$, $\mathbb{E}_{o \sim \pi_\theta}[f(o)] = \mathbb{E}_{o \sim \pi_{\theta_\text{old}}}\left[\frac{\pi_\theta(o)}{\pi_{\theta_\text{old}}(o)} f(o)\right]$. Applying this per token gives the surrogate objective $\rho_t \hat{A}_t$, with $\rho_t = \frac{\pi_\theta(o_t \mid \ldots)}{\pi_{\theta_\text{old}}(o_t \mid \ldots)}$. When $\theta = \theta_\text{old}$, $\rho_t = 1$ and $\nabla_\theta \rho_t = \nabla_\theta \log \pi_\theta(o_t)$, so its gradient is exactly the policy gradient, and it stays a sensible objective as $\theta$ moves away.

Step 2: clipping. The surrogate is only a good approximation while $\pi_\theta$ stays close to $\pi_{\theta_\text{old}}$: the samples tell us nothing about what happens if we move far away. Maximizing it naively would let the optimizer keep pushing $\pi_\theta$ further away from $\pi_{\theta_\text{old}}$ based on data that $\pi_\theta$ didn't generate. So we clip the ratio to $[1 - \epsilon, 1 + \epsilon]$ (with $\epsilon = $ clip_coef), and take the minimum of the clipped and unclipped objectives:

J^{\text{CLIP}}_t = \min\Big[\rho_t \hat{A}_t, \; \text{clip}(\rho_t, 1-\epsilon, 1+\epsilon)\,\hat{A}_t\Big]

If $\hat{A}_t > 0$ (the token was good), the objective increases with $\rho_t$ until $\rho_t = 1 + \epsilon$, after which it's flat, so there's no more gradient pushing the token's probability up.
If $\hat{A}_t < 0$ (the token was bad), the objective increases as $\rho_t$ decreases, until $\rho_t = 1 - \epsilon$.
Because we take the minimum, clipping only ever removes incentives to move further; it never stops us undoing a change that made things worse (e.g. if $\hat{A}_t > 0$ but $\rho_t < 1 - \epsilon$ (i.e the current policy is less likely than the old policy to generate the good token), then we made a mistake updating and should revert, and we can as the unclipped term is the minimum, and thus still has a gradient).

You should implement this below. Remember that the advantages have shape (batch,) (one per sequence), and need to be applied to every token of their sequence. Return the mean over all tokens and sequences. Work with log-probabilities, and only exponentiate the difference (as you did for the KL penalty).

Here's a plot of $J^{\text{CLIP}}$ as a function of the ratio $\rho$, for a positive and a negative advantage (compare Figure 1 of the PPO paper). Each learning phase starts at $\rho = 1$.

eps = 0.2
ratio = np.linspace(0.5, 1.5, 201)
fig, ax = plt.subplots(figsize=(6, 3.5), dpi=150)
for adv, color in [(1.0, "#378ADD"), (-1.0, "#D85A30")]:
    clipped = np.minimum(ratio * adv, np.clip(ratio, 1 - eps, 1 + eps) * adv)
    ax.plot(ratio, clipped, color=color, label=f"$\\hat{{A}} = {adv:+.0f}$")
for x in (1 - eps, 1, 1 + eps):
    ax.axvline(x, color="gray", lw=0.8, ls=":" if x != 1 else "--")
ax.set(xlabel=r"ratio $\rho$", ylabel=r"$J^{\mathrm{CLIP}}$", title=f"Clipped surrogate objective ($\\epsilon = {eps}$)")
ax.legend()
fig.tight_layout()
plt.show()

Now implement the function below.

def calc_clipped_surrogate_objective(
    logprobs: Float[Tensor, "minibatch_size gen_len"],
    old_logprobs: Float[Tensor, "minibatch_size gen_len"],
    advantages: Float[Tensor, " minibatch_size"],
    clip_coef: float,
) -> Float[Tensor, ""]:
    """
    Returns the clipped surrogate objective, suitable for maximisation with gradient ascent.

    Args:
        logprobs:
            log pi_theta of each generated token, under the current policy (with gradients)
        old_logprobs:
            log pi_theta_old of each generated token, under the policy at rollout time
        advantages:
            the advantage of each sequence, which applies to all of its tokens
        clip_coef:
            amount of clipping, denoted by epsilon in the GRPO / PPO papers

    Returns:
        The clipped surrogate objective, averaged over all tokens in the minibatch.
    """
    assert logprobs.shape == old_logprobs.shape and advantages.shape == logprobs.shape[:1]

    raise NotImplementedError()


tests.test_calc_clipped_surrogate_objective(calc_clipped_surrogate_objective)

Solution

def calc_clipped_surrogate_objective(
    logprobs: Float[Tensor, "minibatch_size gen_len"],
    old_logprobs: Float[Tensor, "minibatch_size gen_len"],
    advantages: Float[Tensor, " minibatch_size"],
    clip_coef: float,
) -> Float[Tensor, ""]:
    """
    Returns the clipped surrogate objective, suitable for maximisation with gradient ascent.

    Args:
        logprobs:
            log pi_theta of each generated token, under the current policy (with gradients)
        old_logprobs:
            log pi_theta_old of each generated token, under the policy at rollout time
        advantages:
            the advantage of each sequence, which applies to all of its tokens
        clip_coef:
            amount of clipping, denoted by epsilon in the GRPO / PPO papers

    Returns:
        The clipped surrogate objective, averaged over all tokens in the minibatch.
    """
    assert logprobs.shape == old_logprobs.shape and advantages.shape == logprobs.shape[:1]

    ratio = (logprobs - old_logprobs).exp()
    advantages = advantages[:, None]  # broadcast each sequence's advantage over its tokens

    non_clipped = ratio * advantages
    clipped = ratio.clip(1 - clip_coef, 1 + clip_coef) * advantages

    return t.minimum(non_clipped, clipped).mean()

Memory

We've given you an implementation of the ReplayMemory and ReplayMinibatch classes. These store the data from each rollout phase, and split it into minibatches for the learning phase. Each minibatch stores:

sample_ids, the sampled sequences (we run the current policy on these in the learning phase).
logprobs, the log-probabilities of the generated tokens under $\pi_{\theta_\text{old}}$, i.e. the policy at rollout time, for the importance ratio.
ref_logprobs, the log-probabilities of the generated tokens under $\pi_\text{ref}$, for the KL penalty.
advantages, one per sequence.

We don't need to store the actions separately, since the actions (tokens generated) are contained within the sequences themselves.

@dataclass
class ReplayMinibatch:
    """
    Samples from the replay memory.
    """

    sample_ids: Int[Tensor, "minibatch_size seq_len"]
    logprobs: Float[Tensor, "minibatch_size gen_len"]
    ref_logprobs: Float[Tensor, "minibatch_size gen_len"]
    advantages: Float[Tensor, " minibatch_size"]


class ReplayMemory:
    def __init__(
        self,
        args: GRPO_Args,
        sample_ids: Int[Tensor, "batch_size seq_len"],
        logprobs: Float[Tensor, "batch_size gen_len"],
        ref_logprobs: Float[Tensor, "batch_size gen_len"],
        advantages: Float[Tensor, " batch_size"],
    ):
        """Initializes the replay memory, with all the data generated from the rollout phase at once."""
        assert sample_ids.shape[0] == args.batch_size
        assert logprobs.shape == ref_logprobs.shape == (args.batch_size, args.gen_len)
        assert advantages.shape == (args.batch_size,)

        self.args = args
        self.sample_ids = sample_ids
        self.logprobs = logprobs
        self.ref_logprobs = ref_logprobs
        self.advantages = advantages

    def get_minibatches(self) -> list[ReplayMinibatch]:
        """
        Generates a list of minibatches by randomly sampling from the replay memory. Each sequence
        appears exactly `batches_per_learning_phase` times in total.
        """
        minibatches = []
        for _ in range(self.args.batches_per_learning_phase):
            for indices in t.randperm(self.args.batch_size).reshape(self.args.num_minibatches, -1):
                minibatches.append(
                    ReplayMinibatch(
                        sample_ids=self.sample_ids[indices],
                        logprobs=self.logprobs[indices],
                        ref_logprobs=self.ref_logprobs[indices],
                        advantages=self.advantages[indices],
                    )
                )
        return minibatches

Optimizer & Scheduler

We've given you the optimizer and learning rate scheduler. A few notes:

We use AdamW with maximize=True, since we're maximizing an objective function rather than minimizing a loss function. We only pass it the parameters with requires_grad=True, which will matter when we come to LoRA.
The scheduler uses PyTorch's LambdaLR, which multiplies the base learning rate by lr_lambda(step). We use a linear warmup (to avoid large updates while Adam's moment estimates are still poorly calibrated) followed by a linear decay to final_scale times the base learning rate. We step the scheduler once per phase.

Here lr_lambda is a function mapping the number of steps (i.e. number of times we've called scheduler.step()) to a float which gets multiplied by the base learning rate (i.e. 0.1 means we use 10% of the base LR). There are schedulers other than LambdaLR which have specific built-in behaviour (see documentation page), although this gives you the most flexibility.

Aside - why we use warmup

Warmup is a common strategy early in training, to make sure we don't get excessive updates early on. It seems to work pretty well empirically. Some possible reasons for this are:

It helps avoid large updates when the Adam moving averages of first and second moments are not yet well calibrated.
Early on in training, the gradients might be very large because the model's predictions are nowhere near where they need to be. So an LR warmup is more useful early on, to help avoid massive steps.

We've given you the code you'll be using for returning a custom lr_lambda function with a linear warmup then linear decay. We've also provided code for you in the trainer class's init method below which creates your scheduler. All you need to do is make sure you're stepping it appropriately.

def get_optimizer_and_scheduler(
    args: GRPO_Args, model: nn.Module
) -> tuple[t.optim.Optimizer, t.optim.lr_scheduler.LambdaLR]:
    """
    Creates an AdamW optimizer over the model's trainable parameters, and an LR scheduler that linearly warms up for
    `warmup_steps` phases, and then linearly decays to `final_scale` over the remaining phases.
    """

    def lr_lambda(step: int) -> float:
        assert step <= args.total_phases, f"Step = {step} should be at most total_phases = {args.total_phases}."
        if step < args.warmup_steps:
            return (step + 1) / (args.warmup_steps + 1)
        else:
            return 1 - (1 - args.final_scale) * (step - args.warmup_steps) / (args.total_phases - args.warmup_steps)

    params = [p for p in model.parameters() if p.requires_grad]
    fused = all(p.is_cuda for p in params)  # the whole update step runs as one GPU kernel
    optimizer = t.optim.AdamW(params, lr=args.lr, maximize=True, fused=fused)
    scheduler = t.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)
    return optimizer, scheduler

Training your model

We're now ready to put everything together!

Exercise - complete `GRPO_Trainer`

Difficulty: 🔴🔴🔴🔴🔴

Importance: 🔵🔵🔵🔵🔵

You should spend up to 40-60 minutes on this exercise.

You need to implement three methods:

rollout_phase should sample completions with get_samples and then compute everything the learning phase needs and return it as a ReplayMemory: the log-probabilities of the generated tokens under the current policy and under the reference policy (use get_logprobs_batched for the first, and the get_ref_logprobs method we've given you for the second), the rewards, and the GRPO advantages. Then call self.log_rollout (given) to print some samples. All of this should happen without gradients.
compute_grpo_objective should compute the GRPO objective for a minibatch: run the current model on the minibatch's samples (use get_logits), get the log-probabilities of the generated tokens, and combine the clipped surrogate objective and the KL penalty (make sure you get the sign right!). You can also log things to Weights & Biases here.
learning_phase should iterate over the minibatches from memory.get_minibatches(), and for each one:
- zero the gradients,
- compute the objective,
- backpropagate,
- clip the gradient norm to args.max_grad_norm using t.nn.utils.clip_grad_norm_,
- step the optimizer,
- and increment self.step.

Step the scheduler once at the end of the phase. Return the mean objective.

A few tips / notes before you start:

We keep the reference model behind the method get_ref_logprobs, rather than using self.ref_model directly, because in the LoRA section we'll override it.
For faster feedback loops, don't use wandb until you've stopped getting errors!
We keep the policy in evaluation mode (.eval()) throughout, even while training it. GPT-2 uses dropout (or at least the HuggingFace implementation of it does, unlike the one we built from scratch), and in training mode it would randomly perturb every forward pass, so the log-probabilities we compute in the rollout phase, in the learning phase and for the reference model would all be noisy, and the importance ratio wouldn't even start at 1. .eval() only switches off dropout (gradients still flow), which is standard practice for RL on language models.

class GRPO_Trainer:
    model: PreTrainedModel
    ref_model: PreTrainedModel | None

    def __init__(self, args: GRPO_Args):
        t.manual_seed(args.seed)
        self.args = args
        self.run_name = f"{args.wandb_project_name}__seed{args.seed}__{time.strftime('%Y%m%d-%H%M%S')}"

        self.tokenizer = AutoTokenizer.from_pretrained(args.base_model)
        # dtype=float32 explicitly, since Hugging Face otherwise loads the checkpoint's own dtype (often bfloat16)
        # Both in eval mode, which switches off GPT-2's dropout (see the notes above)
        self.model = AutoModelForCausalLM.from_pretrained(args.base_model, dtype=t.float32).to(device).eval()
        self.ref_model = AutoModelForCausalLM.from_pretrained(args.base_model, dtype=t.float32).to(device).eval()
        self.ref_model.requires_grad_(False)
        self.optimizer, self.scheduler = get_optimizer_and_scheduler(args, self.model)

    def get_ref_logprobs(self, sample_ids: Int[Tensor, "batch seq"]) -> Float[Tensor, "batch gen_len"]:
        """Log-probabilities of the generated tokens under the reference policy."""
        return get_logprobs_batched(self.ref_model, sample_ids, self.args.gen_len, self.args.minibatch_size)

    @t.no_grad()
    def log_rollout(
        self,
        samples: list[str],
        rewards: Float[Tensor, " batch"],
        logprobs: Float[Tensor, "batch gen_len"],
        ref_logprobs: Float[Tensor, "batch gen_len"],
    ) -> None:
        """Prints (and logs to wandb) the mean reward, the KL to the reference policy, and a few samples."""
        # Per-token KL, measured with the k1 estimator averaged over the batch: unbiased, and unlike k3 it doesn't blow up
        # when the policy drifts far from the reference (see the KL section)
        kl = (logprobs - ref_logprobs).mean().item()
        if self.args.use_wandb:
            wandb.log({"mean_reward": rewards.mean().item(), "kl": kl}, step=self.step)

        # The first 3 samples
        table = [[f"{rewards[i].item():.2f}", f"{ref_logprobs[i].sum().item():.2f}", repr(samples[i])] for i in range(3)]
        table = tabulate(table, ["Reward", "Ref logprobs", "Sample"], tablefmt="simple_grid", maxcolwidths=[None, None, 90])
        print(
            f"Phase {self.phase + 1:03}/{self.args.total_phases:03}, Mean reward: {rewards.mean().item():.4f}, "
            f"KL: {kl:.4f}\n{table}\n"
        )

    def rollout_phase(self) -> ReplayMemory:
        """
        Performs a single rollout phase, returning a ReplayMemory object containing the data generated during this
        phase. All forward passes here should be done without gradients.

        Steps of this function are:
            - Generate samples from our model
            - Get the logprobs of the generated tokens under the current & reference policies
            - Compute the rewards and advantages
            - Log the rollouts, and return the data in a ReplayMemory object
        """
        sample_ids, samples = get_samples(
            self.model,
            self.tokenizer,
            prompt=self.args.prompt,
            batch_size=self.args.batch_size,
            gen_len=self.args.gen_len,
            temperature=self.args.temperature,
        )

        raise NotImplementedError()

    def compute_grpo_objective(self, minibatch: ReplayMinibatch) -> Float[Tensor, ""]:
        """
        Computes the GRPO objective function to maximize: the clipped surrogate objective minus the KL penalty.

        Steps of this function are:
            - Get the logits for the samples in the minibatch, under the current policy
            - Get the logprobs of the generated tokens
            - Use these to compute both terms of the objective function, and return it
            - Also optionally log stuff to Weights & Biases (provided)
        """
        raise NotImplementedError()

        if self.args.use_wandb:
            with t.no_grad():
                ratio = (logprobs - minibatch.logprobs).exp()
                clipfrac = ((ratio - 1.0).abs() > self.args.clip_coef).float().mean().item()
            wandb.log(
                dict(
                    total_steps=self.step,
                    lr=self.scheduler.get_last_lr()[0],
                    clipped_surrogate_objective=clipped_surrogate_objective.item(),
                    clipfrac=clipfrac,
                    kl_penalty=kl_penalty.item(),
                ),
                step=self.step,
            )

        return objective

    def learning_phase(self, memory: ReplayMemory) -> float:
        """
        Performs a learning step on `memory`. This involves the standard training steps
        (i.e. zeroing gradient, computing objective function, doing backprop, stepping optimizer).

        You should also remember the following:
            - Clipping grad norm to the value given in `self.args.max_grad_norm`
            - Incrementing `self.step` by 1 for each minibatch
            - Stepping the scheduler (once per calling of this function)

        Returns the average objective function value over the minibatches as a float for logging.
        """
        raise NotImplementedError()

    def train(self) -> None:
        """
        Performs a full training run.
        """
        self.step = 0

        if self.args.use_wandb:
            wandb.init(
                project=self.args.wandb_project_name,
                entity=self.args.wandb_entity,
                name=self.run_name,
                config=self.args,
            )
        runner = tqdm(range(self.args.total_phases))
        for self.phase in runner:
            memory = self.rollout_phase()
            objective = self.learning_phase(memory)
            runner.set_description(f"Objective: {objective:.4f}")

        if self.args.use_wandb:
            wandb.finish()


tests.test_grpo_trainer(GRPO_Trainer, GRPO_Args)

Solution

class GRPO_Trainer:
    model: PreTrainedModel
    ref_model: PreTrainedModel | None

    def __init__(self, args: GRPO_Args):
        t.manual_seed(args.seed)
        self.args = args
        self.run_name = f"{args.wandb_project_name}__seed{args.seed}__{time.strftime('%Y%m%d-%H%M%S')}"

        self.tokenizer = AutoTokenizer.from_pretrained(args.base_model)
        # dtype=float32 explicitly, since Hugging Face otherwise loads the checkpoint's own dtype (often bfloat16)
        # Both in eval mode, which switches off GPT-2's dropout (see the notes above)
        self.model = AutoModelForCausalLM.from_pretrained(args.base_model, dtype=t.float32).to(device).eval()
        self.ref_model = AutoModelForCausalLM.from_pretrained(args.base_model, dtype=t.float32).to(device).eval()
        self.ref_model.requires_grad_(False)
        self.optimizer, self.scheduler = get_optimizer_and_scheduler(args, self.model)

    def get_ref_logprobs(self, sample_ids: Int[Tensor, "batch seq"]) -> Float[Tensor, "batch gen_len"]:
        """Log-probabilities of the generated tokens under the reference policy."""
        return get_logprobs_batched(self.ref_model, sample_ids, self.args.gen_len, self.args.minibatch_size)

    @t.no_grad()
    def log_rollout(
        self,
        samples: list[str],
        rewards: Float[Tensor, " batch"],
        logprobs: Float[Tensor, "batch gen_len"],
        ref_logprobs: Float[Tensor, "batch gen_len"],
    ) -> None:
        """Prints (and logs to wandb) the mean reward, the KL to the reference policy, and a few samples."""
        # Per-token KL, measured with the k1 estimator averaged over the batch: unbiased, and unlike k3 it doesn't blow up
        # when the policy drifts far from the reference (see the KL section)
        kl = (logprobs - ref_logprobs).mean().item()
        if self.args.use_wandb:
            wandb.log({"mean_reward": rewards.mean().item(), "kl": kl}, step=self.step)

        # The first 3 samples
        table = [[f"{rewards[i].item():.2f}", f"{ref_logprobs[i].sum().item():.2f}", repr(samples[i])] for i in range(3)]
        table = tabulate(table, ["Reward", "Ref logprobs", "Sample"], tablefmt="simple_grid", maxcolwidths=[None, None, 90])
        print(
            f"Phase {self.phase + 1:03}/{self.args.total_phases:03}, Mean reward: {rewards.mean().item():.4f}, "
            f"KL: {kl:.4f}\n{table}\n"
        )

    def rollout_phase(self) -> ReplayMemory:
        """
        Performs a single rollout phase, returning a ReplayMemory object containing the data generated during this
        phase. All forward passes here should be done without gradients.

        Steps of this function are:
            - Generate samples from our model
            - Get the logprobs of the generated tokens under the current & reference policies
            - Compute the rewards and advantages
            - Log the rollouts, and return the data in a ReplayMemory object
        """
        sample_ids, samples = get_samples(
            self.model,
            self.tokenizer,
            prompt=self.args.prompt,
            batch_size=self.args.batch_size,
            gen_len=self.args.gen_len,
            temperature=self.args.temperature,
        )

        # Log-probabilities of the generated tokens under the current (= old, from the learning phase's point of view)
        # and reference policies
        logprobs = get_logprobs_batched(self.model, sample_ids, self.args.gen_len, self.args.minibatch_size)
        ref_logprobs = self.get_ref_logprobs(sample_ids)

        # Rewards, and group-relative advantages
        rewards = self.args.reward_fn(samples)
        advantages = compute_grpo_advantages(rewards)

        self.log_rollout(samples, rewards, logprobs, ref_logprobs)

        return ReplayMemory(
            args=self.args,
            sample_ids=sample_ids,
            logprobs=logprobs,
            ref_logprobs=ref_logprobs,
            advantages=advantages,
        )

    def compute_grpo_objective(self, minibatch: ReplayMinibatch) -> Float[Tensor, ""]:
        """
        Computes the GRPO objective function to maximize: the clipped surrogate objective minus the KL penalty.

        Steps of this function are:
            - Get the logits for the samples in the minibatch, under the current policy
            - Get the logprobs of the generated tokens
            - Use these to compute both terms of the objective function, and return it
            - Also optionally log stuff to Weights & Biases (provided)
        """
        logits = get_logits(self.model, minibatch.sample_ids, self.args.gen_len)
        logprobs = get_logprobs(logits, minibatch.sample_ids, self.args.gen_len)

        clipped_surrogate_objective = calc_clipped_surrogate_objective(
            logprobs, minibatch.logprobs, minibatch.advantages, self.args.clip_coef
        )
        kl_penalty = calc_kl_penalty(logprobs, minibatch.ref_logprobs, self.args.kl_coef)
        objective = clipped_surrogate_objective - kl_penalty

        if self.args.use_wandb:
            with t.no_grad():
                ratio = (logprobs - minibatch.logprobs).exp()
                clipfrac = ((ratio - 1.0).abs() > self.args.clip_coef).float().mean().item()
            wandb.log(
                dict(
                    total_steps=self.step,
                    lr=self.scheduler.get_last_lr()[0],
                    clipped_surrogate_objective=clipped_surrogate_objective.item(),
                    clipfrac=clipfrac,
                    kl_penalty=kl_penalty.item(),
                ),
                step=self.step,
            )

        return objective

    def learning_phase(self, memory: ReplayMemory) -> float:
        """
        Performs a learning step on `memory`. This involves the standard training steps
        (i.e. zeroing gradient, computing objective function, doing backprop, stepping optimizer).

        You should also remember the following:
            - Clipping grad norm to the value given in `self.args.max_grad_norm`
            - Incrementing `self.step` by 1 for each minibatch
            - Stepping the scheduler (once per calling of this function)

        Returns the average objective function value over the minibatches as a float for logging.
        """
        total = 0.0
        minibatches = memory.get_minibatches()

        # An inner progress bar under train()'s outer one: leave=False removes it when the phase ends, and
        # delay=0.5 only shows it if the phase takes longer than half a second.
        for minibatch in tqdm(minibatches, desc="learning", leave=False, delay=0.5):
            self.optimizer.zero_grad()
            objective = self.compute_grpo_objective(minibatch)
            objective.backward()
            nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.args.max_grad_norm)
            self.optimizer.step()
            self.step += 1
            total += objective.item()

        self.scheduler.step()
        return total / len(minibatches)

    def train(self) -> None:
        """
        Performs a full training run.
        """
        self.step = 0

        if self.args.use_wandb:
            wandb.init(
                project=self.args.wandb_project_name,
                entity=self.args.wandb_entity,
                name=self.run_name,
                config=self.args,
            )
        runner = tqdm(range(self.args.total_phases))
        for self.phase in runner:
            memory = self.rollout_phase()
            objective = self.learning_phase(memory)
            runner.set_description(f"Objective: {objective:.4f}")

        if self.args.use_wandb:
            wandb.finish()

Once you've implemented your trainer class, you can run the code below to train your model. We recommend you start with the test run below, using a KL coefficient of zero.

Question - with kl_coef=0.0, what results do you think you should reliably get?

With this KL coefficient, the model has no incentive to match the reference distribution, it will only try to maximize the reward. So once it's figured out that it can just output full stops all the time and totally abandon any kind of grammar or coherence, it will do this. Within the first few phases (usually by phase 3 or 4 of the 6), the model should have collapsed into producing reward-maximizing output like "This is......", or something close. Once every sample in the batch is identical, all the GRPO advantages are zero, so the objective reads 0 and training stops changing anything.

# Testing your setup: kl_coef=0.0 (see the dropdown above for explanation)
if RUN_BASE_RLHF:
    model = None  # free the demo model from earlier
    args = GRPO_Args(
        kl_coef=0.0,
        lr=2e-3 if LOW_GPU_MEM else 1e-3,  # gpt2 (small) needs a higher lr to collapse as reliably
        batch_size=64,
        num_minibatches=1,
        batches_per_learning_phase=1,
        total_phases=6,
        warmup_steps=0,
        reward_fn=reward_fn_char_count,
    )
    trainer = GRPO_Trainer(args)
    trainer.train()
else:
    print(f"{RUN_BASE_RLHF=}, skipping test run")

Click to see example output produced from the test run above

Phase 001/010, Mean reward: 1.0625, KL: 0.0000
┌──────────┬────────────────┬─────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                  │
├──────────┼────────────────┼─────────────────────────────────────────────────────────────────────────────────────────┤
│        2 │        -111.16 │ 'This is by far the best Pokemon TCG I have ever played on. Unique art resulting in     │
│          │                │ spectacle is what this system deserves. I have had similar taste in'                    │
├──────────┼────────────────┼─────────────────────────────────────────────────────────────────────────────────────────┤
│        0 │        -147.71 │ "This is going to be rough lead me in with tonight's brewing challenge\n\nToday was our │
│          │                │ first critical base for a flank Brewing challenge!\n\n\n7 breweries"                    │
├──────────┼────────────────┼─────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -125.73 │ 'This is critical to understand because using "Adult Minority Ratings" is not actually  │
│          │                │ being nominated or selected for anything.\n\nI am apolitical, like you and'             │
└──────────┴────────────────┴─────────────────────────────────────────────────────────────────────────────────────────┘

Phase 002/010, Mean reward: 3.5781, KL: 1.0949
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│        6 │        -133.78 │ 'This is a rush transcript. Copy will not function.\n\nFinal release... this has been      │
│          │                │ confirmed to me. Anonymous is drunk and I think only he understands'                       │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│        4 │         -90.56 │ "This is a rush job. I don't want to copy and paste or anything. This is what I pulled it  │
│          │                │ out of. Not sure how it goes."                                                             │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│        5 │         -83.68 │ "This is not what I was expecting. I was expecting it to be sealed, but what I got doesn't │
│          │                │ attach to it. I'm not sure if..."                                                          │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Phase 003/010, Mean reward: 10.3125, KL: 2.8364
┌──────────┬────────────────┬──────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                   │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        9 │        -105.72 │ 'This is not what you are looking for. Should be good.\n\nOpen batch... Should be        │
│          │                │ good.\n\nViewing...\n\nRelated'                                                          │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        8 │        -106.9  │ 'This is not exactly what you were expecting.... should be good.Absolutely no. Should be │
│          │                │ good. Should be good.\n\nInformation on this page'                                       │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│       38 │        -145.35 │ "This is not a rush... >.... response.... Friday email.... final.... cool........        │
│          │                │ stuff....'s... deal...."                                                                 │
└──────────┴────────────────┴──────────────────────────────────────────────────────────────────────────────────────────┘

Phase 004/010, Mean reward: 89.8125, KL: 1.7582
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       90 │         -53.15 │ 'This is.................................................................................. │
│          │                │ ........'                                                                                  │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       90 │         -53.15 │ 'This is.................................................................................. │
│          │                │ ........'                                                                                  │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       90 │         -53.15 │ 'This is.................................................................................. │
│          │                │ ........'                                                                                  │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

...

Phase 010/010, Mean reward: 90.0000, KL: 1.7716
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       90 │         -53.15 │ 'This is.................................................................................. │
│          │                │ ........'                                                                                  │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       90 │         -53.15 │ 'This is.................................................................................. │
│          │                │ ........'                                                                                  │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       90 │         -53.15 │ 'This is.................................................................................. │
│          │                │ ........'                                                                                  │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Once the smoke test has passed and you can see that the model mode collapses to always generating periods, you can do a real run with the KL penalty switched on. Now the model has to be clever, and find a way to launder as many periods as possible in the output text, in a way that the KL penalty doesn't notice (i.e. it should look like text that the base model plausibly could have generated).

if RUN_BASE_RLHF:
    trainer = None
    gc.collect()
    t.cuda.empty_cache()
    # CUDA errors? reduce batch_size or gen_len
    args = GRPO_Args(
        kl_coef=0.25,
        lr=1e-4,
        num_minibatches=8,
        batches_per_learning_phase=2,
        warmup_steps=5,
        total_phases=40,
        reward_fn=reward_fn_char_count,
    )
    trainer = GRPO_Trainer(args)
    trainer.train()
else:
    print(f"{RUN_BASE_RLHF=}, skipping main run")

Click to see example output

Phase 001/100, Mean reward: 1.3438, KL: 0.0000
┌──────────┬────────────────┬───────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                    │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -167.86 │ "This is by no means an absolute proof of Eniksen's yes. However, in statements from Metz │
│          │                │ obviously challenging a business relationship from Hansen opening for Tid"                │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│        3 │        -129.54 │ "This is going to sound weird but please forgive my latin getup style, because we've more │
│          │                │ brain cells to fill...RELATED: How Our Nervous"                                           │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -120.89 │ "This is critical now because we're about to run out of meme. A little after 7PM on       │
│          │                │ Thursday January 20th, we exported the room in full,"                                     │
└──────────┴────────────────┴───────────────────────────────────────────────────────────────────────────────────────────┘

Phase 002/100, Mean reward: 1.3047, KL: 0.0027
┌──────────┬────────────────┬───────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────┤
│        2 │         -93.25 │ 'This is a dungeon update based on the original release. That means it has some new   │
│          │                │ monsters, traps and quests. However, this update has also some metagame'              │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────┤
│        2 │        -130.33 │ 'This is what we meant by "moderate time." There are lots of tips on how to suppress  │
│          │                │ trigger feelings while building requisite time to fight off the insomnia. When'       │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────┤
│        2 │         -45.09 │ "This is a rush transcript. Copy may not be in its final form.\n\nAMY GOODMAN: We end │
│          │                │ today's show with two recent pieces originally restricted to"                         │
└──────────┴────────────────┴───────────────────────────────────────────────────────────────────────────────────────┘

Phase 003/100, Mean reward: 1.4844, KL: 0.0123
┌──────────┬────────────────┬───────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                    │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -115.69 │ 'This is one of the first things you should do after any budget creation exercise by this │
│          │                │ gaming hub. Since almost every budget publisher has done something similar at different   │
│          │                │ times'                                                                                    │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -129.6  │ 'This is a dynamic virtual environment to work with large constant persistent expository  │
│          │                │ views.\n\nThe configuration for libatlas is completely driven by the datastore'           │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│        2 │        -147.08 │ 'This is a person who brings gems of damp alt threads to fresh tears for progressive      │
│          │                │ crybabies/. All hope seems to be lost.\n\nIt seems that'                                  │
└──────────┴────────────────┴───────────────────────────────────────────────────────────────────────────────────────────┘

...

Phase 098/100, Mean reward: 9.8281, KL: 0.4339
┌──────────┬────────────────┬──────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                   │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        0 │         -83.72 │ 'This is an Open Access article distributed under the terms of the Copyright Licensing   │
│          │                │ Terms of Part I of the United States Copyright law, including\n\nForeign Objection'      │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        2 │        -132.33 │ 'This is the first of a larger set of content that you might stumble upon if you refill  │
│          │                │ one of your rolls. This is HELLS GOLD .We qualified for'                                 │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        4 │        -134.45 │ 'This is no joke... for some reason "it was a movie" is really OK. -Yohei Yoshida Senior │
│          │                │ Producer Once Upon a Time US Ep'                                                         │
└──────────┴────────────────┴──────────────────────────────────────────────────────────────────────────────────────────┘

Phase 099/100, Mean reward: 11.6953, KL: 0.4764
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       12 │         -87.92 │ 'This is a conversation between Serval [Blind], willad [Blue Mage] and Black Gold. . . . . │
│          │                │ . . . . . . .'                                                                             │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│       22 │         -44.59 │ 'This is a post which is by no means true . . . . . . . . . . . . . . . . . . . . . .'     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -144.23 │ 'This is the living-memory version of the novella "Feeding Insects," written for AiH4. It  │
│          │                │ wasn\'t friends ported licensing material by'                                              │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Phase 100/100, Mean reward: 11.1094, KL: 0.4208
┌──────────┬────────────────┬──────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                   │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │        -135.13 │ 'This is a guide to educating others in orthopaedic behavior problems. The languages and │
│          │                │ audio/visual aids used in the art are Restorative Movement Seizing'                      │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        1 │         -95.24 │ 'This is going to be kind of a long read. I\'m sure there are readers out there saying   │
│          │                │ something like "numeral…", or "on numerology'                                            │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│        3 │        -118.41 │ "This is a conversation between Dr. Fedora : Sarge vs ghazi's inbrawl is really cool     │
│          │                │ .\n\nDr. Fedora : Sarge game was"                                                        │
└──────────┴────────────────┴──────────────────────────────────────────────────────────────────────────────────────────┘

Some observations on the example run above

In this run, the model's reward went from about 1.3 to around 11 periods per sample by the end, with a per-token KL to the reference model of around 0.4-0.8. Some things to notice:

Spaced-out periods. The model's favourite trick: about a quarter of the samples printed in the second half of training end in a run of . . . . . ., e.g. This is. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . scored 30 with a reference log-probability of only -23 for the whole sample. This is why repetitive outputs are hard to stop with a KL penalty alone: the KL penalty compares the policy to the reference model in the same context, and in the context of . . . ., the reference model predicts another . too. Only the first few tokens of the pattern are penalized.
A favourite prefix. This is a conversation between ... keeps coming back (about 1 in 9 of the printed samples in the second half), since a script tends to lead to short lines with lots of periods.
The period tokens. Even with the KL penalty, the model occasionally finds GPT-2's tokens made of many periods: one sample in phase 82 scored 114, at the cost of a reference log-probability of -119.
High-likelihood boilerplate. Other samples are ordinary text, including text the base model knows almost by heart, like This is a rush transcript. Copy may not be in its final form.\n\nAMY GOODMAN: ... (the standard opening of Democracy Now! transcripts). The KL penalty is tiny for these, and short formal sentences contain a few periods.
Training isn't monotonic. The mean reward swings around a lot in the second half, at one point (phase 65) falling all the way back to its starting value of 1.3 before recovering to around 11.

Look at the KL column too: the policy is allowed to drift from the reference model, but slowly. With kl_coef=0 the model finds the tokens made of many periods (rewards of 90 or more, as in the test run above), and with a larger kl_coef (try 0.6) it should stay more coherent, at the cost of fewer periods.

You can also play around with the parameters - in particular, try a few different prompts. The behaviour of the model (e.g. which kinds of techniques it converges onto for period maximization) or whether it easily mode collapses into insanity can be highly dependent on the prompt!

Some common strategies you should observe include:

Shorter sentences
Repeating U.S. or U.S.A.
Library versions e.g. Python 2.7.12 or the 2.6.0.2 release
Names with initials e.g. C. S. Lewis or titles e.g. Dr. and PhD.
Abbreviations e.g. Data-R.A.R. series or "L.A. Times"
Decimals in numbers e.g. 9.5cm x 7.5 cm
Triple periods e.g. the man . . . the woman . . .

You might also observe increasingly incoherent mode collapse if you train for too long and don't regularize with a high KL penalty. Here are a few that I got:

This is really helpful. The U.S. U.S. U.S. U.S.
This is the A.A.G.A.R.M.A.R.M.A.R.M.A.R.M
This is my mother. . . me. . . . . . . . . . . . . . . . . . . . . . . .

Exercise - use a more complex reward function

Difficulty: 🔴🔴🔴🔴⚪

Importance: 🔵🔵🔵🔵⚪

You should spend up to 30-50 minutes on this exercise.

Note: The following exercises need more VRAM, since the IMDB reward model sits on the GPU alongside the model you're training and the reference model. With gpt2-medium and gen_len=40 (as in the example runs below) they peak at about 18GB, so they fit on a 24GB GPU; with gpt2-small (LOW_GPU_MEM = True) they need less. Don't worry if you can't run them, these exercises are mostly for playing around with the reward model. You've already conceptually gained pretty much everything about RLHF if you've completed the above. We just now replace our toy reward model with something more complex.

We recommend you experiment with a few different reward functions, in particular some sentiment-based reward functions which are based on pretrained text classification models. For example, we might use one of the following:

lvwerra/distilbert-imdb, which was trained to classify IMDB film reviews as positive or negative.
cardiffnlp/twitter-roberta-base-sentiment, which is a model trained on tweets and finetuned for sentiment analysis (categories are positive, neutral and negative).
distilbert-base-uncased-emotion, which was finetuned on the Emotion Dataset for Emotion Recognition Tasks, i.e. it's trained to classify text according to emotional tone (classes are sadness, joy, love, anger, fear and surprise).

Note that for some of these, you should be using prompts which are appropriate for the reward function you're fine-tuning on, e.g. "This movie was really" for the IMDB model. Similarly, you might also want to change other parameters e.g. generation length. You can find a list of other models here. Lastly, note that it's fine to use probabilities rather than logits or logit diffs as your reward signal, since the advantage normalization means that you'll still get a good signal even as the probabilities get close to 1.

We've given you a template below, for creating a reward function from the IMDB sentiment classification model. Your job is to complete this function.

cls_model = AutoModelForSequenceClassification.from_pretrained(
    "lvwerra/distilbert-imdb", dtype=t.float16 if t.cuda.is_available() else t.float32
).to(device)
cls_tokenizer = AutoTokenizer.from_pretrained("lvwerra/distilbert-imdb")


@t.no_grad()
def reward_fn_sentiment_imdb(
    gen_sample: list[str], direction: Literal["pos", "neg"] = "pos"
) -> Float[Tensor, " batch"]:
    """
    Reward function based on sentiment classification probability from the lvwerra/distilbert-imdb
    model.

    Args:
        gen_sample (list[str]): The generated sample to evaluate.
        direction (str): The sentiment of the reward function, either "pos" or "neg".
    """
    assert direction in ["pos", "neg"], "direction should be either 'pos' or 'neg'"

    raise NotImplementedError()


# Some samples taken from the IMDB dataset used to finetune this model
samples = [
    "Just finished watching this movie for maybe the 7th or 8th time, picked it up one night previously viewed at Blockbuster and absolutely loved it, I've shown it to 4 people so far and they have enjoyed it as well.",
    "This was the most original movie I've seen in years. If you like unique thrillers that are influenced by film noir, then this is just the right cure for all of those Hollywood summer blockbusters clogging the theaters these days.",
    "I can't believe that those praising this movie herein aren't thinking of some other film.",
    "This film seemed way too long even at only 75 minutes.",
    "Really, I can't believe that I spent $5 on this movie. I am a huge zombie fanatic and thought the movie might be really good. It had zombies in it right? Was I wrong!",
]
classes = ["pos", "pos", "neg", "neg", "neg"]

sentiment = reward_fn_sentiment_imdb(samples, direction="pos").tolist()
table = [[repr(sample), cls, f"{sent:.4f}"] for sample, cls, sent in zip(samples, classes, sentiment)]
print(tabulate(table, ["Sample", "Classification", "Sentiment"], tablefmt="simple_grid", maxcolwidths=[80, None, None]))
tests.test_reward_fn_sentiment_imdb(reward_fn_sentiment_imdb)

Click to see the expected output

┌──────────────────────────────────────────────────────────────────────────────────┬──────────────────┬─────────────┐
│ Sample                                                                           │ Classification   │   Sentiment │
├──────────────────────────────────────────────────────────────────────────────────┼──────────────────┼─────────────┤
│ "Just finished watching this movie for maybe the 7th or 8th time, picked it up   │ pos              │      0.9937 │
│ one night previously viewed at Blockbuster and absolutely loved it, I've shown   │                  │             │
│ it to 4 people so far and they have enjoyed it as well."                         │                  │             │
├──────────────────────────────────────────────────────────────────────────────────┼──────────────────┼─────────────┤
│ "This was the most original movie I've seen in years. If you like unique         │ pos              │      0.9883 │
│ thrillers that are influenced by film noir, then this is just the right cure for │                  │             │
│ all of those Hollywood summer blockbusters clogging the theaters these days."    │                  │             │
├──────────────────────────────────────────────────────────────────────────────────┼──────────────────┼─────────────┤
│ "I can't believe that those praising this movie herein aren't thinking of some   │ neg              │      0.1925 │
│ other film."                                                                     │                  │             │
├──────────────────────────────────────────────────────────────────────────────────┼──────────────────┼─────────────┤
│ 'This film seemed way too long even at only 75 minutes.'                         │ neg              │      0.0106 │
├──────────────────────────────────────────────────────────────────────────────────┼──────────────────┼─────────────┤
│ "Really, I can't believe that I spent $5 on this movie. I am a huge zombie       │ neg              │      0.018  │
│ fanatic and thought the movie might be really good. It had zombies in it right?  │                  │             │
│ Was I wrong!"                                                                    │                  │             │
└──────────────────────────────────────────────────────────────────────────────────┴──────────────────┴─────────────┘

Solution

@t.no_grad()
def reward_fn_sentiment_imdb(
    gen_sample: list[str], direction: Literal["pos", "neg"] = "pos"
) -> Float[Tensor, " batch"]:
    """
    Reward function based on sentiment classification probability from the lvwerra/distilbert-imdb
    model.

    Args:
        gen_sample (list[str]): The generated sample to evaluate.
        direction (str): The sentiment of the reward function, either "pos" or "neg".
    """
    assert direction in ["pos", "neg"], "direction should be either 'pos' or 'neg'"

    tokens = cls_tokenizer(gen_sample, return_tensors="pt", padding=True, truncation=True).to(device)
    logits = cls_model(**tokens).logits
    return logits.float().softmax(dim=-1)[:, 1 if (direction == "pos") else 0]

Once you've got this working, you can try and perform an actual run on positive / negative sentiment. We recommend generating about 40 tokens per sequence, so the reward model gets a long enough generation to classify the sentiment.

MOVIE_PROMPT = "I thought this movie was"

if RUN_BASE_RLHF:
    trainer = None
    gc.collect()
    t.cuda.empty_cache()
    args = GRPO_Args(
        reward_fn=partial(reward_fn_sentiment_imdb, direction="neg"),
        prompt=MOVIE_PROMPT,
        gen_len=40,
        kl_coef=0.1,
        lr=2e-5 if LOW_GPU_MEM else 1e-5,  # gpt2 (small) plateaus below the target at 1e-5
        warmup_steps=0,
        total_phases=12,
    )
    trainer = GRPO_Trainer(args)
    trainer.train()

Click to see example output

Phase 001/060, Mean reward: 0.2768, KL: 0.0000
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.04 │        -137.39 │ 'I thought this movie was coming soon, and I was right! Now I am increasingly glad I       │
│          │                │ waited. The one preview I managed to catch was from a recently released Netflix channel    │
│          │                │ called SBS Drama Library. Hulu has released a'                                             │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.94 │        -145.39 │ 'I thought this movie was going to be fun. Okay, okay, so get the point created, then,     │
│          │                │ despite brainwashing characters and expounding an idiotic theory which no one knows how he │
│          │                │ came up with, it went'                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.94 │        -170.06 │ 'I thought this movie was critical now because we seem to be faced with so much. A         │
│          │                │ goddamned dune? Committee hearings? First signs of the government in full meltdown?        │
│          │                │ Election day strange.\n\nBut in the'                                                       │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Phase 002/060, Mean reward: 0.4252, KL: 0.0008
┌──────────┬────────────────┬───────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                    │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│     0.03 │        -156.39 │ 'I thought this movie was pretty good before we saw it, but once we saw the scenes, I     │
│          │                │ thought it was met with more love deserved for what the characters do than any budget can │
│          │                │ provide by sans gaming or a huge visual'                                                  │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│     0.95 │        -196.91 │ 'I thought this movie was weird.\n\n\nYou were wearing Karate pants, requisite for those  │
│          │                │ film sets, while White Mike Liv punish you with steel-toed boots like captains putting    │
│          │                │ their recruits on when a pirate raider'                                                   │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -167.82 │ "I thought this movie was going to be terrible.? If I was able to receive bank loan       │
│          │                │ payments when I couldn't work...... – John\n\n\nComments Off on John Culberson Writer for │
│          │                │ Killers, Bad Reputations,"                                                                │
└──────────┴────────────────┴───────────────────────────────────────────────────────────────────────────────────────────┘

Phase 003/060, Mean reward: 0.5091, KL: 0.0050
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.01 │        -150.19 │ 'I thought this movie was remarkable\n\nOverall rating: 86% (Hail 50333)\n\n\nDébut\n\n4th │
│          │                │ broadcast\n\nFirst broadcast on SKY\n\nFirst broadcast on SAT\n\nFirst'                    │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.97 │        -208.14 │ 'I thought this movie was ok. A living man finds out what a real crime really is.\n\nusing │
│          │                │ CGI, an Europeans\'re problems are actually difficult to manage character replacement /    │
│          │                │ realism digital effects\n\n"He could not'                                                  │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.02 │        -154.58 │ 'I thought this movie was honest and true is there anything wrong with that sentiment?A    │
│          │                │ fifth character from this highly-acclaimed comedy adaptation will shoot the new episode in │
│          │                │ Miami resort, Monday.\n\nThe script will be'                                               │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

...

Phase 058/060, Mean reward: 0.9377, KL: 0.0660
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -159.6  │ 'I thought this movie was stupid because anyone can think in any form of personal A&R      │
│          │                │ position and no one has to worry about liking and re-watching bullshit. Well, at least for │
│          │                │ 90% of the supply of boredom'                                                              │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.88 │        -163.2  │ 'I thought this movie was "too sexy" and I don\'t want to see more from Carrie, but K.     │
│          │                │ Huston and Jimmy Corrigan are fast friends, and he plays a chicken chewing ism discussing  │
│          │                │ sex on the'                                                                                │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -123.04 │ "I thought this movie was pretty dull despite being so much fun to make... when I first    │
│          │                │ saw it I couldn't believe I watched it by myself - the original plan was for me and my     │
│          │                │ fiance to choreograph it myself together"                                                  │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Phase 059/060, Mean reward: 0.9216, KL: 0.0545
┌──────────┬────────────────┬──────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                   │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -156.89 │ "I thought this movie was pretty much Guillermo del Toro's lowest cunning bones. Driving │
│          │                │ Jason Bourne was kind of boring and I thought that his stupid maneuvers shouldn't be     │
│          │                │ allowed until after he had fully become a monster."                                      │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│     0.86 │        -116.49 │ "I thought this movie was kind of weird. Spicer, Besançon and Cannon were nowhere to be  │
│          │                │ seen. I don't know what the U.S. government is doing over here, but I thought it really  │
│          │                │ seemed"                                                                                  │
├──────────┼────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤
│     1    │        -149.24 │ "I thought this movie was just awful and damn if I can stand it.\n\n\nI can't care about │
│          │                │ violence, profanity, age and here are some examples of untimely death -\n\n\nDo grep on  │
│          │                │ YouTube"                                                                                 │
└──────────┴────────────────┴──────────────────────────────────────────────────────────────────────────────────────────┘

Phase 060/060, Mean reward: 0.9074, KL: 0.0563
┌──────────┬────────────────┬───────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                    │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -136.6  │ 'I thought this movie was crap.\n\n\nBatman and Deathstroke do really dumb things         │
│          │                │ like:\n\n\n9:00\n\n\nYeah, lets get a hero to beat up a kid for no reason, then Bat, the' │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│     0.85 │        -137.44 │ 'I thought this movie was like "Does Nathan have enough power to survive on his own"? Is  │
│          │                │ Nathan capable of taking care of himself? What happens when he has to let Deepthroat and  │
│          │                │ Alara down and allow them to'                                                             │
├──────────┼────────────────┼───────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -167.75 │ "I thought this movie was going to be a mess in its last isometric shot of a third of the │
│          │                │ map using every reward system compromise that I had tried elsewhere. Thankfully that      │
│          │                │ didn't happen.When it comes to running Mon"                                               │
└──────────┴────────────────┴───────────────────────────────────────────────────────────────────────────────────────────┘

1️⃣ RLHF via GRPO

Learning Objectives

The language model "environment"

States, actions and episodes

Policy gradients and baselines

The GRPO objective

Reverse KL Penalty

GRPO training args

Sampling from a language model

Exercise - implement reward_fn_char_count

Exercise - brainstorm your reward function

Exercise - implement compute_grpo_advantages

Exercise - implement get_logprobs

The KL penalty

A cheaper KL

Exercise - implement calc_kl_penalty

Skipping the entropy bonus

The clipped surrogate objective

Exercise - implement calc_clipped_surrogate_objective

Memory

Optimizer & Scheduler

Training your model

Exercise - complete GRPO_Trainer

Exercise - use a more complex reward function

Exercise - implement `reward_fn_char_count`

Exercise - implement `compute_grpo_advantages`

Exercise - implement `get_logprobs`

Exercise - implement `calc_kl_penalty`

Exercise - implement `calc_clipped_surrogate_objective`

Exercise - complete `GRPO_Trainer`