WIP "Faster" grpo trainer

[edbeeching]

GRPO trainer to train on N + 1 nodes, with 1 node allocated for generation. Very experimental, so expect hard edges!

Usage:

For training, run:

accelerate launch --config_file=recipes/accelerate_configs/zero3.yaml scripts/remote_grpo.py \
    --config recipes/Qwen2.5-1.5B-Instruct/grpo/config_remote.yaml

This will automatically spin up an SGLang server on a separate Slurm node and use it for generation.

For development, first spin up an SGLang sever on a separate node:

python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-1.5B-Instruct   --port=30010 --skip-tokenizer-init --mem-fraction-static 0.7 --host=0.0.0.0 --dp-size=8

Then run training by providing the IP address of the server:

accelerate launch --config_file=recipes/accelerate_configs/zero3.yaml scripts/remote_grpo.py \
    --config recipes/Qwen2.5-1.5B-Instruct/grpo/config_remote.yaml \
    --remote_gen_model_url ip-26-0-160-103

TODO

• Remove hard-coded filepath for temporary checkpoints
• Refactor reference model log probs to happen within generation step
• Implement μ iterations from GRPO
• Validate against TRL

[troy12x]

please my goat push this to main pleaseeee

[troy12x]

my project is waiting for ur code no joke

[troy12x]

hi 👉👈

[troy12x]

please finish this my goat

[qgallouedec]

@troy12x please avoid spamming 🙏 it doesn't help, be sure that we're working hard on this.

[troy12x]

sry i dont mean but i really need you guys to finish this fast

[troy12x]

good luck !