Multi-node slurm training?

[JulioZhao97]

Hello guys, I wonder if there is an example for multi-node training on slurm? The multi-node training example you provided is training on multiple machines, which don't apply on slurm where each the communication of each node is managed by slurm manager.

[yogkul2000]

@JulioZhao97 Try this.

#!/bin/bash

#SBATCH -N 2 #num nodes
#SBATCH --ntasks-per-node=1 #one task per node for torchrun
#SBATCH --gpus-per-node=4  
#SBATCH -c 8 #cpus per task
#SBATCH -t 7-00:00:00
#SBATCH -p general
#SBATCH -q public
#SBATCH --mem=256G #per node
#SBATCH -o slurm_logs/slurm_rlhf.%j.out
#SBATCH -e slurm_logs/slurm_rlhf.%j.err

module load mamba/latest
module load cuda-12.4.1-gcc-12.1.0
module load gcc-12.1.0-gcc-11.2.0
module load ffmpeg-6.0-gcc-12.1.0

source activate swift

export MASTER_PORT=$(( 29500 + RANDOM % 1000 ))
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=ib0

export NNODES=2
export NPROC_PER_NODE=4

export MAX_PIXELS=401408
export ENABLE_AUDIO_OUTPUT=0 
export USE_AUDIO_IN_VIDEO=1 
export VIDEO_MAX_PIXELS=50176 
export FPS_MAX_FRAMES=16

echo "========== JOB CONFIGURATION =========="
echo "SLURM_JOB_ID: $SLURM_JOB_ID"
echo "SLURM_NODELIST: $SLURM_NODELIST"
echo "MASTER_ADDR: $MASTER_ADDR"
echo "MASTER_PORT: $MASTER_PORT"
echo "NNODES: $NNODES"
echo "NPROC_PER_NODE: $NPROC_PER_NODE"
echo "TOTAL_GPUS: $(($NNODES * $NPROC_PER_NODE))"
echo "======================================"

cat > swift_rlhf_launcher.sh << 'EOL'
#!/bin/bash

NODE_RANK=$SLURM_PROCID

echo "Starting training on node $NODE_RANK ($(hostname))"

export CUDA_VISIBLE_DEVICES=0,1,2,3

GRADIENT_ACCUM_STEPS=1  # Adjust as needed

torchrun \
    --master_port $MASTER_PORT \
    --master_addr $MASTER_ADDR \
    --nproc_per_node=$NPROC_PER_NODE \
    --nnodes=$NNODES \
    --node_rank=$NODE_RANK \
    -m swift.cli.rlhf \
    --rlhf_type grpo \
    --model Qwen/Qwen2.5-Omni-7B \
    --reward_funcs custom_accuracy_reward custom_format_reward \
    --reward_weights 1 0.5 \
    --train_type lora \
    --lora_rank 8 \
    --lora_alpha 32 \
    --target_modules all-linear \
    --torch_dtype bfloat16 \
    --dataset '/home/video_qa_dataset_grpo.jsonl' \
    --external_plugins examples/train/grpo/plugin/plugin.py \
    --max_completion_length 768 \
    --num_train_epochs 1 \
    --per_device_train_batch_size 6 \
    --per_device_eval_batch_size 6 \
    --learning_rate 5e-7 \
    --gradient_accumulation_steps $GRADIENT_ACCUM_STEPS \
    --eval_steps 200000000 \
    --save_steps 100 \
    --save_total_limit 2000000 \
    --logging_steps 1 \
    --max_length 4096 \
    --output_dir output_stage_1_math_grpo_multinode \
    --warmup_ratio 0.05 \
    --dataloader_num_workers 4 \
    --dataset_num_proc 4 \
    --num_generations 6 \
    --top_p 0.99 \
    --top_k 50 \
    --temperature 1. \
    --dynamic_sample true \
    --system 'examples/train/grpo/prompt.txt' \
    --deepspeed zero2 \
    --log_completions true
EOL

chmod +x swift_rlhf_launcher.sh

srun --ntasks=$SLURM_NTASKS --ntasks-per-node=1 ./swift_rlhf_launcher.sh

EXIT_CODE=$?
if [ $EXIT_CODE -ne 0 ]; then
    echo "Training failed with exit code $EXIT_CODE"
else
    echo "Training completed successfully"
fi

# Cleanup
rm -f swift_rlhf_launcher.sh