From ecacc17f7e598549c2a8059635956442e4489cbb Mon Sep 17 00:00:00 2001
From: DoubleMathew <mmathew23@gmail.com>
Date: Wed, 19 Nov 2025 21:57:01 -0600
Subject: [PATCH] Remove grpo requirement bs=num_generations (#3609)

* Remove grpo requirement bs=num_generations

* Update rl.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 unsloth/models/rl.py              | 25 ++++++++++++++++++++++++-
 unsloth/models/rl_replacements.py |  1 +
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 3618e55825..bfbd0bf103 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -822,7 +822,30 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         extra_args += check_dr_grpo
 
     # Check GRPO num_generations mismatch
-    if "per_device_train_batch_size" in call_args and "num_generations" in call_args:
+    if (
+        "per_device_train_batch_size" in call_args
+        and "num_generations" in call_args
+        and "steps_per_generation" in call_args
+        and "generation_batch_size" in call_args
+    ):
+        # if world size is not set by accelerate or torchrun at this point it will be 1
+        check_num_generations = (
+            "if steps_per_generation is None and generation_batch_size is None:\n"
+            "    ga = gradient_accumulation_steps\n"
+            "    world_size = int(os.environ.get('WORLD_SIZE', '1'))\n"
+            "    if (ga * world_size * per_device_train_batch_size) % num_generations != 0:\n"
+            "        print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\\n"
+            "We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations))\n"
+            "        per_device_train_batch_size = num_generations\n"
+            "\n"
+        )
+        extra_args += check_num_generations
+    elif "per_device_train_batch_size" in call_args and "num_generations" in call_args:
+        if "steps_per_generation" not in call_args:
+            print(f"Unsloth: Could not find `steps_per_generation` in {trainer_file}")
+        if "generation_batch_size" not in call_args:
+            print(f"Unsloth: Could not find `generation_batch_size` in {trainer_file}")
+
         check_num_generations = (
             "if (per_device_train_batch_size // num_generations) * num_generations != per_device_train_batch_size:\n"
             "    print('Unsloth: We now expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n"
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index d1ade28bb5..8746399afa 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -853,6 +853,7 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss)
 # 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 "divisible by the number of generations" not in RLTrainer_source:
+        # in later trl versions this doesn't exist anymore
         return ""
     if "num_generations" not in RLConfig_source:
         return ""