From fa91cc466fdabb7fee01c5d2e175bf12da834aa0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 14 Feb 2025 15:32:02 -0800
Subject: [PATCH] Fix GRPO bsz

---
 unsloth/models/rl.py              | 16 +++++++++++++++-
 unsloth/models/rl_replacements.py | 24 +++++++++++++++++++++---
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 58b6d8271b..eba1e46a21 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -29,6 +29,7 @@ from .rl_replacements import (
     RL_EXTRA_ARGS,
     RL_FUNCTIONS,
     RL_PRE_ITEMS,
+    RL_CONFIG_CHANGES,
 )
 selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"]
 
@@ -165,8 +166,14 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     if RLTrainer.__name__.startswith("Unsloth"): return
     if RLConfig .__name__.startswith("Unsloth"): return
 
+    # Get old source
+    old_RLTrainer_source = inspect.getsource(RLTrainer)
+    old_RLConfig_source  = inspect.getsource(RLConfig)
+
     all_imports = dir(trainer)
-    imports = [x for x in all_imports if not x.startswith("_")]
+    # imports = [x for x in all_imports if not x.startswith("_")]
+    # Fix _deprecate_arguments not getting imported
+    imports = all_imports
 
     # Get default arguments
     EMPTY = inspect.Parameter.empty
@@ -381,6 +388,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         extra_args += num_proc_check
     pass
 
+    # Edit config with anything extra
+    if trainer_file in RL_CONFIG_CHANGES:
+        process_extra_args = RL_CONFIG_CHANGES[trainer_file]
+        for process_extra_arg in process_extra_args:
+            extra_args += process_extra_arg(old_RLTrainer_source, old_RLConfig_source)
+    pass
+
     # Edit report_to and default it to nothing if max_steps is like 60
 
     # Create RLConfig args
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index d01f6cd45f..fefba2444c 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -16,6 +16,7 @@ __all__ = [
     "RL_EXTRA_ARGS",
     "RL_FUNCTIONS",
     "RL_PRE_ITEMS",
+    "RL_CONFIG_CHANGES",
 ]
 
 import re
@@ -23,9 +24,10 @@ import torch
 import inspect
 from collections import defaultdict
 from unsloth_zoo.rl_replacements import RL_REPLACEMENTS
-RL_EXTRA_ARGS = defaultdict(list)
-RL_FUNCTIONS  = defaultdict(list)
-RL_PRE_ITEMS  = defaultdict(list)
+RL_EXTRA_ARGS     = defaultdict(list)
+RL_FUNCTIONS      = defaultdict(list)
+RL_PRE_ITEMS      = defaultdict(list)
+RL_CONFIG_CHANGES = defaultdict(list)
 
 torch_compile_options = {
     "epilogue_fusion"   : True,
@@ -242,3 +244,19 @@ def grpo_trainer_compute_loss(function_name, function):
     return function
 pass
 RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss)
+
+# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356
+# TRL warns if batch size is not a multiple of num_generations -> fix this.
+def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source):
+    if "multiple of num_generations" not in RLTrainer_source: return ""
+    if "num_generations" not in RLConfig_source: return ""
+
+    check_batch_size = \
+    "div = per_device_train_batch_size // num_generations\n"\
+    "if div * num_generations != per_device_train_batch_size:\n"\
+    "    print('Unsloth: We know expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n'\\"\
+    "          'We will change the batch size of ' + per_device_train_batch_size + ' to the `num_generations` of ' + num_generations')\n"\
+    "    per_device_train_batch_size = num_generations\n"
+    return check_batch_size
+pass
+RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size)