使用NCCL作为PyTorch多节点分布式训练后端时候的经验总结
闲言
最近开始做 LLM/GPT 了,单机多卡训练只能勉强满足需求,能够利用上多机多卡,配合上 InifinBand,会很有大的提升的.这中间的测试来来回回也踩了不少坑,这里做一些简单的记录.
对多机多卡训练的极简介绍
这里简单介绍一下如何启动多机多卡训练.
使用 torchrun
这个是 PyTorch 推荐使用的新方式,支持动态增减节点,自动给节点分配 rank 等.比传统的方式更加灵活且简洁.直接使用 torchrun
命令等价于 python -m torch.distributed.run
.
torchrun \
--nnodes $SLURM_NNODES \
--nproc_per_node $SLURM_GPUS_PER_NODE \
--rdzv_backend c10d \
--rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \
--rdzv_id $SLURM_JOB_ID \
train.py args...
参数解释:
--nnodes
:并行的机器数量.--nproc-per-node
:每个节点上的任务数量,一个 GPU 一个任务,实际上也就是每个节点上的 GPU 数量.--rdzv_backend
:rdzv 后端,一般用 c10d.--rdzv_endpoint
:master 节点的 rdzv 服务地址.分布式训练的每个计算节点会通过这个地址进行通信,并且根据rdzv_id
来标识他们属于同一组作业.--rdzv_id
:rdzv_id
用于标识一个唯一的分布式训练,你可以在相同的机器上启动多个分布式训练实验,只要他们的 ID 不同.
accelerate
accelerate launch \
--dynamo_backend no \
--machine_rank $SLURM_NODEID \
--main_process_ip $MASTER_ADDR \
--main_process_port $MASTER_PORT \
--mixed_precision no \
--multi_gpu \
--num_machines $SLURM_NNODES \
--num_processes $num_processes \
--rdzv_backend c10d \
train.py args...
参数解释:
dynamo_backend
:我们不需要这个,指定为no
只是不想看到 warning.--machine_rank
:节点的 rank,一般设置为环境变量SLURM_NODEID
即可.--main_process_ip
和--main_process_port
:主节点的地址和端口,组成IP:PORT
作为rdzv_endpoint
.--mixed_precision
:混合精度,根据需要设置即可.--multi_gpu
:启用多卡并行训练.--num_machines
:并行训练的节点数量.--num_processes
:accelerate 这里要的是总的任务数,每个 GPU 算一个,也就是总的 GPU 数量.所以这里对应到 torchrun 的方式,应该是SLURM_GPUS_PER_NODE x SLURM_NNODES
--rdzv_backend
:这个与前面 torchrun 命令的选项一样.经过测试,仅设置为c10d
是成功的,且不要设置--same_network
.--config_file
: accelerate 的选项可以通过accelerate config
交互式生成,保存到配置文件中,使用的时候此选项指定.但是命令行选项应该优先级更高.你可以把通用的部分写到配置文件,而其他的在命令行指定.
多节点启动
如果不使用别的工具,一般我们就需要分别在每一台机器上执行上面的命令.但是大多数时候我们都是在集群环境下使用的,这个时候可以使用集群工具来自动帮我们完成这个操作.比较常用工具的是 slurm,也有部分集群使用 PBS,两者区别不大.
一个快速简单的例子
以 PyTorch 官方提供了一个 minGPT-ddp 例子,这个例子足够小,可以在 RTX 2080 显卡上跑起来.完整的源码在 pytorch-examples.将其 clone 下来之后,修改 distributed/minGPT-ddp/mingpt/slurm/job.sh
文件里的设置,然后使用 sbatch job.sh
命令提交即可.这里只需要根据具体使用的 slurm 集群的情况,修改其中的计费账户 account 和使用的分区 partition,应该就可以了.其他的脚本也可以参考此模板来创建 slurm sbatch 脚本.
我在一个 slurm 集群上的 2 个 A800 节点进行了测试,完整的测试命令见 run.sh,输出日志见 logs 下的文件.
即使不使用 slurm 集群,上面的脚本本身也是有效的 bash 脚本,简单修改一下,也可以直接执行:
- 修改
SLURM
开头的环境变量,指定为合适的值. - 删掉
srun --label
,即直接执行对应的命令,而无需通过 srun 启动.
补充
就个人的实际测试而言,这个并行训练也不是每次都成功,可能会遇到训练无法启动,就是卡住了.我不确定这个问题是不是我申请资源的时候每个节点值申请了 1 个 GPU,节点上还有其他用户的任务,导致影响到了我的训练,还是说这个集群的网络配置存在一些什么问题.不过好在单机多卡训练一般都是没有问题的,且上面提到的 slurm sbatch 脚本也都可以正常使用的,仅需要指定申请的节点数量为 1 即可.
关于 InfiniBand,RDMA,RoCE
我们这里不对这些概念做特别详细深入的介绍.只是单纯为了对比哪个更快更好,我们应该优先选哪个而已.
- RDMA (remote direct memory access),也就是远程直接内存访问,是一种绕过远程主机操作系统内核访问其内存中数据的技术.由于不经过 CPU、CPU 缓存或者上下文切换的参与,速度很快.
- InfiniBand 即指他的物理设备(包括网卡、交换机),也指一种协议.这里我们更多的关心他的物理层面.一般地,计算机集群的计算节点会配备 IB 网卡,但是由于机房不一定.IB 提供了一种 RDMA 的实现.
- 由于 IB 需要特定的交换机,更换就需要增加成本.一种解决方式是 RoCE (RDMA over Converged Ethernet).简单的理解就是结合了 IB,但是用是以太网交换机,实现了 RDMA.
从速度和效率上来说:
- 原生的 IB 是最优先考的,但是它需要 IB 交换机,多了一笔费用,部分机房可能没有配备.
- RoCE 是其次的,他只需要用普通的以太网交换机.
- 最次的就是完全没有 IB 设备,直接使用以太网网卡.
一般来说,NCCL 会自动选择合适的设备,但是如果出错,我们需要了解一下如何查看相关信息.
查看节点上 IB 设备
我们可以使用 ibstatus
命令查看 IB 设备的状态.例如:
Infiniband device 'mlx5_0' port 1 status:
default gid: fe80:0000:0000:0000:0ac0:ebff:fefe:428c
base lid: 0x0
sm lid: 0x0
state: 4: ACTIVE
phys state: 5: LinkUp
rate: 100 Gb/sec (4X EDR)
link_layer: Ethernet
Infiniband device 'mlx5_1' port 1 status:
default gid: fe80:0000:0000:0000:946d:ae03:00fc:e51c
base lid: 0x9
sm lid: 0x6
state: 4: ACTIVE
phys state: 5: LinkUp
rate: 200 Gb/sec (4X HDR)
link_layer: InfiniBand
Infiniband device 'mlx5_2' port 1 status:
default gid: fe80:0000:0000:0000:0ac0:ebff:fefe:547c
base lid: 0x0
sm lid: 0x0
state: 1: DOWN
phys state: 3: Disabled
rate: 100 Gb/sec (4X EDR)
link_layer: Ethernet
Infiniband device 'mlx5_3' port 1 status:
default gid: fe80:0000:0000:0000:946d:ae03:00fc:e548
base lid: 0x8
sm lid: 0x6
state: 4: ACTIVE
phys state: 5: LinkUp
rate: 200 Gb/sec (4X HDR)
link_layer: InfiniBand
这里一般只关注状态为 ACTIVE
的设备.可以看到 mlx5_0
的 link_layer
为 Ethernet
,而 mlx5_1
的 link_layer
为 InfiniBand
,且 mlx5_1
的速率更高.这里的设备可以用于设置环境变量 NCCL_IB_HCA
.
简单的理解:
link_layer
为Ethernet
,说明这个 IB 设备的链路层协议为 Ethernet,一般就是使用了 RoCE 协议.link_layer
为InfiniBand
,说明这个 IB 设备的链路层协议为 InfiniBand,也就是原生的 InfiniBand.
再看 ibdev2netdev
的输出:
mlx5_0 port 1 ==> ens13np0 (Up)
mlx5_1 port 1 ==> ib0 (Up)
mlx5_2 port 1 ==> ens17np0 (Down)
mlx5_3 port 1 ==> ib1 (Up)
可以看到这几个 IB 设备映射到了网络设备.
使用 NCCL 作为 PyTorch 分布式训练后端的时候,IB 设备的使用情况分析
下面我们测试一下各个情况的组合.这里我们确认了要申请的 2 个节点的 IB 设备名字都是前面输出的内容,也就是两个节点一致.我们测试的组合包括:
NCCL_IB_DISABLE
设置为 0 或者 1NCCL_IB_HCA
设置为不同的设备.NCCL_SOCKET_IFNAME
设置为不同的设备.
正如上文提到的,本次测试使用的 slurm sbatch 脚本见 job.sh,实际执行命令见 run.sh,输出日志见 slurm 下的文件.
测试结果
- 当
NCCL_IB_DISABLE=0
的时候,NCCL_IB_HCA
设置的值如果不是rdma link
显示的 IB 设备,则 NCCL 会提示找不到 IB 设备,然后回落到 NET/Socket,识别到可用的网络设备,并且实际使用的是 ib0(见日志中的[send] via NET/Socket/1
和NCCL INFO Using
给出的设备列表). - 当
NCCL_IB_DISABLE=0
的时候,NCCL_IB_HCA
设置的值是rdma link
显示的 IB 设备,则 NCCL 会按照环境变量的设置去使用该设备,但是如果该设备状态为 Down,则还是会自动回落到 NET/Socket. - 当
NCCL_IB_DISABLE=0
,其他的环境变量不指定的时候,NCCL 自动识别到可用的 IB 设备,并自动选择使用了速率最高的原生 IB. - 当
NCCL_IB_DISABLE=1
的时候,NCCL 将使用NCCL_SOCKET_IFNAME
指定的设备,如果这设备不可用,则会直接报错. - 当
NCCL_IB_DISABLE=1
,其他的环境变量不指定的时候,NCCL 自动识别到可用的设备 ib0 和 ib1,并自动选择了 ib0.
关于 NCCL 环境变量的小结
我们来小结一下 NCCL 环境变量的设置.
- 需要查看一些详细的调试信息的话,可以设置
NCCL_DEBUG=INFO
. - 确认需要使用 IB 设备和 IB 协议的话,首先要确认
NCCL_IB_DISABLE
不能设置为1
,然后设置NCCL_IB_HCA
为rdma link
输出的的设备名字,且优先选择ibstatus
里显示的 link layer 为 InfiniBand 的设备,其次是 link layer 为 Ethernet 的设备. - 如果不用 IB 设备,则需要设置
NCCL_IB_DISABLE=1
,然后设置NCCL_SOCKET_IFNAME
为ibdev2netdev
显示的 IB 映射到的以太网设备名字. - 最后,如果根本没有 IB,则只能设置
NCCL_IB_DISABLE=1
,并设置NCCL_SOCKET_IFNAME
为ip link show
显示的以太网设备.
使用 LLaMa Factory 训练 LLM
这里其实没有什么好说的了,直接参看上文提到的 slurm sbatch 脚本即可,只需修改对应的 python 脚本和对应的选项参数即可.
参考
以下是本文的一些重要参考资料来源,有兴趣的可以看看.
- sbatch doc:需要用到的 SLURM 相关的环境变量这里都有.
- NCCL doc:NCCL 调试相关的环境变量可以参考这里.
- Fine-tuning Llama 2 70B using PyTorch FSDP:这里提供了一个使用 accelerate 训练的参考,他的 slurm sbatch 脚本也可以参考.
- 如何理解 Nvidia 英伟达的 Multi-GPU 多卡通信框架 NCCL:这个知乎回答也值得看看.