🩺 Dr. GRPO loss

[qgallouedec]

What does this PR do?

This PR supersedes #3231 #3138
Closes #3178

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?
• 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]

Same effective batch size (256)

GPUs	Grad accum steps	Per device batch size
1	2	128
1	4	64
1	8	32
1	16	16
1	32	8
2	1	128
2	2	64
2	4	32
2	8	16
2	16	8
4	1	64
4	2	32
4	4	16
4	8	8
8	1	32
8	2	16
8	4	8

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig

dataset = load_dataset("trl-lib/tldr", split="train[:500]")

# Dummy reward function: count the number of unique characters in the completions
def reward_num_unique_chars(completions, **kwargs):
    return [len(set(c)) for c in completions]

ga = 2
bs = 16

args = GRPOConfig(
    output_dir=f"DrGRPO_bs{bs}_ga_{ga}_4GPU",
    per_device_train_batch_size=bs,
    gradient_accumulation_steps=ga,
    num_train_epochs=1,
    logging_steps=1,
    max_prompt_length=64,
    max_completion_length=64,
    loss_type="drgrpo",
)

trainer = GRPOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    args=args,
    reward_funcs=reward_num_unique_chars,
    train_dataset=dataset,
)
trainer.train()

[lewtun]

LGTM with some nits and a question about what BNPO refers to

[lewtun]

What is bnpo? Would be good to have a reference to where it's defined (I thought we had DAPO as the default loss)

[qgallouedec]

In fact, I realized while doing this PR that it wasn't exactly DAPO that was being used, but a variant of BNPO as defined here :

Let me try to clarify here. Losses per token are normalized by

• GRPO: the length of the sequence
• DAPO: the average sequence length in the group
• BNPO: the average sequence length in the batch
• TRL's BNPO: the average sequence length in the local batch*; this is what I call bnpo in the code, but it's not 100% correct
• Dr GRPO: by the maximum possible length of the completion

*a batch is made up of num_devices * gradient_accumulations local batches

[qgallouedec]

Special cases:
When

• per_device_batch_size==num_generations, TRL's BNPO is equivalent to DAPO
• per_device_batch_size==1, TRL's BNPO is equivalent to GRPO
• gradient_accumualtion_steps==1 and num_devices=1, TRL's BNPO is equivalent to the actual BNPO.

[lkevinzc]

@qgallouedec Thanks for the comprehensive support! A minor comment for your future consideration: Dr. GRPO does not constrain the constant normalizer to be MAX_LEN (although it's easier to just use that). This can affect the update scale (related to your recent tweet https://x.com/QGallouedec/status/1908741708021457357). In fact, different constant of x in the setting in your tweet can be absorbed into the constant normalizer we propose in the paper, and MAX_LEN is a convenient example.

[lewtun]

If I understand correctly, @edbeeching was trying something slightly different in #3231 that did local scaling per batch instead of a global constant. Do you know if there's much difference between the two?

[edbeeching]

They are roughly equivalent, I have closed my PR in favor of this one.

[qgallouedec]

Results mostly match, expect for the the loss and the grad norm, Dr GRPO seems to reduce the range, as expected.