跳到主要内容

使用Slurm进行大模型预训练

概述​

在slurm 启动大模型预训练任务，本演示为如何使用slurm集群部署预训练任务，并使用预训练模型进行推理。

前置条件​

Get Started​

登录集群​

用以下命令登录客户端节点：

kubectl exec -it slurm-client-8549c7576f-99vg9 -- su - user

其中 slurm-client 后面的字符串为系统自动生成，可用 kubectl get pod 命令查看。

注意

在这里我们切换到了 user 用户，因为工作空间被挂载到了 user 的家目录下。如无特殊说明，以后各节所述的操作均以 user 用户登录客户端节点进行。

检查 GPU​

首先可以在每个计算节点上运行 nvidia-smi 检查 Nvidia GPU 的状态，如：

srun -w slurm-worker-0 nvidia-smi

以上命令在 slurm-worker-0 这个节点上运行了 nvidia-smi.

另外可以用已经安装的 pytorch-cuda 包来检查 GPU 状态。编辑以下 python 代码：

#!/usr/bin/env python3

import os
import torch


print(
    'hostname = ', os.uname().nodename,
    ', torch.cuda.is_available() = ', torch.cuda.is_available(),
    ', torch.cuda.device_count() = ', torch.cuda.device_count()
)

将上述内容保存为文件 tourch_cuda_detect.py 并加上可执行权限，在所有节点上运行：

srun -N3 torch_cuda_detect.py

这里我们有 3 个节点，将输出类似如下的结果：

hostname =  slurm-worker-1 , torch.cuda.is_available() =  True , torch.cuda.device_count() =  1
hostname =  slurm-worker-0 , torch.cuda.is_available() =  True , torch.cuda.device_count() =  1
hostname =  slurm-worker-2 , torch.cuda.is_available() =  True , torch.cuda.device_count() =  1

这说明每个节点上有 1 个 GPU. 需注意的是，即使上述方法确认节点上有 GPU, 但如果其类型及个数与 Slurm 配置中不一致，Slurm 也不会登记使用 GPU. 可以用 sinfo -o "%20N %10c %10m %25f %G" 命令查询每个节点上登记的 GPU.

提交任务​

大模型的并行训练常使用 torch.distributed 框架启动训练脚本。可以将其写成 sbatch 脚本以方便使用，举例如下：

#!/bin/bash

#SBATCH --job-name=TorchDist
#SBATCH --output=%x-%j.out

## Resource Request
#SBATCH --nodes=3
#SBATCH --tasks-per-node=1
#SBATCH --cpus-per-task=2
#SBATCH --mem-per-cpu=64M

export GPUS_PER_NODE=4
export MASTER_PORT=9901
export NCCL_SOCKET_IFNAME=eth0
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ENV,INIT,GRAPH
export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_HCA=ib7s
export UCX_NET_DEVICES=ib7s400p0:1,ib7s400p1:1

## Job Steps
srun sh -c 'python -m torch.distributed.run \
        --nnodes ${SLURM_NNODES} \
        --nproc_per_node ${GPUS_PER_NODE} \
        --node_rank ${SLURM_PROCID} \
        --master_addr $(scontrol show hostnames | head -n 1) \
        --master_port ${MASTER_PORT} \
        training_script.py'

其中：

  • training_script.py 是用户用 Python 实现的训练脚本
  • 以 NCCL 开头的环境变量用于设置 NCCL，详情见 NCCL 文档 “Environment Variables”
  • 以 UCX 开头的环境变量用于设置 OpenUCX, 详情见 UCX 文档 “Running UCX”

对这段脚本的说明如下：

  • torch.distributed.run 本身可以在单个节点上启动多个进程，由 --nproc_per_node 参数指定。一般来说，为每个 GPU 启动一个进程即可
  • 有多个节点时，每个节点上需单独运行 torch.distributed.run. 这里用 Slurm 帮我们实现多节点并行启动
  • Slurm 只需要在每个节点上启动一个任务即可，也就是 --tasks-per-node 参数设为 1
  • Slurm 通过环境变量 SLURM_NNODES 给出作业包含的节点数，可以用作 torch.distributed.run 的 --nnodes 参数
  • 对于每一个任务，Slurm 赋与了一个唯一 ID, 由环境变量 SLURM_PROCID 给出，可以把它用作 torch.distributed.run 的 --node_rank 参数
  • torch.distributed.run 把 Rank 为 0 的节点作为 master, master 在 --master_port 指定的端口上监听
  • torch.distributed.run 把 Rank 不为 0 的节点作为 client, client 主动去连接 --master_addr 指定的主机和 --master_port 指定的端口

这里的一个要点是如何确定 --master_addr 参数。

  • Slurm 通过环境变量 SLURM_JOB_NODELIST 给出作业包含的节点的主机名列表，注意这是一个可能包含简写的字符串
  • scontrol show hostnames 命令默认解析环境变量 SLURM_JOB_NODELIST，返回按行分隔的主机名。这个命令只解析语法，并不检查主机是否存在
  • Slurm 从 0 开始为各个任务分配 SLURM_PROCID
  • SLURM_PROCID 为 0 的任务所在的节点总是在 SLURM_JOB_NODELIST 的第一位（实践观察，未见文档说明）

因此可以用脚本中的方法取 --master_addr 参数。