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.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:
- Rejected content
- Rejected score
- Chosen content
- 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.

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:
chosen_keyandrejected_key: These specify the dataset fields used for contrastive learning and must match the column names in the preprocessed dataset.
- 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

- Evaluation

<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:
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
- For both the chosen and rejected responses:
- Concatenate the prompt with the response and ensure the text ends with an
<EOS>token.
- Tokenize the combined text, truncating it if it exceeds the maximum length. If truncation occurs, the
<EOS>token may be lost.
- 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:
RewardModelinherits frombase_pretrained_model, which is a pretrained backbone excluding the output layer.
- A linear classification head
value_head_prefixis added to project the hidden states to scalar reward values.
- In the
forwardmethod, 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
- Without Margin
- 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.
The loss is defined as
loss = -F.logsigmoid(chosen_reward - reject_reward)- With Margin
- Expressed mathematically:
- The derivative of the loss is:
- To understand the strength of the learning signal, we examine the gradient magnitude:
- 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.
The loss becomes:
loss = -F.logsigmoid(chosen_reward - reject_reward - margin)The derivative of is
This is simply a sigmoid curve shifted right by .

Colab example link: OpenRLHF Reward Model SmolLM2-135M-Instruct.
Prev
8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
Next
8.5 OpenRLHF: PPO Tutorial and Code Analysis
Loading...