Deploying SGLang on Slurm with AMD GPUs across multiple nodes (partial GPU usage)

ROCm/AMD practical rollout: resource allocation, environment variables, communication, and pitfalls

This article focuses on deployment concepts and reusable script templates. Module names, partition settings, and NIC devices vary across clusters—adjust accordingly.

What this work focuses on

  • Partial GPUs per node (e.g., two 8-GPU nodes, 4 GPUs used on each): a per GPU-node strategy to set HIP_VISIBLE_DEVICES precisely.
  • Device allocation in SGLang aligned with Slurm GPU IDs, preserving topology across nodes and NUMA domains.
  • Server config clarified/fixed: MASTER_ADDR/PORT, nccl-init-addr, and per-node cache path isolation.
  • Decode throughput experiments on MI210 nodes (8 GPUs per node, two NUMA nodes), measuring same-NUMA vs cross-NUMA performance.

Background & Goal

In HPC clusters, Slurm is a common scheduler. AMD GPUs rely on ROCm and RCCL (PyTorch still uses the NCCL interface) for acceleration and communication. This guide explains how to reliably deploy SGLang inference services in a multi-node setup where each node uses only part of its GPUs.

Nodes & topology (MI210)

  • Each node hosts 8 MI210 GPUs, forming two NUMA nodes.
  • MI210 supports up to 3 Infinity Fabric slots; bandwidth differs significantly between same-NUMA and cross-NUMA connections.
  • Experiments compare decode throughput for same-NUMA vs cross-NUMA placements.

Prerequisites

  • OS/kernel: cluster defaults (e.g., RHEL/CentOS/Ubuntu LTS).
  • Drivers/libs: matching ROCm version; PyTorch built for ROCm (pip/conda with rocm wheel).
  • Communication: RCCL (via NCCL envs), IB or RoCE NIC (e.g., ib0).
  • Scheduling: Slurm configured with GresTypes=gpu, nodes labeled with GPU types (e.g., mi250).

Key points for partial GPU allocation

When using a subset of GPUs per node (e.g., 3/8), ensure:

  • Request exact GPUs via --gres=gpu:<type>:<count> to avoid conflicts.
  • Limit visible devices properly: align HIP_VISIBLE_DEVICES (and possibly CUDA_VISIBLE_DEVICES) with Slurm allocation.
  • Cross-node communication only happens among allocated processes; set WORLD_SIZE/RANK, MASTER_ADDR/MASTER_PORT correctly.

Useful Slurm environment variables

  • SLURM_NODELIST: node list of the job.
  • SLURM_NNODES: number of nodes.
  • SLURM_PROCID / SLURM_NODEID: global/node-local task indices.
  • SLURM_JOB_GPUS: GPUs allocated to this job on the current node (depends on Slurm version).

Recommended environment variables (ROCm/AMD)

# Expose only allocated GPUs (indices provided by Slurm)
export HIP_VISIBLE_DEVICES=${SLURM_JOB_GPUS}

# Communication (RCCL via NCCL envs)
export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=ib0       # adjust to your NIC
export NCCL_IB_DISABLE=0
export NCCL_P2P_DISABLE=0
export RDMAV_FORK_SAFE=1

# Stability/perf tweaks in some environments
export HSA_FORCE_FINE_GRAIN_PCIE=1
export HSA_ENABLE_SDMA=1

Multi-node communication params

MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -n 1)
MASTER_PORT=30000                    # pick a free port
WORLD_SIZE=${SLURM_NNODES}
RANK=${SLURM_PROCID}                 # if 1 task per node, RANK=node index

export MASTER_ADDR MASTER_PORT WORLD_SIZE RANK

Deployment approach & server configuration for SGLang

Typical setup: one process per node, configure TP / PP / DP as needed, and use PyTorch NCCL backend (bridged to RCCL under ROCm). For AMD/ROCm, use the NCCL backend.

  • Explicitly set MASTER_ADDR / MASTER_PORT (or nccl-init-addr) using the head node hostname and a fixed port.
  • Isolate caches per node (e.g., OUTLINES_CACHE_DIR) to avoid sqlite conflicts in multi-node jobs.
SGLang CLI flags may vary across versions. The examples below illustrate parameter organization and env passing—confirm exact flag names for your version.

Example 1: use SGLang server entry directly

python -m sglang.launch_server \
  --model /path/to/model \
  --host 0.0.0.0 \
  --port ${MASTER_PORT} \
  --tensor-parallel-size 3 \        # 3 GPUs per node (example)
  --pipeline-parallel-size 1 \
  --distributed-backend nccl \      # AMD/ROCm typically still uses nccl
  --log-level info

Example 2: wrap with torchrun (convenient for RANK passing)

torchrun --nproc_per_node=1 --nnodes=${WORLD_SIZE} --node_rank=${RANK} \
  --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
  -m sglang.launch_server \
  --model /path/to/model \
  --host 0.0.0.0 --port ${MASTER_PORT} \
  --tensor-parallel-size 3 --pipeline-parallel-size 1 --distributed-backend nccl

Complete sbatch template (multi-node, partial GPUs)

#!/bin/bash
#SBATCH -J sglang-amd
#SBATCH -p compute                 # partition
#SBATCH -N 2                       # 2 nodes
#SBATCH --ntasks-per-node=1        # one process per node
#SBATCH --gres=gpu:mi210:4         # 4 AMD MI210 per node
#SBATCH --cpus-per-task=16
#SBATCH -t 02:00:00
#SBATCH --exclusive

module purge
module load rocm                   # e.g., rocm/5.7
module load python/3.10

# Activate Python env if needed
# source ~/miniconda3/bin/activate sglang-rocm

MASTER_ADDR=$(scontrol show hostnames $SLURM_NODELIST | head -n 1)
MASTER_PORT=30000
WORLD_SIZE=${SLURM_NNODES}
RANK=${SLURM_PROCID}

export MASTER_ADDR MASTER_PORT WORLD_SIZE RANK
if [[ -n "${SLURM_JOB_GPUS}" ]]; then
  export HIP_VISIBLE_DEVICES=${SLURM_JOB_GPUS}
elif [[ -n "${CUDA_VISIBLE_DEVICES}" ]]; then
  export HIP_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}
else
  export HIP_VISIBLE_DEVICES=0,1,2,3
fi
export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=ib0
export NCCL_IB_DISABLE=0
export NCCL_P2P_DISABLE=0
export RDMAV_FORK_SAFE=1
export HSA_FORCE_FINE_GRAIN_PCIE=1

echo "[INFO] Nodes: ${WORLD_SIZE}, Rank: ${RANK}, Master: ${MASTER_ADDR}:${MASTER_PORT}"
echo "[INFO] HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES}"

# Option 1: launch SGLang server directly
python -m sglang.launch_server \
  --model /path/to/model \
  --host 0.0.0.0 \
  --port ${MASTER_PORT} \
  --tensor-parallel-size 4 \
  --pipeline-parallel-size 1 \
  --distributed-backend nccl \
  --log-level info

# Option 2: torchrun (comment out above and enable below)
# torchrun --nproc_per_node=1 --nnodes=${WORLD_SIZE} --node_rank=${RANK} \
#   --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
#   -m sglang.launch_server \
#   --model /path/to/model \
#   --host 0.0.0.0 --port ${MASTER_PORT} \
#   --tensor-parallel-size 3 --pipeline-parallel-size 1 --distributed-backend nccl

Practical script: complete multi-node sbatch (annotated, router + multi-worker)

The snippets below are taken from the full multinode.sbatch (available at the site root for download). It adopts one SGLang worker per node (local TP) + sglang_router handling cross-node DP, which is suitable for partial GPUs per node on MI210 and convenient for same- vs cross-NUMA throughput comparisons.

Download the full sbatch (English, multinode_en.sbatch)

Job header and resource request (example)

#!/usr/bin/env bash
#SBATCH -p faculty
#SBATCH --qos=gtqos
#SBATCH --nodes 2
#SBATCH --ntasks=2
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=64
#SBATCH --mem=256G
#SBATCH -t 12:00:00
#SBATCH -J qwen30b_tp4_dp2
# Note: tune these per your cluster partitions/CPU/memory.

Set per-node visible GPUs precisely

set_local_gpus() {
  if [[ -n "${SLURM_STEP_GPUS:-}" ]]; then
    HIP_VISIBLE_DEVICES="${SLURM_STEP_GPUS}"
  elif [[ -n "${SLURM_JOB_GPUS:-}" ]]; then
    HIP_VISIBLE_DEVICES="${SLURM_JOB_GPUS}"
  else
    HIP_VISIBLE_DEVICES="${HIP_DEVICES_DEFAULT:-0,1,2,3}"
  fi
  export HIP_VISIBLE_DEVICES
  export ROCR_VISIBLE_DEVICES="$HIP_VISIBLE_DEVICES"
  export HSA_VISIBLE_DEVICES="$HIP_VISIBLE_DEVICES"
  export CUDA_VISIBLE_DEVICES="$HIP_VISIBLE_DEVICES"
}
          

Launch per-node worker (local TP) with cache isolation

srun --nodes=1 --nodelist="$node" --gres="gpu:${WORKER_GPUS_PER_NODE}" \
     --ntasks=1 --cpus-per-task="$SLURM_CPUS_PER_TASK" \
     bash -lc "$(declare -f set_local_gpus); set_local_gpus; \
       export OUTLINES_CACHE_DIR=\"/tmp/\$(hostname -s)_cache_${SLURM_JOB_ID}\"; \
       python3 -m sglang.launch_server \
         --tp-size $TP_SIZE --pp-size $PP_SIZE \
         --model-path $MODEL_PATH --host 0.0.0.0 --port $port \
         --disable-radix-cache --enable-dp-attention \
         --device cuda" &
          

Router (cross-node DP) and health checks

python3 -m sglang_router.launch_router \
  --worker-urls "${WORKER_URLS[@]}" --host 0.0.0.0 --port "$ROUTER_PORT" \
  --policy "$ROUTER_POLICY" &

# Wait for port ready
until nc -z localhost "$ROUTER_PORT"; do sleep 0.5; done
curl -s "http://localhost:${ROUTER_PORT}/healthy" | jq .
          
Compared with the reference article: the reference starts a single multi-node server via two srun and parameters like --tp/--nnodes/--node-rank/--nccl-init-addr. Here we use a router + multi-worker pattern, which is more robust for partial GPUs per node and makes resource isolation easier.

Common pitfalls

  • Port conflicts: if MASTER_PORT is occupied, processes may hang across nodes—choose a free port.
  • Wrong NIC: incorrect NCCL_SOCKET_IFNAME yields poor bandwidth or failures—verify IB device (e.g., ib0).
  • GPU visibility mismatch: ensure HIP_VISIBLE_DEVICES matches Slurm allocation.
  • PyTorch/ROCm mismatch: adjust versions to a compatible pair.
  • Mixed nodes causing perf jitter: prefer homogeneous nodes/topology.

Summary

Key ingredients: precise Slurm allocation (partial), correct ROCm/RCCL envs, proper MASTER/RANK/WORLD_SIZE, and suitable parallel strategy for SGLang. With the templates here, you should achieve stable deployment on most AMD clusters.