Confused with Reward Training Data Collator

[fahadh4ilyas]

If I see the Reward Trainer example, the token length of chosen and rejected is varied and not same because the padding only effect to batch of each chosen and rejected. So, if the longest token in chosen is 256 and the longest token in rejected is 128, the token will padded accordingly. But, computing loss seems to compute the difference of logits between chosen and rejected. How is this achieved when the token size between chosen and rejected are not same?

I know the token size will not be same because in RewardTrainer, if you did not specified data_collator, RewardDataCollatorWithPadding will be used with init parameter tokenizer=tokenizer, and max_length=max_length. The problem is padding parameter is not set to 'max_length' and that makes the length of chosen and rejected are not same.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[Eduardo-d9]

Hello, I am encountering this same issue. When i try to train my reward model using the example dataset "trl-lib/ultrafeedback_binarized" I get the following error:

loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
                                     ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (817) must match the size of tensor b (578) at non-singleton dimension 1

I too came to the conclusion that this was happening because the data collator was not padding properly. So i defined my own data_collator using DataCollatorWithPadding and defined a max_length for padding. Then added the data_collator to the RewardTrainer. This fixed the original error but gave me a new one:

lambda x: len(x["input_ids_chosen"]) <= max_length and len(x["input_ids_rejected"]) <= max_length,
                                            ^^^^^^^^^^
NameError: cannot access free variable `max_length` where it is not associated with a value in enclosing scope

I fixed this by setting a max_length directly in the RewardTrainer class. A new error appeared:

raise ValueError(
ValueError: You should supply an encoding or a list of encodings to this method that includes input_ids, but you provided ['chosen', 'rejected', 'score_chosen', 'score_rejected', 'input_ids_chosen', 'attention_mask_chosen', 'input_ids_rejected', 'attention_mask_rejected']

This seems to be caused by my custom data_collator, which looks for input_ids to pad them. However, since we are tokenizing both the chosen and rejected answers separately, different input_ids are necessary.

Here is my code:

from datasets import load_dataset
from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification
import torch
from trl import PPOConfig, PPOTrainer


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct",max_length= 8000)
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct",torch_dtype=torch.bfloat16).to(device)
train_dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="test")
#train_dataset.remove_columns(["score_chosen", "score_rejected"])

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = model.config.eos_token_id


data_collator = DataCollatorWithPadding(tokenizer, padding='max_length',max_length=8000,return_tensors="pt")

training_args = RewardConfig(
    output_dir="Granite_Reward_Model",
    logging_steps=1,
    max_length=8000,
    per_device_train_batch_size=6,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    center_rewards_coefficient=0.01,
    fp16=False,
    bf16=True,
    seed = 42)

trainer = RewardTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    processing_class=tokenizer,
    data_collator=data_collator,
)


trainer.train()

model.save_pretrained("./granite_reward_model")
tokenizer.save_pretrained("./granite_reward_model")

Environment

• trl version: 0.16.0
• transformers version: 4.46.3

Question

• How should DataCollatorWithPadding be used correctly in this context to ensure both chosen and rejected inputs are padded properly?