google / google/tunix

Train GRPO through SGLangJax in Single or Multi Hosts with Pathways

Open
#874 0 comments 0 reactions 1 assignee Claimed by @aolemila View on GitHub
type:feature/enhancement
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**
image

# 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**
Pasted Graphic 9

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.