Train GRPO through SGLangJax in Single or Multi Hosts with Pathways
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 345
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 240
Description
Note:
- Current situation is we succeed to run `python3 scripts/grpo_demo_llama3_qwen2.py` in single and multi hosts with Pathways. The performance needs to optimize.
- Multi hosts locates in the same Pod. Currently we have not tested the multi hosts across Pods.
- If you want to use Pathways, please ensure you have access to GKE. Now allow list is not required.
- How to setup GKE with Pathways? Maybe you can refer to [run_in_pathways.md](https://github.com/sgl-project/sglang-jax/blob/main/docs/features/run_in_pathways.md).
# Single Host with Pathways
**Environment**
- topology: 2x2
- hosts: 1
- machine_type: ct6e-standard-4t
**Instruction**
```bash
# Execute the following script on any hosts where you can connect to GKE.
JAX_PLATFORMS=proxy \
JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 \
python3 scripts/grpo_demo_llama3_qwen2.py \
--num-batches 4 \
--num-test-batches 1 \
--root-dir=/home/gcpuser/aolemila \
--model-version=meta-llama/Llama-3.2-1B-Instruct \
--data-source tfds \
--dataset gsm8k \
--rollout-engine sglang_jax
```
**Expected Result**
# Multi Hosts in the Pod with Pathways
**Environment**
- topology: 4x4
- hosts: 2
- machine_type: ct6e-standard-4t
**Instruction**
```bash
# Execute the following script on any hosts where you can connect to GKE.
JAX_PLATFORMS=proxy \
JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 \
python3 scripts/grpo_demo_llama3_qwen2.py \
--num-batches 4 \
--num-test-batches 1 \
--root-dir=/home/gcpuser/aolemila \
--model-version=meta-llama/Llama-3.2-1B-Instruct \
--data-source tfds \
--dataset gsm8k \
--rollout-engine sglang_jax \
--rollout-tp 8 \
--cluster-setup disaggregated-2-way
```
Note: `--rollout-tp 8` and `--cluster-setup disaggregated-2-way` are required due to SGLangJax lacks to support DP now. It will support DP soon.
**Expected Result**
Contributor guide
Assessment
This issue has not been assessed yet.