Skip to content

Commit a18fbd8

Browse files
BihanBihan RanaAndrey Cheptsov
authored
[Docs]Add Miles Example (#3907)
* [Docs]Add Miles Example * [Docs] Minor Miles training example * [Docs] Minor Miles training example changes * [Docs] Minor Miles training example changes * [Docs] Minor Miles training example changes --------- Co-authored-by: Bihan Rana <bihan@Bihans-MacBook-Pro.local> Co-authored-by: Andrey Cheptsov <andrey.cheptsov@github.com>
1 parent 39d3453 commit a18fbd8

4 files changed

Lines changed: 298 additions & 1 deletion

File tree

‎mkdocs.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,12 +171,14 @@ plugins:
171171
"examples/distributed-training/trl.md": "docs/examples/training/trl.md"
172172
"examples/distributed-training/axolotl.md": "docs/examples/training/axolotl.md"
173173
"examples/distributed-training/ray-ragen.md": "docs/examples/training/ray-ragen.md"
174+
"examples/distributed-training/miles.md": "docs/examples/training/miles.md"
174175
# Single-node and distributed training merged under Training
175176
"docs/examples/single-node-training/trl.md": "docs/examples/training/trl.md"
176177
"docs/examples/single-node-training/axolotl.md": "docs/examples/training/axolotl.md"
177178
"docs/examples/distributed-training/trl.md": "docs/examples/training/trl.md"
178179
"docs/examples/distributed-training/axolotl.md": "docs/examples/training/axolotl.md"
179180
"docs/examples/distributed-training/ray-ragen.md": "docs/examples/training/ray-ragen.md"
181+
"docs/examples/distributed-training/miles.md": "docs/examples/training/miles.md"
180182
"examples/clusters/aws.md": "docs/examples/clusters/aws.md"
181183
"examples/clusters/gcp.md": "docs/examples/clusters/gcp.md"
182184
"examples/clusters/lambda.md": "docs/examples/clusters/lambda.md"
@@ -312,6 +314,7 @@ nav:
312314
- TRL: docs/examples/training/trl.md
313315
- Axolotl: docs/examples/training/axolotl.md
314316
- Ray+RAGEN: docs/examples/training/ray-ragen.md
317+
- Miles: docs/examples/training/miles.md
315318
- Clusters:
316319
- AWS: docs/examples/clusters/aws.md
317320
- GCP: docs/examples/clusters/gcp.md

‎mkdocs/docs/concepts/tasks.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ Jobs on each node communicate using their private IP addresses. Use `DSTACK_MAST
151151
`dstack` is easy to use with `accelerate`, `torchrun`, Ray, Spark, and any other distributed frameworks.
152152

153153
!!! info "Examples"
154-
See the training examples for [TRL](../examples/training/trl.md#distributed-training), [Axolotl](../examples/training/axolotl.md#distributed-training), and [Ray+RAGEN](../examples/training/ray-ragen.md).
154+
See the training examples for [TRL](../examples/training/trl.md#distributed-training), [Axolotl](../examples/training/axolotl.md#distributed-training), [Ray+RAGEN](../examples/training/ray-ragen.md), and [Miles](../examples/training/miles.md).
155155

156156
See the cluster examples for [AWS](../examples/clusters/aws.md), [GCP](../examples/clusters/gcp.md), [Lambda](../examples/clusters/lambda.md), [Crusoe](../examples/clusters/crusoe.md), [Nebius](../examples/clusters/nebius.md), and [NCCL/RCCL tests](../examples/clusters/nccl-rccl-tests.md).
157157

‎mkdocs/docs/examples.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,17 @@ hide:
4949
Fine-tune an agent on multiple nodes with RAGEN, verl, and Ray.
5050
</p>
5151
</a>
52+
53+
<a href="/docs/examples/training/miles"
54+
class="feature-cell">
55+
<h3>
56+
Miles
57+
</h3>
58+
59+
<p>
60+
RL-fine-tune Qwen2.5-32B with Miles.
61+
</p>
62+
</a>
5263
</div>
5364

5465

‎mkdocs/docs/examples/training/miles.md‎

Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
---
2+
title: Miles
3+
description: RL post-training Qwen2.5-32B with Miles, SGLang, Megatron-LM, and Ray across two 8xH100 nodes
4+
---
5+
6+
# Miles
7+
8+
This example shows how to use `dstack` and [Miles](https://github.com/radixark/miles)
9+
for reinforcement learning (RL) post-training of a 32B language model with
10+
[GRPO](https://arxiv.org/abs/2402.03300) across a multi-node cluster.
11+
Miles integrates [SGLang](https://github.com/sgl-project/sglang) for
12+
high-throughput rollouts, [Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
13+
for training, and [Ray](https://docs.ray.io/en/latest/) to coordinate the
14+
trainer and rollout actors across nodes.
15+
16+
Here we fine-tune `Qwen/Qwen2.5-32B-Instruct` on the
17+
[GSM8K](https://huggingface.co/datasets/openai/gsm8k) dataset.
18+
19+
!!! info "Prerequisites"
20+
Multi-node tasks require a [fleet](../../concepts/fleets.md) with
21+
`placement` set to [`cluster`](../../concepts/fleets.md#cluster-placement).
22+
23+
## Run a Ray cluster
24+
25+
### Define a configuration
26+
27+
The [task](../../concepts/tasks.md) below starts Ray on two nodes and prepares
28+
each node by downloading the model and dataset, then converting the checkpoint
29+
to Megatron's `torch_dist` format.
30+
31+
<div editor-title="miles-qwen32b-h100.dstack.yml">
32+
33+
```yaml
34+
type: task
35+
name: miles-qwen32b-h100
36+
nodes: 2
37+
image: radixark/miles:sglang-miles-v0.5.12
38+
env:
39+
- WANDB_API_KEY
40+
- PYTHONPATH=/root/Megatron-LM
41+
- NCCL_DEBUG=INFO
42+
- MODEL_ID=Qwen/Qwen2.5-32B-Instruct
43+
commands:
44+
# 1. Download the model and dataset.
45+
- pip install -U "huggingface_hub[cli]"
46+
- hf download "$MODEL_ID" --local-dir "/root/$(basename "$MODEL_ID")"
47+
- hf download --repo-type dataset openai/gsm8k --local-dir /root/gsm8k
48+
# 2. Convert the Hugging Face checkpoint to Megatron torch_dist.
49+
- |
50+
MODEL_NAME="$(basename "$MODEL_ID")"
51+
cd /root/miles && python tools/convert_hf_to_torch_dist.py \
52+
--swiglu \
53+
--num-layers 64 \
54+
--hidden-size 5120 \
55+
--ffn-hidden-size 27648 \
56+
--num-attention-heads 40 \
57+
--use-rotary-position-embeddings \
58+
--disable-bias-linear \
59+
--add-qkv-bias \
60+
--normalization RMSNorm \
61+
--norm-epsilon 1e-5 \
62+
--rotary-base 1000000 \
63+
--group-query-attention \
64+
--num-query-groups 8 \
65+
--vocab-size 152064 \
66+
--untie-embeddings-and-output-weights \
67+
--hf-checkpoint "/root/$MODEL_NAME" \
68+
--save "/root/${MODEL_NAME}_torch_dist"
69+
# 3. Start Ray.
70+
- |
71+
if [ $DSTACK_NODE_RANK = 0 ]; then
72+
ray start --head --port=6379
73+
else
74+
ray start --address=$DSTACK_MASTER_NODE_IP:6379
75+
fi
76+
ports:
77+
- 8265
78+
resources:
79+
gpu: H100:8
80+
shm_size: 32GB
81+
disk: 1000GB..
82+
volumes:
83+
- /checkpoints:/checkpoints
84+
```
85+
86+
</div>
87+
88+
### Run the configuration
89+
90+
Run the task with [`dstack apply`](../../reference/cli/dstack/apply.md). By
91+
default, `dstack apply` forwards the Ray dashboard port to `localhost:8265`.
92+
93+
<div class="termy">
94+
95+
```shell
96+
$ export WANDB_API_KEY=...
97+
$ dstack apply -f miles-qwen32b-h100.dstack.yml
98+
```
99+
100+
</div>
101+
102+
While `dstack apply` is attached, you can submit Ray jobs through
103+
`localhost:8265`. If you detach or run from another machine, use
104+
[`dstack attach`](../../reference/cli/dstack/attach.md) to re-attach and make
105+
the dashboard port accessible on `localhost`.
106+
107+
> To run on a single node, remove `nodes` or set it to `1`, then submit the job
108+
> with `NUM_NODES=1`. In this case, `placement: cluster` is not required.
109+
110+
## Submit Ray jobs
111+
112+
Install `ray` locally before submitting jobs:
113+
114+
<div class="termy">
115+
116+
```shell
117+
$ pip install ray
118+
```
119+
120+
</div>
121+
122+
The submit script below runs the Miles training job on the Ray cluster. The
123+
model is sharded across all 8 GPUs per node with tensor parallelism, and SGLang
124+
uses the same 8 GPUs per node for rollout.
125+
126+
<div editor-title="submit-miles-train.sh">
127+
128+
```bash
129+
#!/bin/bash
130+
set -euo pipefail
131+
132+
export RAY_ADDRESS=http://localhost:8265
133+
134+
: "${NUM_NODES:?NUM_NODES is not set}"
135+
: "${GPUS_PER_NODE:?GPUS_PER_NODE is not set}"
136+
137+
MODEL_ID="Qwen/Qwen2.5-32B-Instruct"
138+
MODEL_NAME="$(basename "$MODEL_ID")"
139+
HF_CHECKPOINT="/root/$MODEL_NAME"
140+
REF_LOAD="/root/${MODEL_NAME}_torch_dist"
141+
PROMPT_DATA="/root/gsm8k/main/train-00000-of-00001.parquet"
142+
EVAL_PROMPT_DATA="/root/gsm8k/main/test-00000-of-00001.parquet"
143+
INPUT_KEY="question"
144+
LABEL_KEY="answer"
145+
EVAL_DATASET_NAME="gsm8k"
146+
CHECKPOINT_DIR="/checkpoints/${MODEL_NAME}-${EVAL_DATASET_NAME}"
147+
SAVE_INTERVAL=10
148+
WANDB_PROJECT="dstack-miles-RL"
149+
WANDB_GROUP="${MODEL_NAME}-gsm8k-${NUM_NODES}node-${GPUS_PER_NODE}gpu"
150+
WANDB_NAME="rollout-$(date +%Y%m%d-%H%M%S)"
151+
ROLLOUT_GPUS_PER_ENGINE=8
152+
153+
CMD='cd /root/miles && python3 train.py \
154+
--actor-num-nodes '"$NUM_NODES"' \
155+
--actor-num-gpus-per-node '"$GPUS_PER_NODE"' \
156+
--num-gpus-per-node '"$GPUS_PER_NODE"' \
157+
--rollout-num-gpus-per-engine '"$ROLLOUT_GPUS_PER_ENGINE"' \
158+
--sglang-server-concurrency 128 \
159+
--colocate \
160+
--calculate-per-token-loss \
161+
--use-miles-router \
162+
--swiglu \
163+
--num-layers 64 \
164+
--hidden-size 5120 \
165+
--ffn-hidden-size 27648 \
166+
--num-attention-heads 40 \
167+
--use-rotary-position-embeddings \
168+
--disable-bias-linear \
169+
--add-qkv-bias \
170+
--normalization RMSNorm \
171+
--norm-epsilon 1e-5 \
172+
--rotary-base 1000000 \
173+
--group-query-attention \
174+
--num-query-groups 8 \
175+
--vocab-size 152064 \
176+
--untie-embeddings-and-output-weights \
177+
--hf-checkpoint '"$HF_CHECKPOINT"' \
178+
--ref-load '"$REF_LOAD"' \
179+
--prompt-data '"$PROMPT_DATA"' \
180+
--input-key '"$INPUT_KEY"' \
181+
--label-key '"$LABEL_KEY"' \
182+
--apply-chat-template \
183+
--rollout-shuffle \
184+
--rm-type math \
185+
--num-rollout 20 \
186+
--rollout-batch-size 8 \
187+
--n-samples-per-prompt 8 \
188+
--rollout-max-response-len 512 \
189+
--rollout-temperature 1 \
190+
--global-batch-size 64 \
191+
--eval-interval 5 \
192+
--eval-prompt-data '"$EVAL_DATASET_NAME"' '"$EVAL_PROMPT_DATA"' \
193+
--n-samples-per-eval-prompt 1 \
194+
--eval-max-response-len 512 \
195+
--eval-top-k 1 \
196+
--tensor-model-parallel-size 8 \
197+
--sequence-parallel \
198+
--pipeline-model-parallel-size 1 \
199+
--context-parallel-size 1 \
200+
--expert-model-parallel-size 1 \
201+
--expert-tensor-parallel-size 1 \
202+
--use-dynamic-batch-size \
203+
--max-tokens-per-gpu 9216 \
204+
--advantage-estimator grpo \
205+
--use-kl-loss \
206+
--kl-loss-coef 0.00 \
207+
--kl-loss-type low_var_kl \
208+
--kl-coef 0.00 \
209+
--entropy-coef 0.00 \
210+
--eps-clip 0.2 \
211+
--eps-clip-high 0.28 \
212+
--optimizer adam \
213+
--lr 1e-6 \
214+
--lr-decay-style constant \
215+
--weight-decay 0.1 \
216+
--adam-beta1 0.9 \
217+
--adam-beta2 0.98 \
218+
--sglang-mem-fraction-static 0.7 \
219+
--use-wandb \
220+
--wandb-host https://wandb.ai/ \
221+
--wandb-project '"$WANDB_PROJECT"' \
222+
--wandb-group '"$WANDB_GROUP"' \
223+
--wandb-exp-name '"$WANDB_NAME"' \
224+
--attention-dropout 0.0 \
225+
--hidden-dropout 0.0 \
226+
--accumulate-allreduce-grads-in-fp32 \
227+
--attention-softmax-in-fp32 \
228+
--attention-backend flash \
229+
--save '"$CHECKPOINT_DIR"' \
230+
--save-interval '"$SAVE_INTERVAL"''
231+
232+
# GLOO_SOCKET_IFNAME=eth0 is required for multi-node Gloo process group init.
233+
# Without it, Gloo resolves to a loopback address (127.0.1.1) instead of the
234+
# inter-node interface, causing `init_gloo_group()` to timeout.
235+
RUNTIME_ENV_JSON=$(cat <<EOF
236+
{
237+
"env_vars": {
238+
"PYTHONPATH": "/root/Megatron-LM",
239+
"CUDA_DEVICE_MAX_CONNECTIONS": "1",
240+
"NCCL_DEBUG": "INFO",
241+
"GLOO_SOCKET_IFNAME": "eth0"
242+
}
243+
}
244+
EOF
245+
)
246+
247+
ray job submit \
248+
--address="$RAY_ADDRESS" \
249+
--runtime-env-json="$RUNTIME_ENV_JSON" \
250+
-- bash -lc "$CMD"
251+
```
252+
253+
</div>
254+
255+
Submit the job with the same cluster shape as the task:
256+
257+
<div class="termy">
258+
259+
```shell
260+
$ NUM_NODES=2 GPUS_PER_NODE=8 bash submit-miles-train.sh
261+
```
262+
263+
</div>
264+
265+
!!! info "Training parameters"
266+
1. `--tensor-model-parallel-size 8` shards the 32B model across all 8 GPUs
267+
per node.
268+
2. `--rollout-num-gpus-per-engine 8` starts SGLang with TP-8 on each node.
269+
3. `--sglang-server-concurrency` sets how many requests SGLang processes
270+
concurrently.
271+
4. `--max-tokens-per-gpu 9216` sets the per-GPU token budget. Lower this if
272+
Megatron OOMs during training.
273+
5. `--sglang-mem-fraction-static 0.7` sets the SGLang KV cache memory
274+
fraction. Lower this if Megatron OOMs at startup.
275+
276+
Using Ray via `dstack` gives you access to the Ray ecosystem while benefiting
277+
from `dstack`'s provisioning capabilities.
278+
279+
!!! info "What's next"
280+
1. Read about [distributed tasks](../../concepts/tasks.md#distributed-tasks)
281+
and [fleets](../../concepts/fleets.md)
282+
2. See the [SGLang inference](../inference/sglang.md) example
283+
3. Browse Miles' [examples](https://github.com/radixark/miles/tree/main/examples)

0 commit comments

Comments
 (0)