GRPO refactoring #2530
GRPO refactoring #2530
Conversation
mydatascience
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jacoguzo,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
suexu1025 and
vipannalla
as code owners
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
src/MaxText/examples/GRPO_README.md
Outdated
| @@ -0,0 +1,226 @@ | |||
| # GRPO Demo - Unified Training Interface | |||
|
|
|||
| This directory contains a unified interface for running GRPO (Group Relative Policy Optimization) training demos across different model sizes and configurations. | |||
There was a problem hiding this comment.
Should we be using the word "demo"?
Don't we anticipate users to use these scripts directly?
There was a problem hiding this comment.
Yeah let's call it grpo_runner.py which calls in grpo_tunix_trainer.py
src/MaxText/examples/GRPO_README.md
Outdated
| - `grpo_llama3_1_8b_demo_pw.py` - Pathways-based 8B model | ||
| - `grpo_llama3_1_70b_demo_pw.py` - Pathways-based 70B model | ||
|
|
||
| These have been consolidated into a single **unified CLI script** (`grpo_demo.py`) that works with the new **grpo.yml** configuration file. |
There was a problem hiding this comment.
again - should be "demo"?
to me, demo indicates it may not be suitable for production workloads
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
mydatascience
force-pushed
the
universal_grpo
branch
from
149a3dc to
e72906d
Compare
Signed-off-by: Vladimir Suvorov <[email protected]>
src/MaxText/examples/README.md
Outdated
| - **`grpo_llama3_demo.ipynb`** → GRPO training on math dataset | ||
| - **`grpo_demo.py`** → Unified CLI for GRPO training (any model) | ||
|
|
||
| #### GRPO Usage |
There was a problem hiding this comment.
Let's call it
GRPO python script usage
src/MaxText/examples/README.md
Outdated
|
|
||
| ### GRPO Training | ||
|
|
||
| - **`grpo_llama3_demo.ipynb`** → GRPO training on math dataset |
There was a problem hiding this comment.
Since you are using #### for the python script, maybe put #### GRPO colab usage here, and can we call it grpo_llama3_1_8b_demo.ipynb
There was a problem hiding this comment.
Yes put there
| if num_vms >= 2: | ||
| # Multi-VM single host setup | ||
| num_devices = len(devices) | ||
| num_trainer_devices = int(num_devices * 0.5) # 50% for training |
There was a problem hiding this comment.
This is not correct
for pathways we use the following and split out the mesh for trainer and inference if there are multiple hosts present
TRAINER_DEVICES_FRACTION = 0.5
SAMPLER_DEVICES_FRACTION = 0.5
if not using pathways, or if once one host
trainer_devices = devices
sampler_devices = devices
There was a problem hiding this comment.
Changed and made it as params to grpo.yml
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
Signed-off-by: Vladimir Suvorov <[email protected]>
3 participants
Description
Refactoring of grpo. Adding new unified functionality allowing to add models easily
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.