ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove grpo requirement bs=num_generations (#3609)

Browse source 
* 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>
This commit is contained in:
DoubleMathew 2025-11-19 21:57:01 -06:00 • committed by GitHub
commit ecacc17f7e
2 changed files with 25 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

25
unsloth/models/rl.py
View file

@ -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"

1
unsloth/models/rl_replacements.py
View file

@ -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 ""