Implementation:OpenRLHF OpenRLHF ProcessRewardModelTrainer

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

Overview

Concrete tool for training a process reward model (PRM) that predicts step-level rewards for reasoning chains.

Description

The ProcessRewardModelTrainer class implements the training loop for process reward models. It uses PRMLoss to compute cross-entropy loss at placeholder token positions where step-level labels are assigned. The trainer supports both hard labels (token-based +/-) and soft labels (float reward values), MoE auxiliary loss for Mixtral-style models, packing samples via Flash Attention, gradient accumulation, distributed training with DeepSpeed, and WandB/TensorBoard logging. The evaluate method tracks both loss and accuracy metrics.

Usage

Use this trainer for training process reward models that evaluate reasoning quality at each intermediate step. This enables more granular reward signals compared to outcome reward models that only evaluate the final answer.

Code Reference

Source Location

Repository: OpenRLHF
File: openrlhf/trainer/prm_trainer.py
Lines: 1-260

Signature

class ProcessRewardModelTrainer(ABC):
    def __init__(
        self,
        model,
        strategy,
        optim: Optimizer,
        train_dataloader,
        eval_dataloader,
        scheduler,
        max_norm: float = 1,
        batch_size: int = 1,
        max_epochs: int = 2,
        tokenizer=None,
        save_hf_ckpt: bool = False,
        disable_ds_ckpt: bool = False,
    ) -> None: ...

    def fit(self, args, consumed_samples=0, num_update_steps_per_epoch=None): ...
    def evaluate(self, eval_dataloader, steps=0): ...
    def save_logs_and_checkpoints(self, args, global_step, step_bar, logs_dict={}, client_states={}): ...

Import

from openrlhf.trainer.prm_trainer import ProcessRewardModelTrainer

I/O Contract

Inputs

Name	Type	Required	Description
model	nn.Module	Yes	Actor model used as backbone for PRM
strategy	DeepspeedStrategy	Yes	Distributed training strategy
optim	Optimizer	Yes	PyTorch optimizer
train_dataloader	DataLoader	Yes	Training data (ProcessRewardDataset)
eval_dataloader	DataLoader	No	Evaluation data
scheduler	LRScheduler	Yes	Learning rate scheduler
tokenizer	PreTrainedTokenizer	No	Tokenizer (for checkpoint saving)
max_norm	float	No	Gradient clipping max norm (default: 1.0)
batch_size	int	No	Global batch size (default: 1)

Outputs

Name	Type	Description
fit()	None	Trains model in-place, saves checkpoints and logs metrics
evaluate()	None	Computes eval loss and accuracy; logs to WandB/TensorBoard

Usage Examples

Creating and Running PRM Training

from openrlhf.trainer.prm_trainer import ProcessRewardModelTrainer

trainer = ProcessRewardModelTrainer(
    model=model,
    strategy=strategy,
    optim=optimizer,
    train_dataloader=train_dataloader,
    eval_dataloader=eval_dataloader,
    scheduler=scheduler,
    max_norm=1.0,
    batch_size=args.train_batch_size,
    max_epochs=args.max_epochs,
    tokenizer=tokenizer,
    save_hf_ckpt=args.save_hf_ckpt,
    disable_ds_ckpt=args.disable_ds_ckpt,
)

# Run training
trainer.fit(args, consumed_samples=0, num_update_steps_per_epoch=num_steps)

Related Pages

Page Connections

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

Principle

Implementation

Heuristic

Environment