diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index b385dba2eb..cfd9ad8227 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -481,6 +481,28 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): extra_args += num_proc_check pass + # Check for loss_type = dr_grpo and scale_rewards for GRPO + if "loss_type" in call_args and "scale_rewards" in call_args: + check_dr_grpo = \ + "if loss_type.lower() == 'dr_grpo':\n"\ + " loss_type = 'dr_grpo'\n"\ + "elif loss_type.lower() == 'dapo':\n"\ + " loss_type = 'dapo'\n"\ + "if loss_type.lower() == 'dr_grpo':\n"\ + " if scale_rewards == None:\n"\ + " scale_rewards = True\n"\ + " elif scale_rewards == True:\n"\ + " print('The Dr GRPO paper recommends setting `scale_rewards` to False! Will override. Set it to `None` to force False.')\n"\ + " scale_rewards = False\n" + "elif loss_type.lower() = 'dapo'\n"\ + " print('The DAPO paper recommends `mask_truncated_completions = True`')\n"\ + " print('The DAPO paper recommends `epsilon_high = 0.28`')\n"\ + " mask_truncated_completions = True\n"\ + " epsilon_high = 0.28\n"\ + "\n" + extra_args += check_dr_grpo + pass + # Edit config with anything extra if trainer_file in RL_CONFIG_CHANGES: process_extra_args = RL_CONFIG_CHANGES[trainer_file]