Simplify reference model handling in GRPO/RLOO

[albertvillanova]

Simplify reference model handling in GRPO/RLOO.

Follow-up to:

• Simplify reference model handling in DPO #5876

This PR refactors how reference models are handled in both GRPOTrainer and RLOOTrainer, especially to improve support for PEFT (Parameter-Efficient Fine-Tuning) adapters. The main changes clarify and centralize the logic for reference model selection and usage, ensuring consistent behavior across different training setups.

Changes

Reference model handling improvements:

• Refactored the initialization of self.ref_model in both grpo_trainer.py and rloo_trainer.py to:

  • Use a local ref_model variable during setup.
  • Assign self.ref_model at the end based on whether beta is zero, if PEFT is used, or if a reference model needs to be instantiated. This ensures that for PEFT, the base model is used as reference by disabling adapters, and for non-PEFT, an explicit reference model is created if needed.

• Updated all downstream usages to reference the new logic, including dropout disabling, model preparation (e.g., DeepSpeed/FSDP), and callbacks for syncing the reference model.

PEFT adapter support and scoring logic:

• Improved the logic in _generate_and_score_completions to explicitly check for PEFT models and correctly switch adapters when computing reference log probabilities, covering both new and retrained adapters.

These changes make the code more robust and maintainable when working with different model types and distributed training setups, and clarify how the reference model is selected and used throughout training.

---

Note

Medium Risk
Changes which module runs reference KL forwards (PEFT now consistently uses self.ref_model with adapter switching) and when a duplicate ref weights copy is loaded; incorrect assignment would skew KL-regularized RL training.

Overview
GRPO and RLOO trainers now build a separate reference checkpoint only when beta != 0 and the policy is not PEFT; setup uses a temporary ref_model for dropout, fp32 LM-head casting, and DeepSpeed/FSDP prep, then sets self.ref_model in one place: None if KL is off, else the prepared duplicate or self.model for adapter-based reference.

Reference log-probs no longer branch on self.ref_model is not None vs PEFT. They use is_peft_model(self.ref_model) to run use_adapter and forward through self.ref_model (aligned with the training module for PEFT), with a single non-PEFT path that always calls the dedicated or policy model via self.ref_model.

^{Reviewed by Cursor Bugbot for commit 0e267aa. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit e4d9e39. Configure here.}

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