Implementation:OpenRLHF OpenRLHF Train PRM

Knowledge Sources	OpenRLHF Process Reward Models
Domains	Reward_Modeling, Training, CLI
Last Updated	2026-02-07 10:40 GMT

Overview

Concrete tool for launching Process Reward Model (PRM) training from the command line.

Description

The train function is the entry point for PRM training. It initializes a DeepSpeed distributed strategy, loads an Actor model (used as the backbone for step-level reward prediction), prepares ProcessRewardDataset with placeholder token-based labeling, configures optimizer and scheduler, and launches ProcessRewardModelTrainer.fit(). It supports LoRA/QLoRA, packing samples via Flash Attention, gradient checkpointing, and configurable reward tokens for step-level labels.

Usage

Use this entry point to train a process reward model that assigns rewards at each reasoning step (rather than only at the final answer). It is invoked via python -m openrlhf.cli.train_prm or through DeepSpeed launcher.

Code Reference

Source Location

Repository: OpenRLHF
File: openrlhf/cli/train_prm.py
Lines: 1-253

Signature

def train(args) -> None:
    """
    PRM training entry point.

    Args:
        args: Namespace containing pretrain, dataset, placeholder_token,
              reward_tokens, max_epochs, learning_rate, packing_samples, etc.
    """

Import

from openrlhf.cli.train_prm import train

I/O Contract

Inputs

Name	Type	Required	Description
args.pretrain	str	Yes	HuggingFace model name or path
args.dataset	str	Yes	Path to PRM training dataset
args.placeholder_token	str	No	Token marking step boundaries for reward prediction
args.reward_tokens	List[str]	No	Tokens representing positive/negative step labels
args.packing_samples	bool	No	Pack samples using Flash Attention (default: False)
args.learning_rate	float	No	Learning rate (default: 1e-6)
args.max_epochs	int	No	Training epochs (default: 1)
args.lora_rank	int	No	LoRA rank, 0 for full fine-tuning (default: 0)

Outputs

Name	Type	Description
Trained model	files	Saved to args.save_path in HuggingFace format
Checkpoints	files	Saved to args.ckpt_path during training
Logs	WandB/TensorBoard	Training metrics (prm_loss, acc)

Usage Examples

Launch PRM Training

deepspeed --module openrlhf.cli.train_prm \
    --save_path ./ckpt/mistral-7b-prm \
    --logging_steps 1 \
    --eval_steps -1 \
    --micro_train_batch_size 1 \
    --train_batch_size 128 \
    --pretrain mistralai/Mistral-7B-v0.1 \
    --learning_rate 1e-6 \
    --max_epochs 1 \
    --dataset your/prm_dataset \
    --input_key input \
    --label_key label \
    --placeholder_token "<|placeholder|>" \
    --reward_tokens "+" "-" \
    --max_len 2048 \
    --zero_stage 2 \
    --param_dtype bf16

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment