Untitled

#!/bin/bash
#SBATCH --job-name=qwen-multinode
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=1  # Using 1 GPU per node for a 0.6B model
#SBATCH --cpus-per-task=8
#SBATCH --partition=gpu     # Change to your cluster's GPU partition

# 1. Identify the Head Node IP
# We need to tell the workers where to find the Ray head.
head_node=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

echo "Head Node: $head_node ($head_node_ip)"

# 2. Start the Ray Head Node
srun --nodes=1 --ntasks=1 -w "$head_node" \
    ray start --head --node-ip-address="$head_node_ip" --port=6379 \
    --num-gpus=1 --block &

sleep 10

# 3. Start the Ray Worker Node
# This srun targets all nodes EXCEPT the head node.
srun --nodes=1 --ntasks=1 --exclude="$head_node" \
    ray start --address="$head_node_ip:6379" \
    --num-gpus=1 --block &

sleep 20

# 4. Launch vLLM
# For a 0.6B model, TP=1 and PP=2 (splits the layers 50/50 across nodes)
srun --nodes=1 --ntasks=1 -w "$head_node" \
    vllm serve Qwen/Qwen3-0.6B \
    --tensor-parallel-size 1 \
    --pipeline-parallel-size 2 \
    --port 8000 \
    --trust-remote-code