Process-supervised RM Trainer

[gaetanlop]

What does this PR do?

Adding support for process-supervised reward training to TRL as requested in #2110 .

List of papers using PRMs: [1], [2], [3], [4]...

Fixes # (issue)

#2110

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• 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.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?

Who can review?

@lewtun @kashif

[lewtun]

This is awesome @gaetanlop ! Would you like some early feedback on the PR or would you prefer I wait a bit until it's more polished?

[gaetanlop]

Hey @lewtun, thank you for the message. Currently, the only files that are more or less ready are prm_trainer.py and prm_config.py. The rest are just placeholders that I haven’t had the opportunity to work on yet.

Implementing a PRMs seems to be pretty straighforward, it seems to be a token classification task where only prediction for the last token of each step gets assigned a label and other tokens are ignored during loss calculation.

If the dataset isn’t pre-tokenized, I assume it should contain the following columns:

• prompt: Either a string or past messages
• steps: A list of strings
• labels: A list of integers corresponding to the label associated to each step

Are you aware of an HF dataset to train PRMs for the example file? Also, how can I add a new subset to the trl-internal-testing/zen dataset to support stepwise reward models for the unit test of the prm_trainer?

Thanks again for your time!

[gaetanlop]

PR ready for review. I have changed the naming conventions that I used before prm to the suggested naming in #2110 stepwise.

Tests: I created a dummy_dataset but we should add a subset to trl-internal-testing/zen as done in other scripts.
Example: The example is currently using a placeholder for the dataset name as to the best of my knowledge trl didn't release a dataset for stepwise reasoning on HF. We should add this too.

[lewtun]

Thank you for the very clean PR @gaetanlop - this looks great! I've left some minor suggestions regarding the structure, but aside from that and having a smallish dataset in the right format we can sanity check that the accuracy goes up, loss goes down etc I think this is quite close to being ready

[lewtun]

What do you think about picking a subset from PRM800k to test everything works?

You could create a subset in the expected format and then we can merge it with trl-lib/zen :)

[gaetanlop]

I made two pull requests to trl-lib/zen (https://huggingface.co/datasets/trl-lib/zen/discussions/3) to add the subsets to trl-lib.

[qgallouedec]

#2127

[gaetanlop]

Thanks for looking at this @lewtun. Seems like trl-internal-testing/zen is the dataset you are using for testing. I have done a PR to trl-lib/zen, should I also PR trl-internal-testing/zen to add 19 samples of PRM800K for testing or are you handling it on your side (it looks like they are both the same dataset)?

[qgallouedec]

Thanks a lot @gaetanlop! I left some comments.
I think we need to first define a new standard for the dataset required here. I prefer working on it an another PR: #2148.
Once we've merged #2148 we'll need to update this one.

Feel free to share your thoughts in #2148 btw.

[gaetanlop]

Thanks for the reviews @qgallouedec. I will wait that #2148 is completed and merged before resolving your comments. Don't hesitate to ping me here when you have done your decisions regarding the new standard dataset format.