Implementation:OpenRLHF OpenRLHF UnpairedPreferenceDataset init

Knowledge Sources	OpenRLHF KTO
Domains	Data_Processing, Alignment
Last Updated	2026-02-07 10:40 GMT

Overview

Concrete tool for constructing unpaired preference datasets for KTO training with KL divergence estimation support.

Description

The UnpairedPreferenceDataset class extends PyTorch Dataset to prepare data for Kahneman-Tversky Optimization. Unlike paired preference datasets (chosen/rejected), this class handles independently labeled samples where each example has a binary label (desirable=1, undesirable=0). The custom collate_fn is critical: it doubles the batch by appending unmatched prompt-response pairs (prompt[i] + response[i+1]) for KL divergence estimation between the policy and reference models. Supports chat templates and parallel data processing.

Usage

Use this dataset class when training with KTO, which requires unpaired preference data. Each sample should have a prompt, a response, and a binary label indicating desirability. This is the dataset class used by KTOTrainer.

Code Reference

Source Location

• Repository: OpenRLHF
• File: openrlhf/datasets/unpaired_preference_dataset.py
• Lines: 1-143

Signature

def preprocess_data(
    data,
    input_template=None,
    input_key=None,
    output_key=None,
    label_key=None,
    apply_chat_template=None,
) -> Tuple[str, str, int]: ...

class UnpairedPreferenceDataset(Dataset):
    def __init__(
        self,
        dataset,
        tokenizer: Callable,
        max_length: int,
        strategy,
        input_template=None,
        num_processors: int = 8,
    ) -> None: ...

    def process_data(self, data) -> dict: ...
    def __len__(self) -> int: ...
    def __getitem__(self, index) -> Tuple[str, str, int, int]: ...
    def collate_fn(self, item_list) -> Tuple[Tensor, Tensor, Tensor, List[int]]: ...

Import

from openrlhf.datasets.unpaired_preference_dataset import UnpairedPreferenceDataset

I/O Contract

Inputs

Name	Type	Required	Description
dataset	HF Dataset	Yes	Must have columns matching input_key, output_key, and label_key
tokenizer	Callable	Yes	HuggingFace tokenizer
max_length	int	Yes	Maximum sequence length
strategy	DeepspeedStrategy	Yes	Provides args (input_key, output_key, label_key, apply_chat_template)
input_template	str	No	Template string with {} placeholder for formatting prompts

Outputs (collate_fn)

Name	Type	Description
input_ids	Tensor	Tokenized sequences, doubled batch (matched + unmatched for KL) (2*B, seq_len)
attention_mask	Tensor	Attention masks (2*B, seq_len)
labels	LongTensor	Binary labels: 1=desirable, 0=undesirable, -1=unmatched KL pair (2*B,)
prompt_ids_lens	List[int]	Length of prompt tokens for loss masking (2*B,)

Usage Examples

Creating KTO Dataset

from openrlhf.datasets import UnpairedPreferenceDataset
from openrlhf.datasets.utils import blending_datasets

# Load raw data
train_data = blending_datasets(
    args.dataset,
    args.dataset_probs,
    strategy,
    args.seed,
    max_count=args.max_samples,
)

# Create unpaired preference dataset
# Each sample has: input (prompt), output (response), label (0 or 1)
train_dataset = UnpairedPreferenceDataset(
    train_data,
    tokenizer,
    args.max_len,
    strategy,
    input_template=args.input_template,
)

# The collate_fn doubles the batch: first half is matched pairs,
# second half is unmatched pairs for KL estimation
train_dataloader = strategy.setup_dataloader(
    train_dataset,
    args.micro_train_batch_size,
    True,
    True,
    train_dataset.collate_fn,
)

Related Pages

• Environment:OpenRLHF_OpenRLHF_CUDA_GPU_Environment