Use any reward model for online methods

[qgallouedec]

What does this PR do?

This PR allows any reward model to be used with Online DPO, i.e. it removes the requirement to have the same chat template and tokenizer.

The user must now provide reward_processing_class.

trainer = OnlineDPOTrainer(
    model=model,
    reward_model=reward_model,
    args=training_args,
    train_dataset=train_dataset,
    processing_class=tokenizer,
    reward_processing_class=reward_tokenizer,  # <-
)

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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

I think we don't need strip

[qgallouedec]

Results for a gemma reward model

accelerate launch examples/scripts/dpo_online.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --reward_model_path Ray2333/GRM-Gemma-2B-rewardmodel-ft \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 5.0e-7 \
    --logging_steps 10 \
    --output_dir Qwen2-0.5B-OnlineDPO-GRM-Gemma \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 2 \
    --warmup_ratio 0.1 \
    --missing_eos_penalty 1.0 \
    --push_to_hub

https://wandb.ai/huggingface/huggingface/runs/520cnnjl

For ref, with Pair RM judge instead:

accelerate launch examples/scripts/dpo_online.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --judge pair_rm \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 5.0e-7 \
    --logging_steps 10 \
    --output_dir Qwen2-0.5B-OnlineDPO-PairRM \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 2 \
    --warmup_ratio 0.1 \
    --push_to_hub

https://wandb.ai/huggingface/huggingface/runs/ffd4u5wa

[lewtun]

Awesome PR @qgallouedec - this will unlock so many combinations of policy with models on RewardBench 🔥 !

Have you done a test run of e.g. trying to optimise Qwen2.5-0.5B-Instruct with the 7B ArmoRM model?

[lewtun]

Is the reason to use a processing class in case we want to support other modalities beyond text?

[qgallouedec]

Possibly. And since the tokenizer is now called processing_class within trainers, I'd recommend always aligning with it (even if only the textual modality is supported). Unless you have a good reason not to.

[lewtun]

Should we also set truncation=True along with truncation_side="left" to ensure the labels aren't lost on long inputs? We might also need to allow people to set max_length since the context window of the tokenizer might be different from the policy - would that be best stored in the ScriptArguments?

[qgallouedec]

Ok for truncation, and truncation side. Not sure what's the best way to let the user set the max_length. Ok for doing this in a follow-up PR?

[qgallouedec]

1834770

[lewtun]

Sounds good to follow up in separate PR - it should generally be safe for RMs that define the max length implicitly in their config anyway

[lewtun]

Why do we do this? Is it to align a prompt with chosen/rejected?

[qgallouedec]

At this point, we have:

prompts = ["What color is the sky?", "What's the capital of France?"]
completions = ["Blue", "Lyon", "Green", "Paris"]

and later, we need to concat the prompts and the completions to compute the reward

prompt_completion_ids = torch.cat((prompts_ids, completions_ids), dim=1)
_, scores, _ = get_reward(
    self.reward_model, prompt_completion_ids, self.reward_processing_class.pad_token_id, context_length
)

so we need to repeat the prompt.

[qgallouedec]

Have you done a test run of e.g. trying to optimise Qwen2.5-0.5B-Instruct with the 7B ArmoRM model?

ArmoRM is a custom classifier (its code for using it is not standard). So our get_reward function probably won't work for it. However, by modifying the code a little, I still manage to use it, and this is what I get:

https://wandb.ai/huggingface/huggingface/runs/merlfqgx (screenshot to come)

accelerate launch examples/scripts/dpo_online.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --reward_model_path RLHFlow/ArmoRM-Llama3-8B-v0.1 \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 5.0e-7 \
    --logging_steps 10 \
    --output_dir Qwen2-0.5B-OnlineDPO-AutoRM \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 2 \
    --warmup_ratio 0.1 \
    --missing_eos_penalty 1.0 \
    --push_to_hub