[core] Officially Support Reward Modeling

[younesbelkada]

What does this PR do?

With Reward modeling being an important piece of PPO algorithm, it would be cool to support an "official" RewardTrainer in trl.

The RewardTrainer simply inherits from transformers.Trainer, but with some constraints. Users should be responsible to create a paired dataset that contains input_ids_j, input_ids_k, attention_mask_j, attention_mask_k, if they want to use the default RewardDataCollatorWithPadding data collator.
Also I propose to add the possibility to create the PEFT model under the hood, if a user passes a PeftConfig to the Trainer.

This PR adds a first version of it, adds also nice tests and cool documentation about that

TODO: update the README & reward_trainer.mdx file

cc @lvwerra

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[lvwerra]

Generally looks really good and clean to me. Left a few comments to try to make it a bit more user friendly.

[lvwerra]

Maybe @lewtun would also be interested to have a look to see if there is feedback from the H4 team.

[lewtun]

Awesome feature @younesbelkada 🔥

I left some tiny nits and a feature request to compute accuracy by default :)

[lewtun]

Maybe explain what format the raw dataset should have here? E.g. you could use samples of the StackExchange or Anthropic dataset (https://huggingface.co/datasets/Anthropic/hh-rlhf) as a guide

[younesbelkada]

I added few lines in 4bcd96e but not sure if what I said is 100% correct, would love to have a second look here!

[younesbelkada]

Will also add an example now - EDIT added it

[lewtun]

Since accuracy is the most common metric for evaluating reward models, would it make sense to provide it as a default in compute_metrics? E.g. something like this should work:

    def compute_metrics(eval_pred):
        predictions, _ = eval_pred
        # Here, predictions is rewards_chosen and rewards_rejected.
        # We want to see how much of the time rewards_chosen > rewards_rejected.
        predictions = np.argmax(predictions, axis=0)
        labels = np.zeros(predictions.shape)
        return accuracy.compute(predictions=predictions, references=labels)

[younesbelkada]

This would add evaluate as an additional dependency to the library, we can also have it as an optional dependency similar as peft! For me it's totally fine to have it as a core dependency, but l want to hear @lvwerra's opinion to make sure we are aligned on this

[lewtun]

Ah true, I think having it as an optional dep would be the way to go (unless evaluate is already so light that it's deps are covered by the trl core deps)

[lvwerra]

Accuracy is not a very hard metric, maybe we can just build it from scratch here :)

[lvwerra]

Just a few small comments, otherwise this is good to go!

[lvwerra]

Default is accuracy.

[lvwerra]

i am not sure this works: if we overwrite the class method with our own metric the compute metrics in the parent class is never used, no?

what about defining a compute_accuracy function outside the class and pass if compute_metrics from the init is None

[younesbelkada]

Sounds like a great plan!