最近使用了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多机训练模型的方法,不对的地方欢迎指正!