专栏名称: GiantPandaLLM

专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创，每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你，大家一起共同进步(･ω<)☆

一文讲清 NCCL 集合通信原理与优化

GiantPandaLLM · 公众号 · 3D · 2024-09-18 23:49

正文

专家并行，就是通过选择性的激活一部分参数来处理不同的数据，从而解决模型规模增大训练成本成平方级别的增长的问题。

作为基于 Transformer 的 MoE 模型，主要由以下两部分组成：

稀疏 MoE 层：它将 FFN 拆成多个子层，每一个子层被称为Expert。一般来说，这些 Expert 都是 FFN，但是也可以是更复杂的网络，甚至是 MoE 本身。
Router：也被称为 Gating Network，这部分用于决定将哪些 token 被发送到哪些 Expert。

专家并行的思路是将不同的专家分配到不同的 GPU 上，这有助于减少内存消耗并提高训练效率。计算前需要根据路有规则将 Token 通过 All-to-All 通信发送给不同的 Experts 所在的 GPU 进行运算。

并行模式	集合通信	通信量	单步通信次数
数据并行	all-reduce或reduce-scatter 和 all-gather	单GPU 10 GB 以上	1次
张量并行	AllReduce	GB 级别	多次
流水并行	点对点通信	MB	几十次
序列并行	all-gather和reduce-scatter	GB 级别	1次
专家并行	All-to-All	GB 级别	多次

从上表可以看出，不同的并行模式下都离不开集合通信技术，所以说大模型的分布式训练是离不开集合通信的。

为了更好的实现 scaling law 扩大规模，如上图所示，在真实的训练中，我们一般采用多种模式混合并行。那么模型的训练中也处处充满集合通信。

说起集合通信，就不得不提到 MPI。MPI 是集合通信中的元老，是消息传递的平台，是一种标准和规范，而非具体并行语言。真正的实现是在开源实现库里，例如openMPI。

我们知道，一旦我们的程序运行起来，它就会成为独立的进程。 进程拥有独立的执行环境（内存、寄存器、程序计数器等），是操作系统中独立存在的可执行的基本程序单位。

进程间相互独立（内存空间不相交），但是进程间可以相互交换信息。消息传递是指这些信息在进程间的相互交换，是实现进程间通信的唯一方式。

最基本的消息传递操作包括发送消息send、接受消息receive、进程同步barrier、归约reduction等。

MPI 程序一般包含，变量定义，MPI 环境初始化, 执行进程间通信和退出MPI 环境。

MPI系统的通信方式都建立在点对点通信之上，包括阻塞式点对点通信和非阻塞式点对点通信等。

在 OpenMPI 中，有多种集合通信算法的实现,例如 Allreduce, AlltoAll, AllGather等，且每种也都有不同的实现。

例如，MPI中实现了多种AllReduce算法，例如最传统reduce+broadcast的方式，又或者butterfly的方式，或者Ring AllReduce，或者Segmented Ring 的方式。

一个 Allreduce 有多种实现，可以方便在不同的情况下采用不同的 Allreduce 的实现，来加快集合通信的计算。