👎 [GRPO] Adds option to disable dropout

[edbeeching]

What does this PR do?

Adds an option to disable dropout.

The RLOOTrainer disables dropout in policy, ref_model and reward model. This PR adds the option to disable dropout to the GRPOTrainer, which may improve training stability.

The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization
Provides more insight about why this option may improve training stability:

We disable the dropout layers during training, similar to the settings in Ziegler et al. (2019); Huang
et al. (2024). This is important for PPO training, especially because with dropout activated, the log
probabilities of tokens will not be reproducible, making calculating the KL penalty unreliable while
also causing the ratios of the PPO to be not 1s during the first epoch, causing PPO optimization
problems. For consistency, we also disable dropout for SFT and RM training.

[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.

[lewtun]

Nice feature - LGTM.

[lewtun]

I think AutoModelForSequenceClassification is loaded in eval model by default, so technically we don't need this here (happy to keep it though if we want to be safe)

[qgallouedec]

Indeed!

>>> from transformers import AutoModelForSequenceClassification
>>> model = AutoModelForSequenceClassification.from_pretrained("trl-lib/Qwen2-0.5B-Reward", num_labels=1)
>>> model.training
False

[qgallouedec]

In the other trainers this is set to True, maybe we should do the same here?

[qgallouedec]

we only support PreTrainedModel, which are nn.Module (what else could it be?)

Suggested change

	if args.disable_dropout:
	if isinstance(model, nn.Module):
	disable_dropout_in_model(model)
	if self.ref_model is not None and isinstance(self.ref_model, nn.Module):
	disable_dropout_in_model(self.ref_model)
	if args.disable_dropout:
	disable_dropout_in_model(model)
	if self.ref_model is not None:
	disable_dropout_in_model(self.ref_model)

[qgallouedec]

LGTM, in the future we could use a default to True (let's see if it improve stability)