Remove grpo requirement bs=num_generations #3609
Conversation
Summary of ChangesHello @mmathew23, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refactors the batch size handling within the GRPO trainer to remove a previously enforced strict equality between Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to relax the requirement that per_device_train_batch_size must be a multiple of num_generations for GRPO training. The new logic correctly introduces a more flexible check, which only enforces a constraint on the total effective batch size when steps_per_generation and generation_batch_size are not specified. My review identified a key issue: the old, stricter validation logic remains in the codebase and is executed after the new check. This effectively negates the changes in this pull request, as the old constraint is still enforced. I've left a comment detailing this issue and how to resolve it.
| 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 (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" | ||
| "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" | ||
| " per_device_train_batch_size = num_generations\n" | ||
| "\n" | ||
| ) | ||
| extra_args += check_num_generations |
There was a problem hiding this comment.
While this new relaxed check is a good improvement, it appears the old, stricter check from grpo_trainer_fix_batch_size in unsloth/models/rl_replacements.py is still being applied after this new logic.
The old check unconditionally forces per_device_train_batch_size to be a multiple of num_generations, which counteracts the goal of this pull request.
To ensure this change is effective, the old validation logic should be removed from the grpo_trainer_fix_batch_size function.
There was a problem hiding this comment.
No this is incorrect. It only forces the constraint in a specific case.
1 participant
instead of forcing batch size = num_generations we remove the requirement. This allows for gradient accumulation on entire rollout groups.
The constraint is only forced in a scenario where
steps_per_generation is None,generation_batch_size is Noneandbatch_size*gradient_accumulation % num_generations != 0.This has the added benefit of being backward compatible with our GRPO notebooks.
Qwen3: https://colab.research.google.com/drive/11W2iNnXji4orPGLMUwcHVfpHsy605QH2?usp=sharing
GPT: https://colab.research.google.com/drive/118z5T-85I4gsZvDr36yLfGJ4Lya7qP-V?usp=sharing
Gemma Vision: https://colab.research.google.com/drive/1S2VyBOVzsM4gyb7-I-I6BP__zoOmDDJ4?usp=sharing