[Bug] RewardDataset silently keeps preference pairs that become identical after truncation
Summary
RewardDataset independently truncates the chosen and rejected sequences to max_length, then overwrites the final retained token with EOS. If the preference-defining difference lies after the truncation boundary—or exactly at the final retained token—the two model inputs become identical. The dataset still keeps and trains on that pair.
For reward modeling and DPO, identical chosen/rejected inputs carry no preference signal: the pair loss is ln(2) and the main preference gradient is exactly zero. In mixed batches, these constant-loss samples also dilute the gradient and distort accuracy/step accounting.
Suggested severity: medium-high. The bug is reachable through both public RM and DPO CLIs, whose default --data.max_len is 512. Its likelihood depends on the response-length distribution; the shipped shell recipes explicitly use 8192, which substantially reduces it.
Tested on main at 3c3be6234e0cb353e76bb8019947db9dfe99fca7 (v0.11.0).
Root cause
The dataset's preprocessing filter only rejects prompt is None, plus an independent prompt-length check for DPO:
At retrieval time, chosen and rejected text are independently right-truncated and their final retained tokens are forcibly replaced with EOS:
There is no validation after these transformations. The DPO prompt-length filter does not help when the prompt fits but the answers share a long prefix. RM has no equivalent prompt-length filter at all.
The affected losses depend on chosen-minus-rejected values:
Identical inputs therefore make the preference difference zero as a function of the model parameters, producing loss = ln(2) and zero main preference gradient.
Stable reproduction with Qwen3.5
Run from the repository root. MODEL_PATH may be a local Qwen3.5-0.8B directory or the corresponding Hub ID.
import torch
import torch.nn.functional as F
from datasets import Dataset
from types import SimpleNamespace as NS
from transformers import AutoTokenizer
from openrlhf.datasets import RewardDataset
MODEL_PATH = "/path/to/Qwen3.5-0.8B"
tokenizer = AutoTokenizer.from_pretrained(
MODEL_PATH,
trust_remote_code=True,
local_files_only=True,
)
common = "shared token sequence " * 80
raw = Dataset.from_list(
[
{
"prompt": "Question: ",
"chosen": common + "CORRECT",
"rejected": common + "WRONG",
}
]
)
strategy = NS(
args=NS(
data=NS(
prompt_key="prompt",
chosen_key="chosen",
rejected_key="rejected",
apply_chat_template=False,
)
)
)
for is_dpo in (False, True):
dataset = RewardDataset(
raw,
tokenizer,
max_length=32,
strategy=strategy,
is_dpo=is_dpo,
num_processors=None,
)
chosen_ids, _, rejected_ids, _, _ = dataset[0]
print("DPO" if is_dpo else "RM", "identical:", torch.equal(chosen_ids, rejected_ids))
x = torch.tensor(0.0, requires_grad=True)
rm_loss = -F.logsigmoid(x - x)
rm_loss.backward()
print("RM loss/grad:", rm_loss.item(), x.grad.item())
x = torch.tensor(0.0, requires_grad=True)
dpo_loss = -F.logsigmoid(0.1 * ((x - x) - (0.0 - 0.0)))
dpo_loss.backward()
print("DPO loss/grad:", dpo_loss.item(), x.grad.item())Observed on unpatched main:
RM identical: True
DPO identical: True
RM loss/grad: 0.6931471824645996 0.0
DPO loss/grad: 0.6931471824645996 0.0The untruncated chosen and rejected sequences are different; the only distinction is beyond the retained token budget.
There is a second boundary case: if the differing token is the last retained token, the explicit EOS overwrite also removes that distinction even though truncation initially retained it.
CLI reachability and practical relevance
Both public CLIs construct RewardDataset directly, and both default to --data.max_len 512:
As a bounded relevance check—not a full prevalence estimate—I ran the real preprocessing and Qwen3.5 tokenizer over 10 deterministic pages of 100 rows from the public OpenRLHF/preference_dataset_mixture2_and_safe_pku dataset. At max_length=512, 97 of 1,000 sampled rows produced identical final RM and DPO pairs (only two were already identical before truncation). At 8192, none of those sampled rows collapsed. This is why the severity is medium-high rather than high and why the shipped 8192-token example commands should not be described as always affected.
Expected behavior
Preference pairs whose final input_ids and attention_mask are identical should be filtered before dataloader construction/sharding. Pairs whose distinction remains inside the retained token budget must be kept.
Changing global truncation direction is not necessary and could alter training semantics for valid data. Failing later in the collator or loss would also leave dataset length and distributed sharding inconsistent.
Suggested fix
Reuse the exact tokenization/EOS transformation for preprocessing and __getitem__, then set prompt=None when the final chosen/rejected input_ids and masks are equal. The existing dataset filter already removes prompt=None rows.
This adds one-time pair validation during dataset construction while leaving the model input format unchanged. It also avoids duplicating the EOS/truncation logic between validation and retrieval.
Duplicate check
I searched all open and closed issues and all PRs for combinations of RewardDataset truncation, chosen rejected identical, same input preference, zero gradient DPO truncation, and related terms. I found no report or fix for post-truncation pair collapse.
The closest items are not duplicates:
- #409 is an old preprocessing-speed request that led to parallel
dataset.map; it does not validate final pair equality. - #1025 concerns packed-sequence truncation in
unpad_and_slice_tensor, not RM/DPO preference pairs. - #909 concerns chat-template/prompt formatting.
- Open #1247 adds an optional REBEL-style loss and margin plumbing. It still independently truncates chosen/rejected inputs and contains no equality filter. It may require a straightforward rebase because it touches the same file, but it is not a semantic duplicate.
Source: OpenRLHF/OpenRLHF