分布式训练通信NCCL之Ring-Allreduce详解

🎀个人主页： https://zhangxiaoshu.blog.csdn.net
📢欢迎大家：关注🔍+点赞👍+评论📝+收藏⭐️，如有错误敬请指正!
💕未来很长，值得我们全力奔赴更美好的生活！

前言

随着Chat GPT、文生图、多模态等模型的发展，海量的训练数据、超大规模的模型给深度学习带来了日益严峻的挑战，因此，经常需要使用多加速卡和多节点来并行化训练深度神经网络。目前，数据并行和模型并行作为两种在深度神经网络中常用的并行方式，分别针对不同的适用场景，有时也可将两种并行混合使用。数据并行是在不同设备上放置完整的模型，然后将数据划分在每个设备并行计算。这必然会涉及到不同设备之间的数据传输，即，设备通信，在分布式数据并行的情况下，利用 GPU上的 Minibatch进行权重参数梯度的计算，再利用 GPU之间的通信来完成梯度同步，最后更新模型。常见的深度学习框架提供的通信后端主要有Mpi、Gloo、Nccl，其中Nccl通信后端中的Ring-Allreduce通信技术和硬件的P2P通信技术极大的改善了低效的通信传输问题。本文对Ring-Allreduce通信技术和硬件的P2P通信技术进行详细介绍。

文章目录

前言
一、Ring-Allreduce
- 1. Scatter-Reduce阶段
- 2. Allgather阶段
- 3. Ring-Allreduce通信容量分析
二、硬件Direct通信技术
总结

先对分布式训练中数据并行的流程进行一个回顾， 如下图所示，其基本流程包括将整个训练数据集划分为多个小批次，并将这些小批次分配到不同的设备或节点上。每个设备拥有完整的模型副本，独立处理分配给它的数据。在每个训练步骤中，设备执行前向传播、损失计算、反向传播等操作。随后，梯度信息从各设备中聚合，用于更新模型参数。这一过程循环迭代，直至模型达到收敛或事先定义的训练轮数。数据并行的优势在于有效地利用分布式计算资源，加速大规模深度学习模型的训练，提高训练效率。

在这里插入图片描述

可以发现在分布式训练模型的过程中参数的更新之前需要聚合各设备的梯度信息，因此产生了分布式训练过程中的通信需求，而通信的好坏直接影响到整个模型的训练速度。而Nccl作为常见的深度学习框架提供的通信后端，其中Ring-Allreduce通信技术和硬件的P2P通信技术极大的改善了低效的通信传输问题。

一、Ring-Allreduce

Ring-Allreduce是一种以环状拓扑为基础的通信系统。整个体系结构的工作过程见下图，Rank代表了各个 GPU的进程编号，并且梯度信息可以在两个不同的区域中同步传输。在Ring-Allreduce体系结构中，每台计算机都是一个工作节点，按环形排列。
在这里插入图片描述
Ring-Allreduce体系结构的工作过程被分成两个阶段，即Scatter-Reduce和 Allgather。在Scatter-Reduce阶段，完成了数据的分配与并行，各个工作节点之间的数据交换。最后，在每一个节点上都会有一个最终的结果。Allgather阶段实现了数据的整体同步和压缩，每一个工作节点之间都会进行一些最后的处理，这样对于所有节点来说就可以得到一个完整的结果。

1. Scatter-Reduce阶段

Scatter-Reduce阶段：假定这个阶段的目的是求和，在这个系统中有 N个工作结点，每一个结点中的数据量大小都是K，在Scatter-Reduce的后期，每一个结点都有一个包括初始数组和的而且大小相同的矩阵。

具体的，

第一步，每个结点把本设备上的数据分成 N个区块， N是Ring-Allreduce体系结构中的工作节点数目，见下图步骤（1）。
在第二步，在第一次传输和接收结束之后，在每一个结点上累加了其他节点一个块的数据。这样的数据传输模式直到“Scatter-Reduce”阶段结束，见下图步骤（2）。
最后每一个节点上都有一个包含局部最后结果的区块，由（3）中的深色区块表示，这个区块是所有节点相应的位置区块之和。

可以使用Python对Scatter-Reduce阶段的求和过程进行模拟，代码如下：

import numpy as np

def scatter_reduce(data, num_nodes):
    # 假设data是每个节点上的初始数组
    # num_nodes是工作节点数目
    
    # 第一步：每个节点把本设备上的数据分成N个区块
    local_blocks = np.array_split(data, num_nodes)
    
    # 第二步：在每个节点上累加其他节点一个块的数据
    for i in range(num_nodes):
        other_blocks = [local_blocks[j] for j in range(num_nodes) if j != i]
        local_blocks[i] += np.sum(other_blocks, axis=0)
    
    # 第三步：每个节点上都有一个包含局部最后结果的区块
    final_result = np.sum(local_blocks, axis=0)
    
    return final_result

# 示例
num_nodes = 4
data_size_per_node = 5
total_data_size = num_nodes * data_size_per_node

# 生成随机数据作为每个节点上的初始数组
data = np.random.randint(0, 10, total_data_size)

# 模拟Scatter-Reduce过程
result = scatter_reduce(data, num_nodes)

# 打印结果
print("初始数据：", data)
print("最终结果：", result)

2. Allgather阶段

Allgather阶段：每个工作节点将包含最终结果的块数据块交换, 这样所有的结点就会得到一个完整的结果，

Allgather阶段总共包含有数据发送和接收N一1次，不同的是，Allgather阶段并不需要将接收到的值进行累加，而是直接使用接收到的块内数值去替环原来块中的数值。在迭代完第1次这个过程后，每个节点的最终结果的块变为2个，如图3.3步骤（2）所示。
之后会继续这个迭代过程直到结束，使得每一个节点都包含了全部块数据结果。下图为整个Allgather过程，可以从图中看到所有数据传输过程和中间结果值。

同样可以使用Python对Allgather阶段的过程进行模拟，代码如下：

import numpy as np

def allgather(local_blocks, num_nodes):
    all_blocks = [np.empty_like(local_blocks) for _ in range(num_nodes)]
    
    for i in range(num_nodes):
        # 第一次迭代直接复制本地块到目标块
        all_blocks[i][:] = local_blocks[i]
    
    for _ in range(num_nodes - 1):
        # 迭代过程中交换块数据
        for i in range(num_nodes):
            target_node = (i + 1) % num_nodes
            # 发送当前节点的块到目标节点
            np.copyto(all_blocks[target_node], local_blocks[i])
            # 接收目标节点的块到当前节点
            np.copyto(local_blocks[i], all_blocks[target_node])
    
    return all_blocks

# 示例
num_nodes = 4
data_size_per_node = 5
total_data_size = num_nodes * data_size_per_node

# 生成随机数据作为每个节点的初始数组
local_data = np.random.randint(0, 10, (num_nodes, data_size_per_node))

# 模拟Allgather过程
result_blocks = allgather(local_data, num_nodes)

# 打印结果
print("初始数据块：", local_data)
print("Allgather结果块：", result_blocks)

3. Ring-Allreduce通信容量分析

从上述Ring-Allreduce的算法过程可以看到，

Scatter-Reduce 阶段的工作节点会进行数据的同时收发，具体的，在这一阶段共有 $N - 1$ 次通信容量为 $K / N$ 的数据通信过程。
同样的在Allgather阶段的工作节点也会进行数据的同时收发并且共有 $N - 1$ 次通信容量为 $K / N$ 的数据通信。

故，在使用Ring-Allreduce算法改善后，每个节点传输数据总量变为：

$V co mm u ni c a t i o n = 2 \times K \times N - 1/ N$

由上式可知，当工作节点的数量变得很大时，在Ring-Allreduce架构中单个节点的通信数据量近似为 $2\times K$ ,与节点数$N￥没有关系。这不仅在一定程度上相比传统的BS（参数服务器）通信方式减少了通信量，同时在节点数量增大时，具有很好的可扩展性。

二、硬件Direct通信技术

在常见的分布式训练加速设备中，常常是多节点多加速卡的形式，节点也可被称之为主机或CPU，加速卡的种类很多，常见的有GPU、DCU、FPGA等。如下图所示，在单节点多加速卡的情况下，节点和加速卡以及加速卡之间的数据通信依靠PCIe或NVLink实现，多节点多加速卡的情况下，节点之间的数据通信依靠以太网或Infiniband实现。

在这里插入图片描述
在跨节点加速卡通信的过程中，往往需要先将加速卡的数据传输到相应节点的CPU上，然后CPU通过以太网传输数据，之后又将数据传给加速卡，这种数据在节点和加速卡之间频繁移动所造成的通信开销是很大的，鉴于此，英伟达公司发布了 GPU Direct技术，用于提高加速卡之间通信的效率。

在单节点多加速卡通信中，提出P2P（GPU Direct peer-to-peer）技术。如下图（a）（b）所示，它实现了节点内部加速卡的直接通信，即加速卡可以直接访问另一个加速卡的内存并实现数据的直接传输，避免了加速卡的数据复制到节点CPU内存上作为中转。
在多节点多加速卡通信中，提出了GDR(GPU direct RDMA)技术，如下图（c）所示，加速卡和网卡可以直接通过PCIe进行数据交互，避免了跨节点通信过程中内存和CPU的参与。从而实现加速卡可以直接访问其他节点的加速卡内存。