Lazy loaded image8.4 OpenRLHF: Reward Model Training and Code Analysis

A ready-to-use Colab environment is provided, and the source code is included at the end of this tutorial.

8.4.1 Reward Model Training

Prepare the Instructed Model

Before training the reward model, we first need a fine-tuned (SFT) model. This model serves as the basis for reward modeling and later RL training. In this example, we use the SmolLM2-135M-Instruct model. Let’s start by examining this model:
Let’s examine the structure of the model:
The last layer is a linear projection with shape (576, 49152), where 576 is the model's hidden size and 49152 is the vocabulary size. This layer maps the hidden representation of each token to a distribution over all possible output tokens.
<ins/>

Prepare Dataset

We use the preference_dataset_mixture2_and_safe_pku dataset, which contains four columns:
  1. Rejected content
  1. Rejected score
  1. Chosen content
  1. Chosen score
OpenRLHF’s reward model uses pairwise contrastive loss to compare the chosen and rejected responses. Although the dataset includes scores, they are not used in training.
notion image
Here’s one example from the dataset after applying the chat template. Each sample consists of a user prompt and a chosen and rejected responses pair:
Rejected Response:
Chosen Response:
We manually prepare the dataset by splitting it into training and evaluation subsets: 400 training examples and 16 evaluation examples. The split is saved to disk in Hugging Face's format for training.
<ins/>

Training With OpenRLHF

Instead of running deepspeed from the terminal, we invoke it within a Python script using the subprocess module for better integration with the Colab environment.
To avoid CUDA out-of-memory (OOM) issues, we use a batch size of 8 and set the maximum sequence length to 1024.
Explanation of Key Settings:
  1. chosen_key and rejected_key: These specify the dataset fields used for contrastive learning and must match the column names in the preprocessed dataset.
  1. Flash Attention (flash_attn) is disabled: Flash Attention only works on Ampere GPU architectures (RTX 30 series and newer). Since Google T4 GPUs do not support it, we disable it.
After training begins, the model is evaluated every 10 iterations. Below are sample logs from wandb showing training and evaluation progress:
  • Training
    • notion image
  • Evaluation
    • notion image
<ins/>

Load and Test the Reward Model

Unlike standard transformer models, the reward model trained by OpenRLHF cannot be loaded using Hugging Face’s from_pretrained method. Instead, it must be loaded using OpenRLHF’s custom function.
🤔
Why not use from_pretrained? The reward model saved by OpenRLHF doesn't use a standard architecture name like LlamaForCausalLM. Instead, its config.json defines a custom structure RewardModel, which is not recognized by Hugging Face Transformers:
notion image
As a result, from_pretrained cannot reconstruct the model. To load it properly, we must use OpenRLHF’s get_llm_for_sequence_regression.
The loaded reward model ends with a linear head of shape (576, 1), which is a simple classification layer that outputs a scalar reward score:
We now test the model using the same chosen and rejected responses from earlier. The input is tokenized, passed to the model, and the scalar reward score is returned:
The score for the Rejected text is -0.002090, while for the Chosen text is 0.002472. It shows that the model assigns a higher score to the chosen response:
<ins/>

8.4.2 OpenRLHF Reward Model Code Analysis

Note: This analysis is based on version v0.7.4 of the OpenRLHF repository. All code links provided below explicitly reference this release. Be aware that future updates may introduce changes not reflected in this version.

Preparing the Dataset

The __getitem__ method in RewardDataset (view source) processes each data sample as follows:
  1. For both the chosen and rejected responses:
  1. Concatenate the prompt with the response and ensure the text ends with an <EOS> token.
  1. Tokenize the combined text, truncating it if it exceeds the maximum length. If truncation occurs, the <EOS> token may be lost.
  1. Explicitly set the final token to <EOS> to guarantee its existence at the end of the sequence.
Prior to __getitem__, the dataset is preprocessed using a chat template. Afterward, the collate_fn function pads all tokenized sequences in a batch to match the length of the longest sequence.

Model Structure Analysis

The function get_llm_for_sequence_regression calls _get_reward_model, which defines the architecture of the reward model. For clarity, the version below removes the packing_samples logic. For a visual overview, you can check the reward model structure in Section 3.3.3: RLHF + PPO in LLM Alignment – Reward Model.
Key components of the model:
  1. RewardModel inherits from base_pretrained_model, which is a pretrained backbone excluding the output layer.
  1. A linear classification head value_head_prefix is added to project the hidden states to scalar reward values.
  1. In the forward method, input sequences are passed through the base model, and the output is processed by the value head. The resulting values are optionally normalized during inference.
<ins/>

Training Function & Loss Analysis

The core reward model training logic begins at line 131 of rm_trainer.py. Below is a simplified version with auxiliary features removed for clarity.
In each training step, the model receives a pair of inputs (chosen and rejected) and computes their respective reward scores. These scores are then passed to the PairWiseLoss function, which applies a contrastive loss with an optional margin (explained later):
Efficient Forward Pass: OpenRLHF uses a single forward pass for both the chosen and rejected samples via the concatenated_forward function. Inputs are padded and concatenated using concatenated_inputs to maintain alignment across batches. So, when we set the batch size to 8, we are actually passing 16 sequences to the model in each batch (8 chosen and 8 rejected responses).
The PairWiseLoss compares reward values between chosen and rejected samples. If a margin is provided, the loss encourages the chosen response to exceed the rejected one by at least that margin.
Pairwise Contrastive Loss and the Role of Margin
  1. Without Margin
    1. The loss is defined as loss = -F.logsigmoid(chosen_reward - reject_reward)
      • Starting with the loss formula:
      • We can expand the sigmoid:
      • Plugging this into the loss:
      • Minimizing this loss is equivalent to maximizing:
      • Because both the numerator and the denominator are positive, maximizing this term encourages:
        • Making as large and positive as possible.
        • Making as small and negative as possible.
  1. With Margin
    1. The loss becomes: loss = -F.logsigmoid(chosen_reward - reject_reward - margin)
      • Expressed mathematically:
      • The derivative of the loss is:
        • The derivative of is
      • To understand the strength of the learning signal, we examine the gradient magnitude:
        • This is simply a sigmoid curve shifted right by .
      • For the same difference , a larger margin yields a larger gradient, providing a stronger update signal. However, this also makes the optimization target stricter, which could lead to oscillation.
        • notion image
Prev
8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
Next
8.5 OpenRLHF: PPO Tutorial and Code Analysis
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.