🏎️ vllm for Online DPO

[qgallouedec]

What does this PR do?

Use vLLM for generation. 2.2x faster 🚀🚀

Demo:

from datasets import load_dataset
from trl import OnlineDPOConfig, OnlineDPOTrainer, PairRMJudge
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
judge = PairRMJudge()
train_dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

training_args = OnlineDPOConfig(output_dir="Qwen2-0.5B-OnlineDPO-vllm", logging_steps=10, use_vllm=True, gradient_accumulation_steps=8)
trainer = OnlineDPOTrainer(
    model=model, judge=judge, args=training_args, processing_class=tokenizer, train_dataset=train_dataset
)

trainer.train()

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]

we get different results with vllm. probably linked to sampling param. investigating

[qgallouedec]

CI fails because in the latest transformers version release yesterday, transformers uses a python 3.10+ syntax (timeout: float | None = None). I'm not sure why it fails only for the cli test, but I think we can safely ignore it.

[qgallouedec]

To avoid OOM for long prompts

[qgallouedec]

dtype and gpu_memory_utilization are hardcoded but we can still make them as arg in the future.
See here why using torch.float32 is important.

[qgallouedec]

Surprisingly, the precision of the generator model seems to have a pretty high impact on the results:

When you keep the default precision (bfloat16), the results seem to be significantly worst. It may be related to higher noise in generation? Making the temperature vary might help to confirm this intuition. In the meantime, I've hard-coded float32 as model precision. It hurts a bit the speed (you've less space for KV cache) but it's still way faster than without vLLM.

2.2x faster 🚀🚀

[edbeeching]

Awesome addition. If I understand correctly there are now two instances of the model, so this will only work with DDP and small models where memory capacity is not an issue?

[qgallouedec]

there are now two instances of the model

Correct: one trained model and one generation model. Note that the weights of these two models are always the same. I wonder if we could share them instead of duplicating them? @hmellor any idea?

small models where memory capacity is not an issue

Yes, parallelism (not supported yet) should make it possible to relax this constraint.

[hmellor]

vllm-project/vllm#10353 is probably the closest we can currently get to directly accessing vLLM's model, but this PR is expected to be superseded by work done in the V1 engine re-architecture that is ongoing (planned beta-release any day now, on by default by end of January).

vllm-project/vllm#12084 introduces some RLHF features to vLLM, but it follows the OpenRLHF model where training and inference processes live on different GPUs.