使用DeepSpeed在Slurm集群上多机跑大模型

最近使用了DeepSpeed跑Bloom模型，在多机训练时遇到了一点问题，查了一些资料没有发现DeepSpeed在Slurm集群上多机跑的方法笔记，特此记录一下。

我用的是DeepSpeed-Chat的Step 1 (SFT) 的训练code，官方给出的训练示例可以在这里找到。

在Slurm集群中，单机多卡的训练脚本如下：

#!/bin/bash
#SBATCH -J FT-BLOOM  ##提交的任务名称
#SBATCH -o ../log/ft.bloom.snode.log
#SBATCH -p GPU_NAME # 任务提交的分区名称
#SBATCH -N 1   # 任务申请1个节点
#SBATCH --cpus-per-task=16  ## 每个任务使用的cpu核数
#SBATCH --gres=gpu:8   ## 每个节点的gpu数
#SBATCH --mem 256G    ## 每个节点的memory 
###SBATCH -w GPU_NAME09 # 指定运行作业的节点，若不填写系统自动分配节点

OUTPUT=/path/to/output
ZERO_STAGE=3
export TOKENIZERS_PARALLELISM=true 
code_path=/path/to/DeepSpeed/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning
deepspeed $code_path/main.py \
   --data_path knowledge_data \
   --data_split 10,0,0 \
   --sft_only_data_path knowledge_data \
   --data_output_path /path/to/output \
   --model_name_or_path /path/to/bloom-3b \
   --per_device_train_batch_size 1 \
   --per_device_eval_batch_size 1 \
   --max_seq_len 1024 \
   --learning_rate 1e-5 \
   --weight_decay 0.1 \
   --num_train_epochs 5 \
   --gradient_accumulation_steps 1 \
   --lr_scheduler_type cosine \
   --num_warmup_steps 0 \
   --seed 1234 \
   --zero_stage $ZERO_STAGE \
   --deepspeed \
   --output_dir $OUTPUT \
   &> $OUTPUT/training.snode.log

如果不使用slurm集群，DeepSpeed多机跑模型只需要在直接在 deepspeed 后面加一个参数 --hostfile=/path/to/hostfile，其中 hostfile 文件中存放的是节点名称及每个节点使用的gpu数：

...
deepspeed --hostfile=/path/to/hostfile \
$code_path/main.py \
...

/path/to/hostfile中的内容为：

g0003 slots=8
g0004 slots=8

其中每一行代表一个节点，第一个是节点名称，第二个是该节点使用的gpu数。

上述方法只适用于每个节点均允许无密钥ssh的情况，在slurm集群中并不能直接使用，slurm集群中使用多节点的脚本如下：

#!/bin/bash
#SBATCH -J FT-BLOOM
#SBATCH -o ../log/ft.bloom.mnodes.log
#SBATCH -p GPU_NAME
#SBATCH -N 2   # 作业申请 2 个节点
#SBATCH --gres=gpu:8  ## 每个节点使用8个GPU
#SBATCH --cpus-per-task=8 ## 每个任务使用8个cpu核
#SBATCH --ntasks-per-node=8  ## 每个节点上8个任务
#SBATCH --gpus-per-task=1  ## 每个GPU上1个任务
#SBATCH --mem 256G  ## 每个节点使用的内存
#SBATCH -w GPU_NAME[02-03] # 指定运行作业的节点，若不填写系统自动分配节点

OUTPUT=/path/to/output
export GPUS_PER_NODE=8

TASK_NAME="task_name"
export CKPT="$OUTPUT/$TASK_NAME"
mkdir -p $OUTPUT $CKPT 

echo "SLURM_NTASKS=$SLURM_NTASKS"
export HOSTFILE="./hostfile"
echo "[slurm node name] slots=$SLURM_NTASKS" > $HOSTFILE

### 定义节点到ip的映射
declare -A IB0
for i in $(seq 0 node_count); do
    key=$(printf "NODE_NAME%02d" "$i")
    value=$(printf "xx.xx.xx.%02d" $((i)))
    IB0["$key"]="$value"
done
### 这一步定义你的slurm集群中所有的节点与ip的对应关系，用一个dict存储

### 系统自动匹配当前进程的主节点
regex="[0-9]+"
if [[ $SLURM_JOB_NODELIST =~ $regex ]]; then
    extracted_number="${BASH_REMATCH[0]}"
    echo "Extracted number: $extracted_number"
else
    echo "No match found"
fi

MASTER_KEY="NODE_NAME$extracted_number"

### 定义master的地址和端口
export MASTER_ADDR=${IB0[$MASTER_KEY]}
export MASTER_PORT=31217

code_path=/path/to/DeepSpeed/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning
deepspeed --num_nodes $SLURM_NNODES \
   --num_gpus $GPUS_PER_NODE \
   --master_addr $MASTER_ADDR \
   --master_port $MASTER_PORT \
   --hostfile $HOSTFILE \
   --no_ssh_check \
   --launcher SLURM \
   --force_multi \
   $code_path/main.py \
   --data_path knowledge_data \
   --data_split 10,0,0 \
   --sft_only_data_path knowledge_data \
   --data_output_path /path/to/output \
   --model_name_or_path /path/to/bloom-1b7 \
   --per_device_train_batch_size 1 \
   --per_device_eval_batch_size 1 \
   --max_seq_len 2048 \
   --learning_rate 1e-5 \
   --weight_decay 0.1 \
   --num_train_epochs 10 \
   --gradient_accumulation_steps 16 \
   --lr_scheduler_type cosine \
   --num_warmup_steps 0 \
   --seed 1234 \
   --zero_stage $ZERO_STAGE \
   --deepspeed \
   --output_dir $CKPT \
   --gradient_checkpointing \
   &> $OUTPUT/training.log

上面的脚本就是DeepSpeed在Slurm集群中多机运行的脚本，但是直接运行脚本会报错，local_rank无法通过args自动传参，导致没有分布式初始化，需要对/path/to/DeepSpeed/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning 的 main.py 做一下修改，把环境变量传参进去，在 main.py 中添加：

### 添加环境变量
world_size = os.environ['SLURM_NTASKS']  
node_id = os.environ['SLURM_NODEID']   
rank = os.environ['SLURM_PROCID']   
local_rank = int(os.environ['SLURM_LOCALID'])

os.environ['RANK'] = os.environ['SLURM_PROCID']
os.environ['WORLD_SIZE'] = os.environ['SLURM_NTASKS']
os.environ['MASTER_PORT'] = os.environ['MASTER_PORT']
os.environ['LOCAL_RANK'] = os.environ['SLURM_LOCALID']

### main函数中添加：
args.local_rank = local_rank

其中Slurm环境变量含义如下：

SLURM_NTASKS  ## 总任务数量，可用作world_size
SLURM_NODEID  ## 节点下标，即node_id
SLURM_PROCID  ## 全局进程id，可用作global rank
SLURM_LOCALID ## 局部进程id，可用作local_rank

以上就是在Slurm集群中使用DeepSpeed Chat多机训练模型的方法，不对的地方欢迎指正！