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 !

[arthurPignetOwkin]

Hello all, thank for your work.

I would like to checkout this branch and give it a try.
Can you point me toward the best commit to checkout ? Otherwise i will go for the latest.

Cheers,

[min-xu-et]

maybe a naive question: why is "faster" in quotes in the title of this PR? Is it because there is no overlapping between the training and inference step which limits the speed of this multi-node setup?

[binary-husky]

I'm trying to do a similar thing, a difference is that I use VLLM for generation and use NCCL to transfer model parameters:

huggingface/trl#3094

[edbeeching]

closing in favor of #533
@arthurPignetOwkin & @min-xu-et please refer to that implementation which will be merged soon.