Notes about AI分布式训练通信原语

本文将mark下AI分布式训练通信原语的相关notes。内容主要转载自GPU售前知识—AI分布式训练通信原语

在大型AI模型训练中，单机单GPU的配置已无法满足算力与内存的需求。随着模型参数规模的不断扩大，分布式训练已成为深度学习发展的必由之路。而支撑分布式训练最基础的能力正是分布式训练通信原语（Communication Primitives）。

分布式训练通信原语的概念

“通信原语”并非指某个具体的框架或API，而是指各种分布式系统中高度抽象、反复出现的基础通信与协同操作。它们定义了在训练过程中多设备、多进程之间如何交换数据、同步状态、维持一致性。无论你使用PyTorch、TensorFlow还是JAX，也无论采用数据并行、模型并行还是混合并行策略，底层实现都依赖这些AI分布式训练通信原语的高效实现。可以说，通信原语的性能直接决定了分布式训练效率的上限。

常见分布式训练通信原语详解

基于实际训练中最常见的通信模式，我们在此介绍八种核心AI训练通信原语：广播Broadcast、分散Scatter、聚合Gather、全聚合AllGather、规约Reduce、规约分散ReduceScatter、全规约AllReduce以及全交换All-to-All。

广播Broadcast

广播Broadcast是一对多通信原语，即一个发送者对应多个接收者。在集群中，某个节点（或GPU）将其数据广播至所有其他节点。例如将GPU0的数据同步至GPU1、GPU2与GPU3。

广播Broadcast是一对多数据同步机制，能将一个GPU的数据同步至所有其他GPU。其应用场景包括：

并行数据参数初始化，确保所有GPU的初始参数一致性；
AllReduce组合（广播Broadcast+规约Reduce）中的广播操作；
分布式训练参数服务器架构中，主节点向工作节点广播Broadcast数据，再由工作节点将数据规约Reduce回主节点的广播环节。

分散Scatter

分散Scatter与广播Broadcast虽同属一对多通信原语，但存在本质区别。分散与广播的核心差异在于：

广播Broadcast：向所有节点发送相同的完整数据。
分散Scatter：将数据拆分后分别发送给不同节点。

例如，将一份数据DATA拆分为DATA-A、DATA-B、DATA-C与DATA-D：

DATA-A → GPU0
DATA-B → GPU1
DATA-C → GPU2
DATA-D → GPU3

分散Scatter的逆操作是聚合Gather，其典型应用场景包括：

规约分散ReduceScatter组合中的分散Scatter操作；
模型并行中，初始化时模型参数分散至不同GPU的环节。

聚合Gather

聚合Gather是多对一通信原语，即多个数据发送者对应一个数据接收者。它将集群中多个节点的数据汇集到单个节点上。各节点将自身数据发送至主节点，聚合完成后主节点即拥有所有节点的数据。例如：

GPU0 → DATA-A
GPU1 → DATA-B
GPU2 → DATA-C
GPU3 → DATA-D
最终主节点将获得完整的DATA-A/B/C/D数据集。

聚合Gather是多对一数据收集方法，将多张GPU卡的数据收集到单张GPU卡上。其逆操作对应分散Scatter，应用场景即规约分散ReduceScatter组合中的分散Scatter环节。

全聚合AllGather

全聚合AllGather是多对多通信原语，实现多张GPU卡之间的数据同步全收集。可将其视作聚合与广播操作的组合，即：

首先，将各节点数据收集至主节点（Gather）；
随后，将完整数据广播至所有节点（Broadcast）。
最终Allgather的结果是每个节点都拥有完整数据。

全聚合AllGather的逆操作是规约分散ReduceScatter，其典型应用场景包括：

并行模型训练场景；
在并行模型训练中，前向计算所需的参数需要全量同步，此时需要通过AllGather将拆分在不同GPU上的参数同步到单张GPU上才能进行前向计算。

规约Reduce

规约（Reduce）是多对一通信原语，用于将多个节点的数据进行归约Reduce计算并将结果汇总至单个节点。常见的规约Reduce操作包括：

求和（SUM）
求积（PROD）
最大值/最小值（MAX/MIN）
逻辑与/或（LAND/LOR）
按位与/或/异或（BAND/BOR/BXOR）
最大值位置/最小值位置（MAXLOC/MINLOC）

这些规约Reduce操作通常需要GPU或加速卡硬件支持对应的算子（operator）才能实现高性能。
规约Reduce的典型应用场景包括：

全规约AllReduce中的规约Reduce环节；
规约分散ReduceScatter组合通信中的规约Reduce操作；
分布式训练参数服务器架构中，Master节点先将模型参数广播Broadcast至所有Worker节点，各Worker完成计算后通过规约Reduce操作将结果聚合回Master节点。

规约分散ReduceScatter

规约分散ReduceScatter是多对多通信原语，具有多个数据发送者与多个数据接收者。它在集群所有节点上沿相同维度执行规约Reduce操作，随后将结果分发至集群每个节点。规约分散ReduceScatter等效于在每个节点上执行一次规约Reduce操作后，再进行一次分散Scatter分发操作。其逆操作是全聚合AllGather。

规约分散ReduceScatter是多对多数据操作，采用”先规约Reduce后分发”的方式。它首先对各GPU卡上的数据进行规约（例如求和），再将数据分发（分散Scatter）处理。其应用场景包括：

规约分散ReduceScatter可同时适用于数据并行与模型并行场景：
数据并行中规约分散ReduceScatter+全聚合AllGather组合内的规约分散ReduceScatter操作；
模型并行中前向AllGather后反向计算中的规约分散ReduceScatter环节。

全规约AllReduce

全规约AllReduce是多对多通信原语，具有多个数据发送者与多个数据接收者。它在集群所有节点上执行相同的规约Reduce操作，并将各节点的规约结果发送至每个节点。全规约AllReduce可通过在主节点执行规约Reduce+广播Broadcast或通过规约分散ReduceScatter+全聚合AllGather组合实现。

全规约AllReduce是多对多数据规约操作，将集群中所有GPU的数据规约（如求和）后分发至各GPU。其应用场景包括：

AllReduce可适用于数据并行场景；
各类通信拓扑中的数据并行场景，例如环式全规约（Ring AllReduce）与树状全规约（Tree AllReduce）中的全规约AllReduce实现。

全交换All-to-All

全交换All-to-All操作中，每个节点将数据分散至集群所有节点，同时各节点也从集群所有节点收集数据。全交换All-to-All是全聚合AllGather的扩展形式，其区别在于：全聚合AllGather中不同节点从给定节点收集相同数据，而全交换All-to-All中不同节点从给定节点收集不同数据，本质上实现了数据转置操作。

全交换All-to-All是多对多转置操作，将数据从所有GPU卡转置至集群中每张GPU卡。其主要应用场景包括：

全交换All-to-All适用于模型并行场景；
模型并行中的矩阵转置环节；
从数据并行到模型并行的矩阵转置过程。

从通信原语到通信库：以NCCL为例

NCCL（NVIDIA Collective Communications Library）是为GPU间通信提供原语的库。它具有拓扑感知能力，可轻松集成到各类应用中。NCCL实现了集体通信与点对点收发原语，并非完整的并行编程框架，而是专注于加速GPU间通信的工具包。

NCCL提供的集体通信原语包括AllReduce、Broadcast、Reduce、AllGather、ReduceScatter、AlltoAll、Gather与Scatter。此外，NCCL还支持点对点收发通信，可灵活实现Scatter、Gather或All-to-All操作。

通过NCCL，通信原语在概念层面被转化为可用的库函数。在传统CUDA编程中，集体通信通常通过多次CUDA内存拷贝操作与本地规约计算内核组合实现；NCCL将每个集体通信原语封装在单个内核中，同步处理通信与计算，从而实现更高效的同步并降低资源开销。

总结

在分布式深度学习训练中，训练框架本身通常不会直接操作底层通信网络，而是通过高性能通信库完成参数同步、梯度规约等操作。这些通信库既屏蔽了复杂的硬件细节，又提升了训练性能。不同的AI芯片和加速卡厂商会提供专属的通信库或扩展方案，例如硬件感知的MPI实现或自研CCL库，以优化底层网络的使用。

实际分布式训练集群中的网络互联形态极其多样：既可能采用以太网、InfiniBand、RoCE v1/v2等通用协议，也可能使用NVLink等专用或半专用高速互联协议。这就需要通信库能够结合厂商提供的SDK和接口，针对特定硬件进行优化，例如CUDA-aware MPI、NCCL、NVSHMEM等。同时，还需要根据实际的网络拓扑和节点布局，选择最合适的通信策略与算法。

参考资料:

文章目录