Lazy loaded image8.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, the vllm engine 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:
  1. Reward Model (not updated during PPO): trained after obtaining the SFT model.
  1. Actor Model: initialized from the SFT model and trained during PPO.
  1. Critic Model: initialized from the Reward Model and trained during PPO.
  1. 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:
  1. SFT Model: This model is used to initialize both the Actor and Reference models.
    1. In this tutorial, we use SmolLM2-135M-Instruct, which was tested in Section 8.4.
  1. Reward Model: This model is used to initialize both the Critic and Reward models.
    1. 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:
Training Plot: reject and chosen reward values, loss, accuracy, learning rate, and training progress
Training Plot: reject and chosen reward values, loss, accuracy, learning rate, and training progress
During training, the evaluation phase tracks the standard deviation and mean of rewards, validation loss, accuracy, and training step:
Evaluation Plot: reward mean/std, validation loss, and evaluation accuracy
Evaluation Plot: reward mean/std, validation loss, and evaluation accuracy

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:
    • notion image
      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:
    • notion image
      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:
notion image
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
  1. Load the SmolLM2-135M-Instruct model as a language model (LM) and the SmolLM2-135M-rm-60k as a reward model (RM).
  1. 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.
  1. Combine the prompt and response using RM’s chat_template.
  1. 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:
  1. 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.
  1. 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 weights
        • sleep(level=2): releases both KV cache and weights
  1. tensor_parallel_size=1: No tensor parallelism (single GPU use).
  1. 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().
  1. 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 = -1 or top-p = 1 sampling: 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 EOS token).
  • 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:
  1. working_dir: Ray requires a working_dir to 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_dir is automatically deleted. Therefore:
        • Paths such as save_path and ckpt_path must be absolute (e.g., /home/user/ppo_training/save) to preserve saved models.
        • Ray synchronizes working_dir across 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_data should be specified as an absolute path or located in Hugging Face.
  1. xxx_num_gpus_per_node: These parameters default to 8, assuming multi-GPU setups. For Colab or any single-GPU environment, be sure to set them to 1.
  1. Critic and Reward Model Initialization:
      • If critic_pretrain_path is 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.
  1. micro_train_batch_size vs. train_batch_size:
      • Although train_batch_size appears in the parameter list, it is not used in PPO training.
      • Only micro_train_batch_size is 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.
  1. prompt_max_len and generate_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_len applies to the full sequence.
  1. vLLM Parameters: vllm_gpu_memory_utilization, vllm_enable_sleep, and enforce_eager: These were discussed in detail in the previous section. OpenRLHF uses the same vLLM settings when initializing the vLLM engine for PPO.
  1. colocate_all_models and offload_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.
  1. 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 is 0.01)
        • kl_estimator: OpenRLHF supports k1, k2, and k3 KL computation methods. A detailed explanation of KL loss and estimators is provided in Section 8.5.4.
        • kl_target and kl_horizon: Used for adaptive KL, adjusting the KL penalty over time. Check the code implementation of AdaptiveKLController for more details.
  1. save_value_network:
      • By default, the critic model is not saved after training.
      • To ensure it is saved, set the save_value_network flag to True.
      • 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:
  1. With KL regularization: It prevents the actor model from drifting too far from the original reference model by applying a KL penalty.
  1. Without KL regularization: The actor model is updated purely based on its gradient. To disable KL loss, comment out the --use_kl_loss flag.
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:
notion image
notion image

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:
  1. With KL Regularization:
    1. 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:
      1. Actor model:
        1. Critic model:
          1. Reference and Reward Models: defined in launcher.py
            1. And:
          Shared Model Interface: BasePPORole
          All 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_handlers holds 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:
          1. Experience Sampling – using the actor model and vLLM
          1. 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_size number of data), the actor’s weights are broadcast to vLLM engines for updated response sampling.
          📌
          freezing_actor_steps: Value Warmup Strategy
          The 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:
          1. Response Generation via vLLM
          1. Model Evaluation: calling actor, reference, reward, and critic models
          1. Advantage and Return Computation

          Key Steps in Experience Generation
          1. 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.
          1. 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_probs over 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
          • 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.
          1. k1: unbiased, high-variance
          1. k2: biased, lower-variance. OpenRLHF mentions that k2 is nearly equivalent to k1.
          1. 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:
          1. Runs a forward pass on the actor model.
          1. Computes the policy loss and KL loss.
          1. Combining them into a total loss.
          1. 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 size
            • L: sequence length
            • C: 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 EOS token 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...
          Article List
          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.