diff --git a/mkdocs/docs/concepts/tasks.md b/mkdocs/docs/concepts/tasks.md index ea1715aa8..0d82e5d2a 100644 --- a/mkdocs/docs/concepts/tasks.md +++ b/mkdocs/docs/concepts/tasks.md @@ -179,6 +179,118 @@ Jobs on each node communicate using their private IP addresses. Use `DSTACK_MAST For convenience, `~/.ssh/config` is preconfigured with these options, so a simple `ssh ` is enough. For a list of nodes IPs check the `DSTACK_NODES_IPS` environment variable. +### Node groups + +A task can define multiple node groups. Each group has its own `nodes` count, +`resources`, `commands`, and `ports`. + +
+ +```yaml +type: task +name: ray-cluster + +python: 3.12 + +groups: + - name: head + nodes: 1 + commands: + - pip uninstall -y ray && pip install -U "ray[default]" + - ray start --head --port=6379 --block + resources: + cpu: 2 + memory: 4GB.. + ports: + - 8265 + + - name: workers + nodes: 2 + commands: + - pip uninstall -y ray && pip install -U "ray[default]" + - ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379 --block + resources: + gpu: H100:8 +``` + +
+ +Commands in any group can reference the internal IP address of any node in the run via +`${{ groups[i].nodes[j].IP_ADDRESS }}`, where `i` is the index of the group in `groups` and `j` is +the index of the node within that group. + +> `groups[0].nodes[0]` is the run's master node — it is what `DSTACK_MASTER_NODE_IP` resolves to. + +Currently, only `resources`, `commands`, and `ports` can be configured per node group. [`groups`](../reference/dstack.yml/task.md#groups) and top-level `nodes` are mutually exclusive.Support for other properties is coming soon. + +??? info "Prefill/decode example" + Node groups can mix CPU and GPU roles. This SGLang prefill/decode split uses a CPU + router (`groups[0]`, the master) and GPU workers. `startup_order: workers-first` + starts prefill and decode before the router. + +
+ + ```yaml + type: task + name: prefill-decode + image: lmsysorg/sglang:v0.5.10.post1 + env: + - HF_TOKEN + - MODEL_ID=zai-org/GLM-4.5-Air-FP8 + + startup_order: workers-first + groups: + # Router (CPU) — master node; wires prefill + decode by IP + - name: router + nodes: 1 + commands: + - pip install smg + - | + echo "prefill=${{ groups[1].nodes[0].IP_ADDRESS }}" + echo "decode=${{ groups[2].nodes[0].IP_ADDRESS }}" + smg launch \ + --pd-disaggregation \ + --prefill http://${{ groups[1].nodes[0].IP_ADDRESS }}:8000 8998 \ + --decode http://${{ groups[2].nodes[0].IP_ADDRESS }}:8000 \ + --prefill-policy cache_aware \ + --host 0.0.0.0 --port 8000 + ports: + - 8000 + resources: + cpu: 4 + + - name: prefill + nodes: 1 + commands: + - | + python -m sglang.launch_server \ + --model-path $MODEL_ID \ + --disaggregation-mode prefill \ + --disaggregation-transfer-backend nixl \ + --host 0.0.0.0 --port 8000 \ + --disaggregation-bootstrap-port 8998 + resources: + gpu: H200 + + - name: decode + nodes: 1 + commands: + - | + python -m sglang.launch_server \ + --model-path $MODEL_ID \ + --disaggregation-mode decode \ + --disaggregation-transfer-backend nixl \ + --host 0.0.0.0 --port 8000 + resources: + gpu: H200 + ``` + +
+ +!!! info "Examples" + See the [Ray+RAGEN](../examples/training/ray-ragen.md) example for running a Ray cluster, + and the [NCCL/RCCL tests](../examples/clusters/nccl-rccl-tests.md) example for running `mpirun` with node groups. + ### Resources When you specify a resource value like `cpu` or `memory`, @@ -460,7 +572,7 @@ If you don't assign a value to an environment variable (see `HF_TOKEN` above), | `DSTACK_NODE_RANK` | The rank of the node | | `DSTACK_MASTER_NODE_IP` | The internal IP address of the master node | | `DSTACK_NODES_IPS` | The list of internal IP addresses of all nodes delimited by "\n" | - | `DSTACK_MPI_HOSTFILE` | The path to a pre-populated MPI hostfile | + | `DSTACK_MPI_HOSTFILE` | The path to a pre-populated MPI hostfile. The file lists GPU nodes as ` slots=` and CPU nodes as `` | | `DSTACK_WORKING_DIR` | The working directory of the run | | `DSTACK_REPO_DIR` | The directory where the repo is mounted (if any) | diff --git a/mkdocs/docs/examples/clusters/nccl-rccl-tests.md b/mkdocs/docs/examples/clusters/nccl-rccl-tests.md index 775e2c8b8..d2ad32dd3 100644 --- a/mkdocs/docs/examples/clusters/nccl-rccl-tests.md +++ b/mkdocs/docs/examples/clusters/nccl-rccl-tests.md @@ -12,7 +12,7 @@ This example shows how to run [NCCL](https://github.com/NVIDIA/nccl-tests) or [R ## Running as a task -Here's an example of a task that runs AllReduce test on 2 nodes, each with 4 GPUs (8 processes in total). +Here's an example of a task that runs AllReduce test on 2 nodes, each with 4 GPUs (8 processes in total), using [node groups](../../concepts/tasks.md#node-groups). === "NCCL tests" @@ -22,33 +22,38 @@ Here's an example of a task that runs AllReduce test on 2 nodes, each with 4 GPU type: task name: nccl-tests - nodes: 2 - startup_order: workers-first stop_criteria: master-done env: - NCCL_DEBUG=INFO - commands: - - | - if [ $DSTACK_NODE_RANK -eq 0 ]; then - mpirun \ - --allow-run-as-root \ - --hostfile $DSTACK_MPI_HOSTFILE \ - -n $DSTACK_GPUS_NUM \ - -N $DSTACK_GPUS_PER_NODE \ - --bind-to none \ - /opt/nccl-tests/build/all_reduce_perf -b 8 -e 8G -f 2 -g 1 - else - sleep infinity - fi + + groups: + - name: master # The name property is optional + nodes: 1 + commands: + - | + mpirun \ + --allow-run-as-root \ + --hostfile $DSTACK_MPI_HOSTFILE \ + -n $DSTACK_GPUS_NUM \ + -N $DSTACK_GPUS_PER_NODE \ + --bind-to none \ + /opt/nccl-tests/build/all_reduce_perf -b 8 -e 8G -f 2 -g 1 + resources: + gpu: nvidia:1..8 + shm_size: 16GB + + - name: workers + nodes: 1 + commands: + - sleep infinity + resources: + gpu: nvidia:1..8 + shm_size: 16GB # Uncomment if the `kubernetes` backend requires it for `/dev/infiniband` access #privileged: true - - resources: - gpu: nvidia:1..8 - shm_size: 16GB ``` @@ -65,7 +70,6 @@ Here's an example of a task that runs AllReduce test on 2 nodes, each with 4 GPU type: task name: rccl-tests - nodes: 2 startup_order: workers-first stop_criteria: master-done @@ -77,35 +81,50 @@ Here's an example of a task that runs AllReduce test on 2 nodes, each with 4 GPU env: - NCCL_DEBUG=INFO - OPEN_MPI_HOME=/usr/lib/x86_64-linux-gnu/openmpi - commands: - # Setup MPI and build RCCL tests - - apt-get install -y git libopenmpi-dev openmpi-bin - - git clone https://github.com/ROCm/rccl-tests.git - - cd rccl-tests - - make MPI=1 MPI_HOME=$OPEN_MPI_HOME - - # Preload the RoCE driver library from the host (for Broadcom driver compatibility) - - export LD_PRELOAD=/mnt/lib/libbnxt_re-rdmav34.so - - # Run RCCL tests via MPI - - | - if [ $DSTACK_NODE_RANK -eq 0 ]; then - mpirun --allow-run-as-root \ - --hostfile $DSTACK_MPI_HOSTFILE \ - -n $DSTACK_GPUS_NUM \ - -N $DSTACK_GPUS_PER_NODE \ - --mca btl_tcp_if_include ens41np0 \ - -x LD_PRELOAD \ - -x NCCL_IB_HCA=mlx5_0/1,bnxt_re0,bnxt_re1,bnxt_re2,bnxt_re3,bnxt_re4,bnxt_re5,bnxt_re6,bnxt_re7 \ - -x NCCL_IB_GID_INDEX=3 \ - -x NCCL_IB_DISABLE=0 \ - ./build/all_reduce_perf -b 8M -e 8G -f 2 -g 1 -w 5 --iters 20 -c 0; - else - sleep infinity - fi - - resources: - gpu: MI300X:8 + + groups: + - name: master # The name property is optional + nodes: 1 + commands: + # Setup MPI and build RCCL tests + - apt-get install -y git libopenmpi-dev openmpi-bin + - git clone https://github.com/ROCm/rccl-tests.git + - cd rccl-tests + - make MPI=1 MPI_HOME=$OPEN_MPI_HOME + + # Preload the RoCE driver library from the host (for Broadcom driver compatibility) + - export LD_PRELOAD=/mnt/lib/libbnxt_re-rdmav34.so + + # Run RCCL tests via MPI + - | + mpirun --allow-run-as-root \ + --hostfile $DSTACK_MPI_HOSTFILE \ + -n $DSTACK_GPUS_NUM \ + -N $DSTACK_GPUS_PER_NODE \ + --mca btl_tcp_if_include ens41np0 \ + -x LD_PRELOAD \ + -x NCCL_IB_HCA=mlx5_0/1,bnxt_re0,bnxt_re1,bnxt_re2,bnxt_re3,bnxt_re4,bnxt_re5,bnxt_re6,bnxt_re7 \ + -x NCCL_IB_GID_INDEX=3 \ + -x NCCL_IB_DISABLE=0 \ + ./build/all_reduce_perf -b 8M -e 8G -f 2 -g 1 -w 5 --iters 20 -c 0; + resources: + gpu: MI300X:8 + + - name: workers + nodes: 1 + commands: + # Setup MPI and build RCCL tests + - apt-get install -y git libopenmpi-dev openmpi-bin + - git clone https://github.com/ROCm/rccl-tests.git + - cd rccl-tests + - make MPI=1 MPI_HOME=$OPEN_MPI_HOME + + # Preload the RoCE driver library from the host (for Broadcom driver compatibility) + - export LD_PRELOAD=/mnt/lib/libbnxt_re-rdmav34.so + + - sleep infinity + resources: + gpu: MI300X:8 ``` diff --git a/mkdocs/docs/examples/training/miles.md b/mkdocs/docs/examples/training/miles.md index 59451d150..7ff038e25 100644 --- a/mkdocs/docs/examples/training/miles.md +++ b/mkdocs/docs/examples/training/miles.md @@ -24,61 +24,97 @@ Here we fine-tune `Qwen/Qwen2.5-32B-Instruct` on the ### Define a configuration -The [task](../../concepts/tasks.md) below starts Ray on two nodes and prepares -each node by downloading the model and dataset, then converting the checkpoint -to Megatron's `torch_dist` format. +The [task](../../concepts/tasks.md) below starts Ray on two GPU nodes using +[node groups](../../concepts/tasks.md#node-groups) and prepares each node by +downloading the model and dataset, then converting the checkpoint to Megatron's +`torch_dist` format.
```yaml type: task name: miles-qwen32b-h100 -nodes: 2 image: radixark/miles:sglang-miles-v0.5.12 env: - WANDB_API_KEY - PYTHONPATH=/root/Megatron-LM - NCCL_DEBUG=INFO - MODEL_ID=Qwen/Qwen2.5-32B-Instruct -commands: - # 1. Download the model and dataset. - - pip install -U "huggingface_hub[cli]" - - hf download "$MODEL_ID" --local-dir "/root/$(basename "$MODEL_ID")" - - hf download --repo-type dataset openai/gsm8k --local-dir /root/gsm8k - # 2. Convert the Hugging Face checkpoint to Megatron torch_dist. - - | - MODEL_NAME="$(basename "$MODEL_ID")" - cd /root/miles && python tools/convert_hf_to_torch_dist.py \ - --swiglu \ - --num-layers 64 \ - --hidden-size 5120 \ - --ffn-hidden-size 27648 \ - --num-attention-heads 40 \ - --use-rotary-position-embeddings \ - --disable-bias-linear \ - --add-qkv-bias \ - --normalization RMSNorm \ - --norm-epsilon 1e-5 \ - --rotary-base 1000000 \ - --group-query-attention \ - --num-query-groups 8 \ - --vocab-size 152064 \ - --untie-embeddings-and-output-weights \ - --hf-checkpoint "/root/$MODEL_NAME" \ - --save "/root/${MODEL_NAME}_torch_dist" - # 3. Start Ray. - - | - if [ $DSTACK_NODE_RANK = 0 ]; then - ray start --head --port=6379 - else - ray start --address=$DSTACK_MASTER_NODE_IP:6379 - fi -ports: - - 8265 -resources: - gpu: H100:8 - shm_size: 32GB - disk: 1000GB.. + +groups: + - name: head # The name property is optional + nodes: 1 + commands: + # 1. Download the model and dataset. + - pip install -U "huggingface_hub[cli]" + - hf download "$MODEL_ID" --local-dir "/root/$(basename "$MODEL_ID")" + - hf download --repo-type dataset openai/gsm8k --local-dir /root/gsm8k + # 2. Convert the Hugging Face checkpoint to Megatron torch_dist. + - | + MODEL_NAME="$(basename "$MODEL_ID")" + cd /root/miles && python tools/convert_hf_to_torch_dist.py \ + --swiglu \ + --num-layers 64 \ + --hidden-size 5120 \ + --ffn-hidden-size 27648 \ + --num-attention-heads 40 \ + --use-rotary-position-embeddings \ + --disable-bias-linear \ + --add-qkv-bias \ + --normalization RMSNorm \ + --norm-epsilon 1e-5 \ + --rotary-base 1000000 \ + --group-query-attention \ + --num-query-groups 8 \ + --vocab-size 152064 \ + --untie-embeddings-and-output-weights \ + --hf-checkpoint "/root/$MODEL_NAME" \ + --save "/root/${MODEL_NAME}_torch_dist" + # 3. Start Ray. + - ray start --head --port=6379 --block + ports: + - 8265 + resources: + gpu: H100:8 + shm_size: 32GB + disk: 1000GB.. + + - name: workers + nodes: 1 + commands: + # 1. Download the model and dataset. + - pip install -U "huggingface_hub[cli]" + - hf download "$MODEL_ID" --local-dir "/root/$(basename "$MODEL_ID")" + - hf download --repo-type dataset openai/gsm8k --local-dir /root/gsm8k + # 2. Convert the Hugging Face checkpoint to Megatron torch_dist. + - | + MODEL_NAME="$(basename "$MODEL_ID")" + cd /root/miles && python tools/convert_hf_to_torch_dist.py \ + --swiglu \ + --num-layers 64 \ + --hidden-size 5120 \ + --ffn-hidden-size 27648 \ + --num-attention-heads 40 \ + --use-rotary-position-embeddings \ + --disable-bias-linear \ + --add-qkv-bias \ + --normalization RMSNorm \ + --norm-epsilon 1e-5 \ + --rotary-base 1000000 \ + --group-query-attention \ + --num-query-groups 8 \ + --vocab-size 152064 \ + --untie-embeddings-and-output-weights \ + --hf-checkpoint "/root/$MODEL_NAME" \ + --save "/root/${MODEL_NAME}_torch_dist" + # 3. Start Ray. + # groups[0].nodes[0] is the head node in the `head` group + - ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379 --block + resources: + gpu: H100:8 + shm_size: 32GB + disk: 1000GB.. + volumes: - /checkpoints:/checkpoints ``` @@ -104,7 +140,7 @@ While `dstack apply` is attached, you can submit Ray jobs through [`dstack attach`](../../reference/cli/dstack/attach.md) to re-attach and make the dashboard port accessible on `localhost`. -> To run on a single node, remove `nodes` or set it to `1`, then submit the job +> To run on a single node, use one group with `nodes: 1`, then submit the job > with `NUM_NODES=1`. In this case, `placement: cluster` is not required. ## Submit Ray jobs diff --git a/mkdocs/docs/examples/training/ray-ragen.md b/mkdocs/docs/examples/training/ray-ragen.md index 73e8749e8..fd72ba563 100644 --- a/mkdocs/docs/examples/training/ray-ragen.md +++ b/mkdocs/docs/examples/training/ray-ragen.md @@ -25,38 +25,53 @@ The task below runs a Ray cluster on an existing fleet: type: task name: ray-cluster -nodes: 2 - -env: -- WANDB_API_KEY image: whatcanyousee/verl:ngc-cu124-vllm0.8.5-sglang0.4.6-mcore0.12.0-te2.2 -commands: - - wget -O miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh - - bash miniconda.sh -b -p /workflow/miniconda - - eval "$(/workflow/miniconda/bin/conda shell.bash hook)" - - git clone https://github.com/RAGEN-AI/RAGEN.git - - cd RAGEN - - bash scripts/setup_ragen.sh - - conda activate ragen - - cd verl - - pip install --no-deps -e . - - pip install hf_transfer hf_xet - - pip uninstall -y ray - - pip install -U "ray[default]" - - | - if [ $DSTACK_NODE_RANK = 0 ]; then - ray start --head --port=6379; - else - ray start --address=$DSTACK_MASTER_NODE_IP:6379 - fi - -# Expose Ray dashboard port -ports: - - 8265 - -resources: - gpu: 80GB:8 - shm_size: 128GB +env: + - WANDB_API_KEY + +groups: + - name: head # The name property is optional + nodes: 1 + commands: + - wget -O miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh + - bash miniconda.sh -b -p /workflow/miniconda + - eval "$(/workflow/miniconda/bin/conda shell.bash hook)" + - git clone https://github.com/RAGEN-AI/RAGEN.git + - cd RAGEN + - bash scripts/setup_ragen.sh + - conda activate ragen + - cd verl + - pip install --no-deps -e . + - pip install hf_transfer hf_xet + - pip uninstall -y ray + - pip install -U "ray[default]" + - ray start --head --port=6379 --block + ports: + - 8265 + resources: + gpu: 80GB:8 + shm_size: 128GB + + - name: workers + nodes: 1 + commands: + - wget -O miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh + - bash miniconda.sh -b -p /workflow/miniconda + - eval "$(/workflow/miniconda/bin/conda shell.bash hook)" + - git clone https://github.com/RAGEN-AI/RAGEN.git + - cd RAGEN + - bash scripts/setup_ragen.sh + - conda activate ragen + - cd verl + - pip install --no-deps -e . + - pip install hf_transfer hf_xet + - pip uninstall -y ray + - pip install -U "ray[default]" + # groups[0].nodes[0] is the head node in the `head` group + - ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379 --block + resources: + gpu: 80GB:8 + shm_size: 128GB # Save checkpoints on the instance volumes: diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 575fe12c6..3621e27f4 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -82,7 +82,7 @@ tasks, and services: ``` - `DSTACK_NODES_IPS`{ #DSTACK_NODES_IPS } – The list of internal IP addresses of all nodes delimited by `"\n"`. -- `DSTACK_MPI_HOSTFILE`{ #DSTACK_MPI_HOSTFILE } – The path to a pre-populated MPI hostfile that can be used directly as `mpirun --hostfile $DSTACK_MPI_HOSTFILE`. +- `DSTACK_MPI_HOSTFILE`{ #DSTACK_MPI_HOSTFILE } – The path to a pre-populated MPI hostfile that can be used directly as `mpirun --hostfile $DSTACK_MPI_HOSTFILE`. The file lists GPU nodes as ` slots=` and CPU nodes as ``. ## Server