hfreduce | 高性能的多卡并行通信工具

High-Flyer March 04, 2020

hfreduce 是幻方AI自研的高性能的多卡并行通信工具，是幻方“萤火二号”计算存储分离后，计算服务中的重要一环，根据萤火二号的集群特性专为计算节点设计的 allreduce 工具。本质上 hfreduce 相当于 PyTorch 中的 DistributedDataParallel（DDP），只不过使用 CPU 做加法运算以计算总梯度，而不是调用其他的集体通信库（CCL），比如说 NCCL，传递梯度到不同的显卡上，再各自计算总梯度。

为什么要这么设计呢？本期文章将为大家解密，详细介绍幻方AI设计 hfreduce 的背景与理念。文章会阐明 hfreduce 对模型训练整体加速的效果以及具体适用的场景。

背景

幻方AI之所以采用 CPU 来做加法运算，计算总梯度，是因为萤火二号上的计算节点配置比较特殊。

幻方的主要AI场景是金融行为分析、自然语言处理、生物分子结构预测等。在这些场景中，基本是数据规模大而模型大小适中。换句话说，在A100显卡40G的显存中，完全可以装得下一个完整的模型和批次样本数据。因此，模型的加速主要是依赖大量的数据并行，让尽可能多的显卡参与训练，再同步梯度。

正是因为上述数据并行的应用场景，随着集群规模不断变大，我们发现使用 NCCL 进行多 GPU 之间的集体通信，梯度信息多次在 PCIe 上遍历就会成为一个性能瓶颈。此外，我们无法使用在交换机上做加法运算计算梯度的方法，因为萤火二号集群中会同时运行大量的训练任务，它们各自需要做reduce以计算总梯度，而交换机之前并不支持在同一个网络内进行如此多路的reduce计算。

框架设计

近些年深度学习模型逐渐尝试更大的参数空间来处理大规模的数据，使用多个GPU一同加速训练成为很多AI从业者的必然选择。NCCL是英伟达发布的用于实现多GPU之间集体通信的工具，在PCIe，NVLink上可以实现较高的通信速度。NCCL借助CPU，可以传播，计算不同显卡上的梯度信息，对并行训练中的多个显卡进行梯度信息的同步。在有NVLink的条件下，在同节点不同显卡内做reduce计算对 NCCL 来说轻而易举，但如果没有NVLink，梯度信息多次在 PCIe 上遍历就会成为一个性能瓶颈。

既然我们是因为没有 NVLink 而放弃了 NCCL，那么把梯度数据都先拷贝到主内存里然后在 CPU 上做加法运算计算总梯度也就是很自然的事了，因为这样可以避免梯度数据在多个 GPU 之间进行多次拷贝。hfreduce 首先从所有并行的GPU里一次性收集梯度信息，然后在CPU里进行reduce计算，完成后将计算出来的梯度一次性广播到所有 GPU，从而将 PCIe 流量减半。此外，CPU 只用于将所有梯度信息相加汇总，与 NCCL 不同，免去了在 GPU 上再计算梯度的开销。

如上图所示，与 NCCL 一样，我们也采用两棵树的拓扑结构进行节点间的通信。对于节点比较多的情况，它能比环形拓扑更好，因为节点之间的对数距离会更短。我们设计让树结构更加平衡一些，所以理论上它能比 NCCL 的树结构稍微快一些。

在CPU 上的加法运算也进行了精心设计。所有加载到内存中的梯度信息都在 NUMA 节点之间进行交互，绑定到 CPU 内核的线程从同一个 NUMA 节点中挑选出梯度数据并将它们相加。 AUX 256 指令优于 AUX 512 指令，因为 AMD CPU 不支持后者，如果所有内核都在执行指令，则会导致 CPU 频率受限。

由于所有的加法运算都是在 CPU 内核上执行的，我们观察到内存带宽是瓶颈，而不是CPU的算力。因此，建议使用更大的样本批次，因为读取它们可能会减少内存带宽的竞争。 PCIe 的宽松排序可能会提高Reduce的性能，但实际上更快的Reduce可能会减慢样本读取的速度，从而减慢整个训练的速度。 Reduce的性能很重要，但就目前来说，该方案仍然只是一个局部最优解。

使用案例

我们给hfreduce设计了Python的接口，使其可以轻松地加入 PyTorch 的训练中。该库能自动处理节点地址交换、梯度归零和梯度的最终平均。更重要的是，该库能在模型代码中直接嵌入，以便在计算梯度后立即开始Reduce计算。 hfreduce的使用基本上就像创建一个 Reducer 实例一样简单。

具体的，在萤火二号上训练模型时，可以参考如下案例代码：

import hfreduce.torch as hfr

model = Model().cuda()
loss = Loss(model)
optimizer = Optimizer(model)

reducer = hfr.AsyncReduceFloat(registry_host, port, proc_rank, procs_per_node, node_rank, node_cnt, model)

for sample in sampels:
    reducer.zero_grad() # instead of optimizer.zero_grad(), call reducer.zero_grad()
    
    model(sample)
    loss.backward()
    
    reducer.synchronize()
    
    optimizer.step()

性能测试

我们使用 PyTorch 内嵌的 VGG16 模型进行实验，测试hfreduce的性能。

如上图所示，hfreduce整体耗时更少，明显比没有NVLink的NCCL更快。使用 PyTorch 进行模型训练，即使在后台使用 PyTorch DDP 调用 NCCL 进行 allreduce，hfreduce 仍然胜过 NCCL。