🐦‍🔥 6x faster GRPO with multi-step optimization

[qgallouedec]

What does this PR do?

This PR implements the multi-step trick. It allows the generation to be reused several times, speeding up training.

It requires to implement the importance sampling and the clipping logic.

- per_token_loss = -torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
+ coef_1 = torch.exp(per_token_logps - old_per_token_logps)
+ coef_2 = torch.clamp(coef_1, 1 - self.epsilon, 1 + self.epsilon)
+ per_token_loss1 = coef_1 * advantages.unsqueeze(1)
+ per_token_loss2 = coef_2 * advantages.unsqueeze(1)
+ per_token_loss = -torch.min(per_token_loss1, per_token_loss2)

Results with various mu (note that mu=1 matches the current implementation in main):

For reproduction:

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer


dataset = load_dataset("trl-lib/tldr", split="train")

# Dummy reward function: the closer the completion is to 20 characters, the higher the reward
def reward_len(completions, **kwargs):
    return [-abs(100 - len(completion)) for completion in completions]

def main():
    num_iterations=4
    training_args = GRPOConfig(
        output_dir=f"Qwen2.5-0.5B-GRPO-2899-μ={num_iterations}",
        logging_steps=5,
        gradient_accumulation_steps=4,
        per_device_train_batch_size=4,
        num_generations=8,
        max_prompt_length=64,
        max_completion_length=64,
        log_completions=True,
        max_steps=200,
        num_iterations=num_iterations,
    )
    trainer = GRPOTrainer(
        model="Qwen/Qwen2.5-0.5B",
        reward_funcs=reward_len,
        args=training_args,
        train_dataset=dataset,
    )
    trainer.train()


if __name__ == "__main__":
    main()

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]

Should we set another default? eg 4?

[edbeeching]

if num_iterations == 1, I think we do a redundant forward pass here as the same thing will be recalculated at the optimization step. Would it be better to only calculate this if num_iterations > 1, and then at the optimization step ?
For example:

if num_iterations==1:
    old_per_token_logps = per_token_logps.detach()
else:
    old_per_token_logps = inputs["old_per_token_logps"]

[qgallouedec]

Indeed, that's why this version is slightly slower with num_iterations=1:

Let me add this optimization

[qgallouedec]

7a6c086

[kashif]

clearer when one realizes that the integers are the indices of the prompts!

[qgallouedec]

clearer when one realizes that the integers are the indices of the prompts!

I'll add it, thanks for the feedback