Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

RLHF Pipeline › Stage 3: PPO Fine-Tuning

Michael BrenndoerferDecember 29, 20258 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 280RLHF Pipeline › Stage 3: PPO Fine-Tuning

Reading section 5 of 11 · Stage 3: PPO Fine-Tuning · 8 minSection 5 of 11 · 8 min

Section 5 of 11 · 8 min read

RLHF Pipeline › Stage 3: PPO Fine-TuningOne section · Sections in this chapter5 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

With the SFT model and reward model ready, we can now run PPO optimization. This stage adjusts the policy to maximize expected reward while staying close to the SFT model's behavior. The challenge here is delicate: we want the model to improve according to the reward signal without losing the coherent language abilities it acquired during pretraining and SFT.

Think of PPO fine-tuning as the difference between a musician who learns jazz improvisation and one who simply memorizes jazz recordings. The SFT model knows how to play in the right style. The reward model has internalized the judgment of experienced listeners. PPO optimization uses that judgment signal to push the musician toward improvising in ways that experienced listeners prefer, while the KL penalty ensures the musician does not completely abandon the musical training that makes their playing coherent in the first place. The constraint is not a limitation but a necessity: a musician who has thrown out all musical structure in pursuit of applause is no longer a musician.

The PPO stage runs for far fewer gradient steps than pretraining or even SFT, often fewer than 10,000 steps compared to hundreds of billions of tokens for pretraining. This brevity is both a feature and a constraint. The short training duration means PPO cannot introduce catastrophically wrong behaviors if properly regularized, but it also means the policy cannot make dramatic departures from the SFT starting point. The alignment effect is best thought of as a refinement of existing capabilities rather than an acquisition of new ones. The capabilities must already exist in the SFT model; PPO merely tunes the probability distribution over responses to favor the higher-quality ones.

The PPO Training LoopLink Copied

As we detailed in the PPO for Language Models chapter, each training iteration involves a carefully orchestrated sequence of steps that together enable stable policy improvement:

Sampling: Generate responses from the current policy
Reward computation: Score responses using the reward model
Advantage estimation: Compute advantages using GAE
Policy update: Optimize the clipped surrogate objective. This iterative process gradually shifts the policy's behavior toward responses that score higher according to the reward model, while the various stability mechanisms in PPO prevent the optimization from taking steps that are too large or in harmful directions.

In[17]:

Code

import torch
import torch.optim


class RLHFTrainer:
    """Complete RLHF trainer implementing the PPO training loop."""

    def __init__(
        self, policy_model, ref_model, reward_model, tokenizer, config
    ):
        self.policy = policy_model
        self.ref_model = ref_model  # Frozen copy of SFT model
        self.reward_model = reward_model
        self.tokenizer = tokenizer
        self.config = config

        self.optimizer = torch.optim.AdamW(
            self.policy.parameters(),
            lr=config["learning_rate"],
            weight_decay=config["weight_decay"],
        )

    def generate_responses(self, prompts, max_length=256):
        """Generate responses from current policy."""
        self.policy.eval()

        responses = []
        log_probs_list = []

        with torch.no_grad():
            for prompt in prompts:
                input_ids = self.tokenizer.encode(prompt, return_tensors="pt")

                output_ids = []
                log_probs = []

                for _ in range(max_length):
                    outputs = self.policy(input_ids)
                    next_token_logits = outputs.logits[:, -1, :]
                    probs = torch.softmax(next_token_logits, dim=-1)

                    next_token = torch.multinomial(probs, num_samples=1)

                    token_log_prob = torch.log(probs[0, next_token.item()])

                    output_ids.append(next_token.item())
                    log_probs.append(token_log_prob.item())

                    input_ids = torch.cat([input_ids, next_token], dim=1)

                    if next_token.item() == self.tokenizer.eos_token_id:
                        break

                responses.append(self.tokenizer.decode(output_ids))
                log_probs_list.append(log_probs)

        return responses, log_probs_list

Computing the Complete RewardLink Copied

The total reward used to update the policy combines two competing objectives. We want to maximize the reward model's score, which represents human preferences, but we also want to prevent the policy from straying too far from the reference model, which represents stable, coherent language generation. The KL divergence penalty provides this regularization, creating a tug-of-war that encourages improvement without catastrophic drift.

We'll explore the KL divergence penalty in detail in the next chapter, but the basic formulation captures this balance mathematically:

R_{total}(x, y) = R_{RM}(x, y) - \beta \cdot D_{KL}(\pi_\theta || \pi_{ref})

where:

$R_{total}(x, y)$ : the combined reward used to update the policy
$x$ : the input prompt
$y$ : the generated response
$R_{RM}(x, y)$ : the preference score from the reward model
$\beta$ : the KL penalty coefficient controlling regularization strength
$D_{KL}$ : the Kullback-Leibler divergence between the two distributions
$\pi_\theta$ : the current policy model
$\pi_{ref}$ : the reference model (frozen SFT model)

The KL coefficient $\beta$ acts as a dial controlling the trade-off between reward maximization and behavioral stability. A larger $\beta$ keeps the policy closer to the reference model, preserving more of the original capabilities but potentially limiting how much the model can improve. A smaller $\beta$ allows more aggressive optimization toward higher rewards, but risks the policy finding reward model exploits or losing coherence. This formulation ensures that while we maximize the preference score, we maintain the linguistic coherence and knowledge of the original model.

In practice, $\beta$ is one of the most sensitive hyperparameters in the entire RLHF pipeline. Values too close to zero produce unstable training because the policy has no incentive to stay in the region of coherent language. Values too large make the training too conservative, and the policy barely moves from its SFT initialization despite many optimization steps. Typical values range from 0.01 to 0.1 for most configurations, though some practitioners use adaptive schemes that increase $\beta$ automatically when the KL divergence exceeds a threshold. This provides a soft constraint instead of a fixed penalty.

Out[18]:

Visualization

Stacked bar chart showing how total reward is composed of reward model score minus KL penalty across training steps. — Reward composition in RLHF. The total reward (blue line) balances the raw preference score (green) against a KL divergence penalty (red). As the policy diverges from the reference model, the growing penalty prevents catastrophic drift.

In[19]:

Code

import torch


def compute_rewards_with_kl_penalty(
    policy_model,
    ref_model,
    reward_model,
    prompts,
    responses,
    tokenizer,
    kl_coef,
    device,
):
    """Compute total reward including KL penalty."""

    full_texts = [p + r for p, r in zip(prompts, responses)]
    tokens = tokenizer(full_texts, padding=True, return_tensors="pt").to(device)

    with torch.no_grad():
        rm_rewards = reward_model(tokens["input_ids"], tokens["attention_mask"])

    with torch.no_grad():
        policy_outputs = policy_model(
            tokens["input_ids"], attention_mask=tokens["attention_mask"]
        )
        ref_outputs = ref_model(
            tokens["input_ids"], attention_mask=tokens["attention_mask"]
        )

        policy_logprobs = torch.log_softmax(policy_outputs.logits, dim=-1)
        ref_logprobs = torch.log_softmax(ref_outputs.logits, dim=-1)

        token_kl = (
            policy_logprobs.exp() * (policy_logprobs - ref_logprobs)
        ).sum(dim=-1)
        kl_penalty = token_kl.sum(dim=-1)

    total_rewards = rm_rewards - kl_coef * kl_penalty

    return total_rewards, rm_rewards, kl_penalty

PPO Update StepLink Copied

The policy update uses the clipped surrogate objective from the PPO Algorithm chapter. This objective function represents the core mechanism that enables stable policy improvement: rather than directly maximizing expected reward, which could lead to catastrophically large updates, PPO constrains how much the policy can change in a single step. The clipping mechanism ensures that even if the advantage estimates suggest a large improvement, the actual policy update remains bounded.

In[20]:

Code

import torch


def ppo_update(
    policy_model, optimizer, batch, clip_epsilon=0.2, entropy_coef=0.01
):
    """Perform PPO policy update with clipped objective."""
    policy_model.train()

    states = batch["states"]
    actions = batch["actions"]
    old_log_probs = batch["old_log_probs"]
    advantages = batch["advantages"]
    returns = batch["returns"]

    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

    outputs = policy_model(states)
    logits = outputs.logits

    log_probs = torch.log_softmax(logits, dim=-1)
    action_log_probs = log_probs.gather(-1, actions.unsqueeze(-1)).squeeze(-1)

    ratios = torch.exp(action_log_probs - old_log_probs)

    surr1 = ratios * advantages
    surr2 = torch.clamp(ratios, 1 - clip_epsilon, 1 + clip_epsilon) * advantages
    policy_loss = -torch.min(surr1, surr2).mean()

    probs = torch.softmax(logits, dim=-1)
    entropy = -(probs * log_probs).sum(dim=-1).mean()
    entropy_loss = -entropy_coef * entropy

    total_loss = policy_loss + entropy_loss

    optimizer.zero_grad()
    total_loss.backward()
    torch.nn.utils.clip_grad_norm_(policy_model.parameters(), max_norm=1.0)
    optimizer.step()

    return {
        "policy_loss": policy_loss.item(),
        "entropy": entropy.item(),
        "mean_ratio": ratios.mean().item(),
        "clip_fraction": (
            (ratios < 1 - clip_epsilon) | (ratios > 1 + clip_epsilon)
        )
        .float()
        .mean()
        .item(),
    }

The entropy bonus term in the loss function deserves particular attention. By adding a small positive reward for entropy (the entropy coefficient $0.01$ in the code above multiplied by the average entropy of the action distribution), we discourage the policy from collapsing to deterministic outputs too quickly. Think of entropy as a measure of how many different responses the model considers at each step: high entropy means the model is still exploring a diverse set of continuations, while low entropy means it has converged to nearly always predicting the same tokens. Without the entropy bonus, PPO training tends to produce overly deterministic policies that perform poorly on novel prompts where the training distribution does not provide clear guidance.

The advantage normalization step (subtracting the mean and dividing by the standard deviation) is a standard variance reduction technique that makes training more stable. Raw advantages can have high variance because they depend on the absolute reward values, which may shift considerably across training steps. Normalizing to zero mean and unit variance at each batch ensures that the effective learning rate remains consistent even as the reward distribution changes during training.