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

qgallouedec · 2025-02-19T00:14:12Z

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 · 2025-02-20T12:41:52Z

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 · 2025-02-20T12:47:40Z

trl/trainer/grpo_config.py

@@ -224,6 +228,14 @@ class GRPOConfig(TrainingArguments):
            "training speed."
        },
    )
+    num_iterations: int = field(
+        default=1,


Should we set another default? eg 4?

edbeeching · 2025-02-20T12:54:47Z

trl/trainer/grpo_trainer.py

@@ -638,6 +731,10 @@ def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[s
        logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens

        with torch.inference_mode():
+            old_per_token_logps = self._get_per_token_logps(


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

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

Let me add this optimization

…into multi-step-grpi

kashif · 2025-02-20T13:55:22Z

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

qgallouedec · 2025-02-20T14:00:57Z

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

I'll add it, thanks for the feedback

alexxchen · 2025-02-27T05:53:01Z

@qgallouedec I think we should consider adding another term to calculate $$KL(\pi_{\theta}, \pi_{\theta_{old}})$$ , because this is the KL term that effectively constrains offline updates, whereas, $$KL(\pi_{\theta}, \pi_{ref})$$ is only used to limit the overall model update. Alternatively, should we set sync_ref_model=True by default and set ref_model_sync_steps=1

qgallouedec · 2025-02-27T08:37:50Z

Not sure to get why you want to sync the ref model every step. It is equivalent to have no ref model, the KL will be always 0

alexxchen · 2025-02-27T08:54:45Z

Not sure to get why you want to sync the ref model every step. It is equivalent to have no ref model, the KL will be always 0

Sorry, my mistake. I meant that setting ref_model_sync_steps = num_iterations might help stabilize the multi-step update. However, this would result in the loss of regularization provided by the original reference model.

* Add num_updates and epsilon parameters to GRPOConfig and GRPOTrainer * test sampler * update the loss computation * fix eval sampler * should work now * buffer inputs with grad accum * optimize when num_iterations == 1 * test * minor comment removal and fix log metric * beta position * clarify comment [ci skip] * clarify sampler doc [ci skip] * fix collision with eval logging * clarify

Alex-Songs · 2025-03-13T07:45:14Z

@qgallouedec Can num_iterations not be used together with num_train_epochs? I found that setting num_train_epochs=2, the larger the num_iterations, the more iterations there are and the slower the training.

SonuDixit · 2025-03-31T13:22:34Z

Hi Team,
Here, num_iterations($\mu$) optimisation steps are made on the same completion(trajectory). Please see the attached image from PPO paper -

Attaching algorithm from DeepseekMath paper -

I think, for line 11 in the above Algorithm - a different batch size $B$ is used such that $\mu * B \approx |D_b|$. A total of $\mu$ updates are being done using the batch size $B$.

Please see the PPO algorithm -

I think $k$ in the PPO-Algo corresponds to $\mu$ in the GRPO. PPO Authors have mentioned hyperparams in appendix A, wherein it can be seen that K*M is less than or approximately equal to NT.

qgallouedec · 2025-03-31T14:52:29Z

Thanks @SonuDixit can you open a dedicated issue for this? Just copy-paste you comment

SonuDixit · 2025-03-31T16:52:50Z

Thanks @SonuDixit can you open a dedicated issue for this? Just copy-paste you comment

#3197

* Add num_updates and epsilon parameters to GRPOConfig and GRPOTrainer * test sampler * update the loss computation * fix eval sampler * should work now * buffer inputs with grad accum * optimize when num_iterations == 1 * test * minor comment removal and fix log metric * beta position * clarify comment [ci skip] * clarify sampler doc [ci skip] * fix collision with eval logging * clarify

Add num_updates and epsilon parameters to GRPOConfig and GRPOTrainer

c076794

qgallouedec mentioned this pull request Feb 19, 2025

Is mini-batch updates implemented in GRPO of trl ? #2903

Closed

qgallouedec and others added 5 commits February 19, 2025 11:23

test sampler

ee19148

update the loss computation

0bf4f23

fix eval sampler

9ce9b98

should work now

a178fd9

buffer inputs with grad accum

946d1f6

qgallouedec changed the title ~~Multi-step GRPOTrainer~~ 🐦‍🔥 6x faster GRPO with multi-step optimization Feb 20, 2025

Merge branch 'main' into multi-step-grpi

035e4be

qgallouedec commented Feb 20, 2025

View reviewed changes

qgallouedec requested review from kashif, edbeeching, lewtun and plaguss February 20, 2025 12:53

qgallouedec marked this pull request as ready for review February 20, 2025 12:53

edbeeching reviewed Feb 20, 2025

View reviewed changes

qgallouedec mentioned this pull request Feb 20, 2025

Why does the loss start at 0 when I train GRPO, and then possibly increase? huggingface/open-r1#239

Closed

qgallouedec added 7 commits February 20, 2025 13:07

optimize when num_iterations == 1

7a6c086

Merge branch 'multi-step-grpi' of https://github.com/huggingface/trl …

70eaf11

…into multi-step-grpi

test

7ab20eb

minor comment removal and fix log metric

088c745

beta position

a8cc776

clarify comment [ci skip]

11d90c8

clarify sampler doc [ci skip]

8fe6264

kashif approved these changes Feb 20, 2025

View reviewed changes

idanshen added a commit to idanshen/multi_ref that referenced this pull request Feb 20, 2025

fix bug in grpo, based on huggingface/trl#2899

53c9484

fix collision with eval logging

7e2b7ce

clarify

3f2219d

qgallouedec merged commit e5ae703 into main Feb 20, 2025
14 checks passed

qgallouedec deleted the multi-step-grpi branch February 20, 2025 18:51

joey00072 mentioned this pull request Feb 21, 2025

feat: Add cliprange to GRPO loss #2739

Closed

5 tasks

jamesbraza mentioned this pull request Feb 27, 2025

[GRPO] Gradient mode issue when num_iterations > 1 #2953

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

Uh oh!

qgallouedec commented Feb 19, 2025 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Feb 20, 2025

Uh oh!

qgallouedec Feb 20, 2025

Uh oh!

edbeeching Feb 20, 2025 •

edited

Loading

Uh oh!

qgallouedec Feb 20, 2025

Uh oh!

qgallouedec Feb 20, 2025

Uh oh!

kashif commented Feb 20, 2025 •

edited

Loading

Uh oh!

qgallouedec commented Feb 20, 2025

Uh oh!

Uh oh!

alexxchen commented Feb 27, 2025

Uh oh!

qgallouedec commented Feb 27, 2025

Uh oh!

alexxchen commented Feb 27, 2025

Uh oh!

Alex-Songs commented Mar 13, 2025

Uh oh!

SonuDixit commented Mar 31, 2025

Uh oh!

qgallouedec commented Mar 31, 2025

Uh oh!

SonuDixit commented Mar 31, 2025

Uh oh!

Uh oh!

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

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

Uh oh!

Conversation

qgallouedec commented Feb 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Feb 20, 2025

Uh oh!

qgallouedec Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

edbeeching Feb 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

qgallouedec Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

qgallouedec Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

kashif commented Feb 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

qgallouedec commented Feb 20, 2025

Uh oh!

Uh oh!

alexxchen commented Feb 27, 2025

Uh oh!

qgallouedec commented Feb 27, 2025

Uh oh!

alexxchen commented Feb 27, 2025

Uh oh!

Alex-Songs commented Mar 13, 2025

Uh oh!

SonuDixit commented Mar 31, 2025

Uh oh!

qgallouedec commented Mar 31, 2025

Uh oh!

SonuDixit commented Mar 31, 2025

Uh oh!

Uh oh!

qgallouedec commented Feb 19, 2025 •

edited

Loading

edbeeching Feb 20, 2025 •

edited

Loading

kashif commented Feb 20, 2025 •

edited

Loading