Check rewards shapes in RewardTrainer

[ioverho]

What does this PR do?

While trying to fine-tune a reward model from a SFT trained model, I kept running into the following error:

  File ".../train_rm.py", line 459, in train_rm
    trainer.train()
  File ".../.venv/lib/python3.11/site-packages/transformers/trainer.py", line 2240, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/transformers/trainer.py", line 3745, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/trl/trainer/reward_trainer.py", line 281, in compute_loss
    loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
                                     ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (84) must match the size of tensor b (146) at non-singleton dimension 1

I interpreted this to mean that the sequence lengths of the chosen and rejected chats were incompatible. The issue persisted even with a batch size of 1. I initially tried fixing this with a custom RewardDataCollatorWithPadding, but this raised more problems.

This PR fixes this issue, by padding the chosen and rejected chats to the same sequence length. Specifically, it finds the longest sequence in either batch, and pads all other sequences to that length. This PR also allows for the user to pass custom data collators using the data_collator argument in the Trainer's __init__.

This happens because the model used still output a score per token, and not for the entire sequence. This can happen when the user passes a model with an LM head or token classification head, instead of one with a sequence classification head. Simply put, the output of a reward model should be a scalar, not a sequence. However, since the code is otherwise dimension-agnostic, it only fails here.

This PR tries to prevent this issue by adding a warning when passing the first output of the model has an unexpected shape (i.e., more than 2 dimensions), and raises a descriptive warning when the model output is completely incompatible (i.e., the rewards have different sequence lengths).

This PR also allows for using custom data collators in the RewardTrainer.

Specifics

trainer.reward_config.RewardConfig

1. Added a pad_to_multiple_of parameter. This was already implemented in the RewardDataCollatorWithPadding, but never actually provided when defining the default data collator. This also matches the behavior of the trainer.sft_config.SFTConfig config class.

trainer.reward_trainer.RewardTrainer

1. Replaced the various references to max_length with references to the args. Currently, the max_length variable is only defined when no data collator is provided (L178), making it impossible to pass a custom data collator despite the function and docs suggesting this is possible
2. Passed the added pad_to_multiple_of parameter to the instantiation of the default data collator.
3. On the first training step, raises a warning if the shapes of the chosen and rejected rewards have different lengths. The output should be 2 dimensional. If this differs, it might indicate the model output does not reflect sequence level classification.
4. Added a check for compatible shapes before computing the loss, for example because the rewards have different sequence lengths. The error suggests checking the model configuration.

Fixes

#3101
Custom DataCollator Bug in RewardTrainer

#3166
The size of tensor a (882) must match the size of tensor b (568) at non-singleton dimension 1

#686
Confused with Reward Training Data Collator

#376
The issue of reward_model output length

Submission Checklist

• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
  • I did not open an issue, but this has been discussed before
• Did you make sure to update the documentation with your changes?
  • Added some minimal documentation to RewardDataCollatorWithPadding, should be covered by autodoc
• Did you write any new necessary tests?
  • Tests added under test_reward_trainer.py

This is my first commit to trl or another HuggingFace library. I've tried to keep the changes as minimal as possible. Please let me know if anything else is required of me.

Updates

• Fixed a bug where only the length of the chosen sequence was checked
• Added some tests
• Realized the mistake was mine, my model was still using its LM head, not the reward regression head (mea culpa)