LLM Learning Roadmap
✨ Awesome-Anything
🖼️ Digital Image Processing
🍃 LLM Components
🌱 LLM Pre-training
☘️ LLM Post-Training
🍀 LLM Popular Models
🪴 LLM Applications
🌿 LLM Optimization
🌾 LLM Compression
🌵 LLM Hands-on Practice
🌴 LLM Must-read Papers
🌳 LLM Q&A
🐝 VLM Image Encoders
📝 MISC.
8.5 OpenRLHF: PPO Tutorial and Code Analysis
This section provides a hands-on tutorial and code walkthrough for PPO (Proximal Policy Optimization) training for LLM using OpenRLHF, following the concepts in Section 3.3.3: RLHF + PPO in LLM Alignment. Before diving into this tutorial, it's essential to have completed the previous two sections:
⚠️ Resource Requirements:PPO training is resource-intensive. Due to its high demands on both CPU and GPU memory, free-tier Colab environments are insufficient. For example, thevllmengine fails to start with just 12.7 GB of CPU memory while it’s invoked by the OpenRLHF.To run this tutorial successfully, you will need a Colab Pro environment with access to an A100 GPU.
A Colab notebook is provided: OpenRLHF PPO Tutorial and Code Analysis.
Quick Recap: PPO Training Setup
PPO-based RLHF training requires four distinct models:
- Reward Model (not updated during PPO): trained after obtaining the SFT model.
- Actor Model: initialized from the SFT model and trained during PPO.
- Critic Model: initialized from the Reward Model and trained during PPO.
- Reference Model (not updated): a static copy of the SFT model used to compute KL divergence penalties.
Model Preparation
Before starting PPO training, you need two pretrained models:
- SFT Model: This model is used to initialize both the Actor and Reference models.
In this tutorial, we use SmolLM2-135M-Instruct, which was tested in Section 8.4.
- Reward Model: This model is used to initialize both the Critic and Reward models.
We'll use AI-Roadmap/SmolLM2-135M-rm-60k, which we trained in Section 8.4.
8.5.1 Test the Reward Model and SFT Model
Testing the Reward Model
The reward model plays a crucial role in PPO training. It provides the signal that drives policy updates, guiding the Actor (SFT model) to generate higher-quality responses based on human preferences. Since no high-quality, small-scale, open-source reward model currently exists, we trained a 135M reward model based on the SmolLM2-135M-Instruct model. You can find detailed training information in the model card
This reward model achieves an accuracy of 76.4% on the test set.
Training and Evaluation Plots: The following plot illustrates the training process, including reward signals, loss, accuracy, learning rate, and global training step:

During training, the evaluation phase tracks the standard deviation and mean of rewards, validation loss, accuracy, and training step:

Reward Model Evaluation on Holdout Data
You can try out this reward model using the Colab notebook below:
👉 Test the Reward Model
To evaluate the reward model, we use the same dataset as during training: preference_dataset_mixture2_and_safe_pku
- The first 480k samples were used for training.
- We evaluate on 800 samples from indices 480,000 to 480,800.
The evaluation function (code above) computes several metrics:
- Accuracy: proportion of times the chosen response receives a higher reward than the rejected one
- Mean rewards: for both chosen and rejected responses
- Response lengths: to analyze correlations with rewards
Evaluation Results
- The reward model correctly rewards the preferred response 77% of the time.
- Chosen responses have both higher rewards and longer lengths on average.
Visualizing Reward Model Behavior
- Reward Distribution: Chosen vs. Rejected. The histogram below shows the reward distribution for both chosen and rejected responses:

As expected, chosen responses tend to receive higher reward values.
- Correlation: Reward vs. Response Length. The histogram below shows the reward distribution for both chosen and rejected responses:

This trend is common in reward models trained on human preference data, where longer responses often provide more informative content.
<ins/>
8.5.2 Prepare the Dataset and Test the SFT Model
Prepare the Dataset
For PPO training, we use the prompt-collection-v0.1 dataset. This dataset contains approximately 179,000 prompts:

To simplify our tutorial and reduce resource consumption, we manually select a small subset of the dataset for training and testing.
Note: While OpenRLHF can support distributed setups (running the actor, critic, reference, and reward models on different GPU), this tutorial assumes a single-GPU setup for simplicity.
Dataset Preparation Code. In the snippet below, we randomly shuffle and select a subset of the full dataset:
Output:
Test the SFT Model
To assess the performance of the SmolLM2-135M-Instruct model (the SFT model), we evaluate its responses using the SmolLM2-135M-rm-60k reward model.
Evaluation Workflow
- Load the SmolLM2-135M-Instruct model as a language model (LM) and the SmolLM2-135M-rm-60k as a reward model (RM).
- For each test prompt:
- Format it using LM’s
chat_template. - Tokenize and generate a response using LM.
- Decode the generated response using the same tokenizer.
- Combine the prompt and response using RM’s
chat_template.
- Tokenize the full input and compute a reward using RM.
Note: In our example, the LM and RM share the same tokenizer and chat template since the reward model was trained using the SFT model as a base.
Performance Note on Token-by-Token Generation
Using Hugging Face’s
model.generate() can be extremely slow, especially for longer sequences. This method generates one token at a time, which becomes computationally expensive when scaled:- 160 prompts × 400 tokens per output = 64,000 decoding iterations
To optimize inference speed, we use vLLM, which enables batched and parallelized generation.
However, vLLM may not fully release GPU or CPU memory after use in the Colab environment. To manage memory effectively, it is highly recommended to restart the Colab runtime between blocks to clear the vLLM and Torch cache.
Evaluation Code and Results: The following script loads the test prompts, generates responses using vLLM, and evaluates them using the reward model:
Output:
This indicates that the SFT model generates coherent and reward-aligned responses, with an average reward of 0.1049 over 160 prompts.
Understanding vLLM Parameters
vLLM includes several configuration options that influence performance and memory usage. Below are explanations of key parameters:
enforce_eager=True- Forces PyTorch to use eager execution (good for debugging, slower in practice).
- Disabling it allows vLLM to compile a CUDA graph, which is faster but consumes significantly more memory (10+ GB on first use).
- Ideal for stable batch workloads with consistent input shapes.
enable_sleep_mode=True- vLLM does not release GPU memory by default to reduce latency between calls.
- This option allows manual memory release via:
sleep(level=1): releases KV cache, keeps model weightssleep(level=2): releases both KV cache and weights
tensor_parallel_size=1: No tensor parallelism (single GPU use).
gpu_memory_utilization=0.8- Reserves 80% of GPU memory in advance (for vLLM’s custom memory manager).
- This memory is retained unless manually released with
sleep().
dtype="bfloat16": Ensure your GPU supports BF16 (A100 and newer). Older GPUs (e.g., V100, T4) simulate BF16 using float32, which consumes more memory and reduces speed.
Sampling Parameters and Generation Strategy
To ensure consistent and high-quality outputs, we configure generation as follows:
top_k = 1: Enables greedy decoding, selecting only the most probable token at each step.
top_k = -1ortop-p = 1sampling: Enables random sampling, which introduces variability but often reduces reward performance.
Greedy vs. Random Sampling
Setting | Avg Reward | Reward Variance |
top_k = -1 | < -0.1 | High variance, different average reward for each run |
top_k = 1 | 0.1049 | Consistant value, able to reproduce |
In practice, greedy decoding (top_k=1) produces more reliable outputs for evaluation and is recommended for stable benchmarking.
<ins/>
8.5.3 OpenRLHF PPO Training
OpenRLHF’s PPO training implementation is more complex than the traditional RL PPO algorithm due to its use of advanced infrastructure for scalability and efficiency. It integrates:
- vLLM for fast response generation (sampling)
- Ray for distributed training across processes or machines
This complexity can make the codebase harder to understand. The following explanation and pseudo-code are based on OpenRLHF v0.7.5.
PPO Training: Simplified Pseudo Code
Below is a simplified version of the OpenRLHF PPO training loop. It abstracts away low-level details like model parallelism, memory handling, and distributed coordination.
Key Notes:
- vLLM handles response generation efficiently, allowing batched token generation without step-by-step decoding.
- KL divergence is applied between the actor and reference log-probabilities to constrain policy shifts.
- Rewards are adjusted with the KL penalty. Only the final token gets the actual reward from the reward model, and intermediate steps use a KL penalty only.
- The critic model learns to predict value estimates across response tokens (excluding the final
EOStoken).
- GAE (Generalized Advantage Estimation) is used to compute the advantage in our example.
- Mini-batching is applied twice:
- During rollout (
micro_rollout_batch_size) - During training (
micro_train_batch_size)
<ins/>
Ray and Training Parameters
Configuring OpenRLHF’s training parameters can be challenging for starters, especially when using Ray for distributed execution. In this tutorial, we focus on single-GPU setups, such as Colab environments. Multi-GPU or multi-node training introduces additional complexity, requiring careful choosing of parameters.
Ray introduces several unique behaviors and configuration requirements. Below is a summary of key considerations when using Ray with OpenRLHF:
working_dir: Ray requires aworking_dirto act as the root directory for execution. This directory can be an absolute or relative path, but keep the following in mind:- All relative paths in the Python code executed by Ray are resolved relative to
working_dir. - All files generated during training are stored here.
- After training completes,
working_diris automatically deleted. Therefore: - Paths such as
save_pathandckpt_pathmust be absolute (e.g.,/home/user/ppo_training/save) to preserve saved models. - Ray synchronizes
working_diracross nodes, so it should be kept small. - Do not store your dataset in
working_dir, or Ray may fail due to excessive size. So,prompt_datashould be specified as an absolute path or located in Hugging Face.
xxx_num_gpus_per_node: These parameters default to8, assuming multi-GPU setups. For Colab or any single-GPU environment, be sure to set them to1.
- Critic and Reward Model Initialization:
- If
critic_pretrain_pathis not set, and the reward model is local (not accessed via API), OpenRLHF will initialize the critic model from the reward model. - If the reward model is accessed via API, and no local model is available, OpenRLHF will fall back to initializing the critic model from the SFT model.
micro_train_batch_sizevs.train_batch_size:- Although
train_batch_sizeappears in the parameter list, it is not used in PPO training. - Only
micro_train_batch_sizeis used (as shown in the pseudo code) to define the training batch size for the actor and critic. - This value should be carefully tuned based on your GPU memory.
prompt_max_lenandgenerate_max_len:prompt_max_len: Maximum token length of input prompts. Prompts longer than this will be truncated by the tokenizer.generate_max_len: Maximum number of tokens generated by vLLM during response generation.- The total sequence length (prompt + response) is therefore
prompt_max_len + generate_max_len. - This differs from reward model training, where
max_lenapplies to the full sequence.
- vLLM Parameters:
vllm_gpu_memory_utilization,vllm_enable_sleep, andenforce_eager: These were discussed in detail in the previous section. OpenRLHF uses the same vLLM settings when initializing the vLLM engine for PPO.
colocate_all_modelsandoffload_states:colocate_all_models: Indicates that all four models (actor, reference, reward, critic) are run on a single device.- During experience collection, OpenRLHF clears the GPU cache after running each model, but retains model weights in memory.
- During training, it also clears the cache between steps, retaining weights to save reload time.
offload_states: If enabled (and running with DeepSpeed ZeRO Stage 3), OpenRLHF will offload model weights to the CPU during idle phases, then reload them during training. You can read the implementation code.
- KL Regularization:
use_kl_loss,init_kl_coef,kl_estimator,kl_target,kl_horizon - By default, OpenRLHF does not apply KL loss unless
use_kl_loss=True. To enable KL divergence between the actor and reference models, set this flag. - Other parameters:
init_kl_coef: Initial KL coefficient (default value is0.01)kl_estimator: OpenRLHF supportsk1,k2, andk3KL computation methods. A detailed explanation of KL loss and estimators is provided in Section 8.5.4.kl_targetandkl_horizon: Used for adaptive KL, adjusting the KL penalty over time. Check the code implementation of AdaptiveKLController for more details.
save_value_network:- By default, the critic model is not saved after training.
- To ensure it is saved, set the
save_value_networkflag toTrue. - The critic model will be saved to
save_path + "_critic".
OpenRLHF PPO Training
This section demonstrates how to launch PPO training using OpenRLHF on a single GPU (e.g., Colab). We experiment with two training setups:
- With KL regularization: It prevents the actor model from drifting too far from the original reference model by applying a KL penalty.
- Without KL regularization: The actor model is updated purely based on its gradient. To disable KL loss, comment out the
--use_kl_lossflag.
Launching PPO Training via Ray: The following script sets up a local Ray head node, defines runtime parameters, and submits a PPO training job. This configuration is adapted for single-node, single-GPU environments.
Training Process:
.png?table=block&id=1f826e5a-7de0-80ee-bbf4-fd3cc5372cc9&t=1f826e5a-7de0-80ee-bbf4-fd3cc5372cc9)

Load and Test the RLHF Model
The evaluation procedure for the RLHF-trained model is nearly identical to the one used for the SFT model. To avoid duplication, we won’t repeat the code here. You can refer to the Colab notebook for full implementation details.
We evaluate two versions of the RLHF model:
- With KL Regularization:
- Without KL Regularization:
Summary Comparison
Model | SFT model | RLHF model with KL | RLHF model without KL |
Average reward | 0.1049 | 0.1201 | 0.1790 |
Average response length | 460.02 | 483.32 | 471.57 |
Interpretation:
- After PPO training, both RLHF models outperform the SFT model. The response length also increases, particularly when KL regularization is used.
- Disabling KL regularization results in the highest reward. This suggests that while KL regularization helps maintain stability and conservativeness, it can slow learning or limit exploration during training. This is also the key reason that DAPO removes KL regularization (see Section 3.4.5 DAPO).
<ins/>
8.5.4 OpenRLHF PPO Code Analysis
In this section, we provide a detailed walkthrough of the actual PPO training implementation in OpenRLHF, based on the simplified pseudo-code introduced earlier. The analysis is based on OpenRLHF v0.7.5 and reflects the complexity of the real system, which is significantly more intricate than the simplified logic.
We won’t cover the dataloader components (e.g., prompt dataloader, experience list dataloader), as they are conceptually similar to those discussed in previous chapters.
Ray and Parallelization
OpenRLHF uses Ray to distribute all four core models across different processes or even machines. Each model is defined as a Ray remote class, which enables it to run asynchronously and in parallel:
- Reference and Reward Models: defined in launcher.py
And:
Shared Model Interface:
BasePPORoleAll four models inherit from the common base class
BasePPORole and implement a method called init_model_from_pretrained(), which initializes the model from a pretrained checkpoint. This consistent interface allows OpenRLHF to manage models uniformly, regardless of their specific roles.Model Group Coordination:
PPORayActorGroup. The PPORayActorGroup class combines the four model roles into a unified group:Here:
- The term “actor” refers generically to each model (not just the policy model).
self._actor_handlersholds Ray remote object references to each model instance.
.remote()is called to execute model initialization asynchronously.
- This design allows each model to be initialized in parallel, potentially on different machines or GPUs.
Synchronization with
ray.get(). Finally, in train_ppo_ray.py (line 154), all these remote model initializations are launched and synchronized using ray.get():This code shows the typical pattern in OpenRLHF PPO training:
- Launch remote method calls in a batched, asynchronous manner.
- Collect the resulting object references.
- Call
ray.get()at a synchronization point to wait for all remote tasks to complete.
This approach is used throughout the training pipeline to ensure scalability and parallel execution across compute resources.
<ins/>
Overall PPO Trainer
To focus on the essential logic behind training, we omit auxiliary functions and user input validation. The core training loop begins at
ppo_trainer.py line 161. The PPO trainer divides the workflow into two main stages:- Experience Sampling – using the actor model and vLLM
- Policy Optimization – training the actor and critic models using PPO
The actor and critic models append their training data in parallel, and then synchronize via
ray.get().PPO Training Logic:
ppo_train(). Within the ppo_train() method, OpenRLHF manages training behavior depending on whether all models are colocated on the same device. We simplify the function by removing conditions related to ZeRO Stage 3 offloading and focusing only on the key logic:Parallelism and Synchronization
- If
colocate_all_models=True, OpenRLHF waits for the critic model training to complete before starting actor model training.
- If the models are not colocated, the actor and critic can be trained in parallel, and synchronization is deferred until both finish.
- After the actor model finishes training using all the experience in the list (
rollout_batch_sizenumber of data), the actor’s weights are broadcast to vLLM engines for updated response sampling.
freezing_actor_steps: Value Warmup StrategyThe
freezing_actor_steps parameter introduces a warm-up period for the critic (value) model, during which the actor model is frozen and not updated. By default, this is set to -1 (i.e., no warm-up).This idea was proposed by the VC-PPO paper, as discussed in Section 3.4.6: VAPO – Mitigating Value Model Bias.
The purpose of this strategy is to reduce value function bias by pretraining the critic before using it to compute advantage estimates that guide the actor’s updates.
This aligns with techniques in advanced PPO variants like VAPO, where reliable value targets improve stability in long-sequence reward optimization.
Experience Maker
The
Experience class in OpenRLHF defines the structure used to store and pass training data during PPO optimization. You can view its definition here:Although the KL divergence is computed during experience creation (as the difference between actor and reference models), it is only used to calculate rewards. The KL penalty applied during backpropagation is recalculated dynamically using the current actor model.
How Experience Is Built
The full experience creation pipeline is implemented in the
make_experience_list() function. This function contains the three main stages of experience creation:- Response Generation via vLLM
- Model Evaluation: calling actor, reference, reward, and critic models
- Advantage and Return Computation
Key Steps in Experience Generation
generate_samples: This function calls the vLLM engine to generate responses for the given prompts. It also prepares input sequences, performs padding, and computes attention and action masks.
make_experience: This step involves calling all four models in batch:- Reward model: assigns a scalar reward to each response.
- Actor model: computes
action_log_probsover response tokens. - Critic model: outputs value estimates.
- Reference model: computes baseline
action_log_probs.
The outputs are stored in
Experience objects, along with metadata (KL divergence, reward, length, etc).3.
compute_advantages_and_returns: This function computes the advantage estimates using the selected strategy (default: GAE, Generalized Advantage Estimation). Steps include:- Recomputing rewards with KL penalties via
compute_reward
- Computing GAE using
get_advantages_and_returns(see Section 3.2.11 PPO Algorithm).
- Normalizing the advantages across the batch
If another estimator is configured (e.g.,
reinforce_baseline, rloo), the appropriate method is used instead of GAE.Reward Computation Logic
The
compute_reward() function determines how the KL penalty and scalar reward are merged into a final reward tensor:- For most timesteps, the reward is just
KL.
- At the final action step, the base reward from the reward model is added.
This mirrors the strategy described in Section 3.2.11, where only the final
EOS action receives the actual scalar reward, and prior tokens are penalized solely based on KL.<ins/>
KL Loss in OpenRLHF
OpenRLHF provides three KL divergence estimators,
k1, k2, and k3. The default option is k1, while k3 is used in GRPO. See Section 3.4.1 GRPO - GRPO Loss Analysis for more information about k1 and k3.Let
This represents the difference in log-probabilities between the actor and reference models.
k1: unbiased, high-variance
k2: biased, lower-variance. OpenRLHF mentions thatk2is nearly equivalent tok1.
k3: unbiased, reduced variance
In the following code,
log_probs_base is , and log_probs is .Actor Training and Loss Calculation
In the previous section on the Overall PPO Trainer, we saw that before calling the
fit() function, OpenRLHF first appends experiences to both the actor and critic models. These models share a common fit() function implementation:This function clears the CUDA cache before and after training, calls the training logic (
ppo_train()), and finally resets the replay buffer.Inside
ppo_train(): The ppo_train() method builds a dataloader from the replay buffer (which stores collected experiences), then iterates over it using the configured batch size (micro_train_batch_size) and number of epochs (max_epochs).After training, it aggregates the per-batch statistics (stored in
status) by averaging them across epochs.Note on
strategy.all_reduce()This function aggregates metrics across all GPUs (with default operator
mean), using PyTorch’s all_reduce to ensure consistent model behavior in distributed settings. For example, output on 4 processes (each on 1 GPU), each rank (process) starts with:- Rank 0:
[1.0]
- Rank 1:
[2.0]
- Rank 2:
[3.0]
- Rank 3:
[4.0]
After
all_reduce(SUM), all ranks get [10.0]The Training Step: The actual training logic for each batch occurs inside
training_step(), which:- Runs a forward pass on the actor model.
- Computes the policy loss and KL loss.
- Combining them into a total loss.
- Performs backpropagation and optimizer updates.
actor_loss: Computed using PPO’s clipped objective.
kl_loss: Computed using the selected KL estimator (e.g.,k1,k2,k3).
total_loss = actor_loss + kl_loss * kl_ctl: Combined loss used for backpropagation.
Policy Loss Formula: The PPO policy loss used here follows the standard clipped objective from the original PPO method:
The implementation of this formula is in the
PolicyLoss class:<ins/>
Critic Model Structure Analysis
Brief Recap: In Section 8.4.2 – OpenRLHF Reward Model Code Analysis: Model Structure, we examined how the reward model modifies the standard language model architecture by replacing the original output head with a linear projection head of shape
(hidden_size, 1).- The original language model produces outputs of shape
(B, L, C), where: B: batch sizeL: sequence lengthC: vocabulary size
- The reward model, by contrast, outputs values of shape
(B, L, 1), and is simplified to(B, L).
- During forward inference, only the reward corresponding to the final
EOStoken is used to represent the total sequence reward, reducing the output to shape(B,).
Critic vs Reward Model Behavior: The critic model in OpenRLHF is built on the same model structure as the reward model, which is why it can be initialized directly from a reward model checkpoint. However, its forward behavior is different:
- Instead of extracting the final token’s value (as the reward model does), the critic:
- Produces value estimates at all time steps (excluding the final EOS token).
- Uses these values to compute advantages during PPO training.
- Only the tokens corresponding to the generated response (as marked by
action_mask) are valid for training.
Critic Training and Loss Calculation
The critic model shares much of its structure and training flow with the actor model. However, its
training_step() method differs slightly, focusing on value regression.In this step:
- The model predicts value estimates for each token in the generated response.
- These predictions are compared with target returns, which are computed using Generalized Advantage Estimation (GAE).
- The difference is minimized via a value loss function, typically mean squared error (MSE).
Value Loss Function: The critic uses the standard PPO value loss (described in
Section 3.3.3 – RLHF + PPO in LLM Alignment: Critic Model.), optionally with clipping to prevent excessive deviation from previous value predictions.
If
eps_clip is set (by default, it is set as 0.2), the new values are clipped to a range around the old values. Then, two versions of the squared error are calculated:surr1: using clipped values
surr2: using raw values
- The final loss is the maximum of both.
Prev
8.4 OpenRLHF: Reward Model Training and Code Analysis
Next
9.1 Technical Reports Explained
Loading...