🐾 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

[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]

@qgallouedec the term used for the labels column in standard_stepwise_supervision and prm800K differs, prm800K uses labels while standard_stepwise_supervision uses label.

Thanks for spotting it. Generation script fixed in c7cf42f, I just pushed the fixed dataset zen/47aee34.

[qgallouedec]

@gaetanlop about the tokenize function, it would be better to

• have it as a static method
• using it with a single example (not a batch).

In fact take example from the implementation of DPOTrainer. I can fix this quickly. I offer to do a PR on your branch, wdyt?

[gaetanlop]

@qgallouedec yes of course, as you prefer, I was following the implementation done in the KTOTrainer. You can open a PR otherwise I will make the changes later today

[qgallouedec]

Should we add the separator token between the prompt and the first step? If you don't (like the current code) you get something like:

prompt = "This is my prompt."
completions = ["This is my first step.", "This is my second step."]
separator = "\n"

# Processing here

result == "This is my prompt.This is my first step. This is my second step."
#                           ^💀

[qgallouedec]

gaetanlop#1

I still need to add the collator and then it's ready. No collator is needed in fact.

[qgallouedec]

First trained model: https://huggingface.co/qgallouedec/Qwen2-0.5B-Reward

[gaetanlop]

@qgallouedec Thank you for the refactoring work on the tokenize_row function. I have made some adjustments to ensure proper handling of special tokens. Also, I refined the label creation process and updated the tokenize_row function to support truncation based on both max_length and max_completion_ids. I have added some tests to confirm that the updated tokenize_row function behaves as intended.

I also made some experiments. The model gets 99.8% accuracy after just a few steps... It might just be predicting True all the time, I will need to double check