👨‍👨‍👧‍👧 GRPO

[qgallouedec]

What does this PR do?

from datasets import load_dataset
from peft import LoraConfig
from trl import GRPOConfig, GRPOTrainer

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

training_args = GRPOConfig(
    output_dir="Qwen2-0.5B-GRPO",
    learning_rate=1e-5,
    logging_steps=10,
    gradient_accumulation_steps=16,
    max_completion_length=128,
)
trainer = GRPOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_model="weqweasdas/RM-Gemma-2B",
    args=training_args,
    train_dataset=dataset,
    peft_config=LoraConfig(task_type="CAUSAL_LM"),
)

trainer.train()

Fixes # (issue)

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.

[edbeeching]

typo advatages

[edbeeching]

Clean implementation, I had to check the KL term as it if different from PPO. But it appears to match the paper.

[qgallouedec]

temporary pinning to a previous version, the latest one doesn't work

[SonuDixit]

Hi Team,

1. Initially KL loss is zero(its minimum), since the ref and initial model are the same. So, the backward_gradient = 0 from this loss

2. If we update only once, the first loss term (= sum of advantage for all generations) is zero. In this case, the gradient of loss(=prob*(sum_of_advantages)) with respect to the prob is zero. Hence, all the backward gradients should be zero for SGD(without weight decay) optimisation.

From 1,2 above, the backward gradient is zero, if we update policy only once. The network may not update if we use simple optimiser like SGD.

Correction -
First objective term $l= \sum_i p_i a_i$ , where $\sum_i a_i = 0$ and $p_i = \frac{exp(x_i)}{\sum_j exp(x_j)}$ is the softmax output.
$\frac{\partial l}{\partial x_1} = p_1 (a_1- \sum_i a_i p_i) = p_1(a_1 - l)$ where I have used $\frac{\partial p_i}{\partial x_1} = p_i(\mathbb{1}{(i = 1)}-p_1)$
Now, even if the loss term is zero, one such case all $p_i$ are equal, we still have the gradient. We should be able to train with SGD.

[XiaofengZHOU]

Hi Team,

1. Initially KL loss is zero(its minimum), since the ref and initial model are the same. So, the backward_gradient = 0 from this loss
2. If we update only once, the first loss term (= sum of advantage for all generations) is zero. In this case, the gradient of loss(=prob*(sum_of_advantages)) with respect to the prob is zero. Hence, all the backward gradients should be zero for SGD(without weight decay) optimisation.

From 1,2 above, the backward gradient is zero, if we update policy only once. The network may not update if we use simple optimiser like SGD.

I have the same question

[zachluo]

@SonuDixit

1. Initially KL loss is zero(its minimum), since the ref and initial model are the same. But after that, the initial policy model will be updated (let's assume the loss is not zero) and the reference model (initialized from SFT model) will not be updated at all.
2. the first loss term is not zero as pi / pi_old = 1 or torch.exp(per_token_logps - per_token_logps.detach()) = 1. Also, advantage is not zero neither so that per_token_loss = -torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) will not become zero.

[zachluo]

The policy model only has a single update following each exploration stage.

Are we going to implement the case that the policy model has multiple updates following each exploration stage?