一次性讲明白，如何搞定一个可以支持多芯混合训练的 AI 集群

2024-04-30
上海
本文字数：4425 字
阅读完需：约 15 分钟

由于外部环境的变化，适用于大模型训练任务的 GPU 整体规模无法继续增长。这些存量 GPU 组成的集群，仍然是当前加速大模型训练的主要 AI 算力来源。同时，各类国产 AI 芯片开始大规模投入实际生产任务。在未来一段时间内，数据中心的 AI 算力将保持多种芯片并存的现象。

但是，当前基础大模型训练所需要的最大 AI 算力集群规模，已经从单一集群千卡逐步提升至万卡量级。同时，很多智算中心已经部署的 GPU 集群，通常是十几台至数百台服务器不等，难以满足未来行业大模型训练的需求。

所以，在已有 AI 算力集群的基础上，构建由 GPU、昆仑芯、昇腾等不同芯片混合组成的单一集群，为大模型训练提供更大 AI 算力，成为了一个自然的选择。

大家都知道，成功建设一个全部由 GPU 芯片组成的集群，让他们作为一个整体可以高效率地跑起来，就已经足够复杂。（详情可参考：AI 大底座，大模型时代的答卷）

如果还要在 GPU 集群中再加上其他类型的 AI 芯片，他们各自讲着不同的语言，拥有完全不同的能力，让他们像同一种芯片组成的集群一样，实现多芯混合训练，加速单一大模型训练任务，那真是复杂到头了。

这里的挑战，和一个经常出现在团建中的竞技项目「多人多足」类似：大家肩并肩站成一排，人和人之间依次绑着小腿，大家齐头向前跑冲向终点。如果这个时候参赛队中有些人只会外语，有些人步子迈得很长……

1 如何建立和加速一个 GPU 集群

为了让大家对「如何搞定一个支持多芯混合训练的 AI 集群」有更清晰的理解，我们首先以 GPU 集群为例，简单介绍建立和加速一个 AI 集群的三个关键方面。

1.1 实现 GPU 互联互通

为了建设一个多卡集群，首先要完成 GPU 卡在物理层面进行连接。在单台服务器内 8 块 GPU 卡通过 NVLink 连接。不同服务器之间的 GPU 卡通过 RDMA 网络连接。

在完成网络的搭建后，借助 NVIDA 开发的集合通信库 NCCL，GPU 就能通过网络实现相互通信完成数据同步，使得训练任务可以一轮一轮地往下推进，直到完成大模型的训练。

1.2 制定分布式并行策略

为了加速大模型训练任务，我们需要将这个任务拆分到集群的所有 GPU 中，使得这些 GPU 能够共同完成任务。这就是常说的分布式并行策略。分布式并行策略有很多种，比如从训练数据维度进行切分的数据并行，按照模型的不同层面进行切分的流水线并行等。

我们需要依据集群的物理拓扑和大模型的参数，找到最优的分布式并行策略，充分发挥集群效能。在任务运行过程中，集群中所有 GPU 步骤保持一致，同时开始计算，同时开始通信，不存在一些 GPU 空转等待另外一部分 GPU 的情况。

1.3 部署 AI 加速套件

按照分布式并行策略被拆分开的模型和数据，将会以一个个算子的形式部署在 GPU 进行计算。为了加速 GPU 对算子的计算过程，我们还需要一个 AI 加速套件。

这个 AI 加速套件需要包含数据加载、算子计算（各种 CUDA 库）、多卡通信（ NCCL 集合通信库）等各个方面的优化。比如采用数据预取策略，使得 I/O 的过程和 GPU 上的计算充分并行起来；使用 NVIDIA 优化后的 GPU 算子或者全新的算子加速 GPU 卡的计算效率；更新 NCCL 能力提升 GPU 卡相互通信的效率。

我们这里总结一下，为了建设一个能够高效训练大模型的集群，需要在卡间和机间建立高效的互联互通，将大模型训练任务按照合适的并行策略拆分到 GPU 卡中，最后通过各种优化方法，加速 GPU 对算子的计算效率，完成大模型训练。

2 建立不同芯片集群的差异

当前，在数据中心的多芯算力的运用方式上，主流仍然是采用一种芯片对应一个集群的思路，这需要根据每一种芯片的特点进行量身定制。

参照上文提到的三个方面，一起来看看基于昆仑芯和昇腾 910B ，建设和加速这些集群的差异。

在互联互通上，昆仑芯服务器内部通过 XPU Link 进行连接，服务器之间通过标准的 RDMA 网卡进行连接，卡和卡之间使用 XCCL 通信库进行相互通信。昇腾 910B 服务器内部通过 HCCS 进行连接，服务器之间通过华为自研的内置 RDMA 进行连接，卡和卡之间使用 HCCL 通信库进行相互通信。
在并行策略上，NVIDIA GPU 和昆仑芯采用单机 8 卡的部署方式，昇腾 910B 则是机内 16 卡分为 2 个 8 卡通信组。这意味着在 AI 框架下形成不同的集群拓扑，需要有针对性地制定分布式并行策略。
在 AI 加速套件上，由于 GPU、昆仑芯、昇腾等芯片在计算能力，显存大小，I/O 吞吐，通信库等均存在差异，故需要面向具体芯片进行特定优化。最后的结果，就是每一种芯片，有一个各自对应的算子库，以及相应的加速策略。

3 建立和加速多芯混合集群的挑战和方案

现在，我们回到今天本文讨论的重点：建设一个支持多芯混合训练的 AI 集群，并加速运行一个大模型训练任务。（这里有个背景知识：要完成一个大模型的训练任务，只能是在单一集群中完成，而不能拆分到不同集群中进行。）

为了实现这个目标，在前文提到的三个维度我们都将遇到挑战：

不同类型的卡的物理连接方式和集合通信库是完全不同的，这导致他们之间无法直接互联互通。
我们需要为跨芯片的互联互通设计一套新的物理网络架构，并配套相应的集合通信库方案，使得通信效果达到最优。
在集群建设完成后，由于不同种类卡的性能不一致，这就导致在单一芯片集群中，基于均匀切分计算量制定分布式并行策略的方法，将无法确保不同芯片可以按照相同节奏计算和通信，导致算力浪费。
我们需要为此找到一套最优的分布式并行策略，使得这些不同种类的芯片能够在计算和通信的节奏上保持一致，使得集群算力效能最大化。
由于不同卡的加速方法存在差异，生态能力和调优策略发展程度不一，这将导致集群中的各类算力效能无法充分发挥，即 MFU 未达到理想值（ MFU 名词解释详见文末）。
我们需要为所有芯片提供统一的加速抽象以便屏蔽这些差异性，同时又能提供最优的 AI 加速套件，确保各种芯片均可以达到 MFU 理想值，实现高效能运行。

3.1 跨芯片的互联互通，构建多芯混合集群

传统的观点认为，不同芯片是很难互联互通，无法支撑大模型训练。从上文可知，NVIDIA GPU、昆仑芯、昇腾 910B 的物理连接方式，以及使用的集合通信库都不一样。

百度百舸为了实现跨芯的互联互通，使用了 CPU 转发来实现跨昇腾 910B 子集群和 GPU 子集群的连接。借助百度自研的集合通信库 BCCL，可以实现 GPU、昆仑芯等标准 RDMA 设备的互联互通，使得通信效果达到最优。

3.2 自适应并行策略搜索，提升多芯混合训练任务的整体效能

传统的分布式并行策略，都是按照等分的方式将大模型和训练数据拆开。其中确保这种等分方式有效的前提，在于集群中的芯片是同型号。这样在算量相同的情况下，所有卡都可以同节奏地运行。

假设集群中存在两种以上的卡，按照等分的方式制定分布式并行策略，则存在高性能卡等待其他卡的过程，产生算力的浪费。所以在同一个多芯集群中，我们需要按照不同芯片子集群的算力对比，分配合适的算量，将过去均匀切分的分布式并行策略改成按芯片算力大小适配的非均匀分布式并行策略方式。

在确定了非均匀并行策略的大方针后，还需要解决的就是具体怎么分，分多少的问题，确定分布式并行策略的最优解。这个最优解包括：采用什么样组合的分布式并行策略；在不同芯片的子集群中分配多少算量，比如分配多少训练数据，多少模型层数等。

百度百舸的 AI 加速套件 AIAK–LLM 实现了针对单一任务多种芯片的自适应并行策略搜索功能，通过计算各种并行策略所需要的计算量、存储量、通信量以及不同芯片的计算和 I/O 效率等（这里的所有数值均源自百度基于数十年 AI 技术积累产生的手册：AI 芯片效能矩阵图谱），从而快速计算出最优的任务切分策略，保证在各种芯片配比下的单一集群，在运行大模型多芯混合训练任务时整体效能最大化。

3.3 Accelerator 抽象，屏蔽硬件差异充分发挥不同芯片的算力效能

传统的加速方案，如上文所说，均是面向特定芯片进行优化。如果同一个大模型任务，运行在多芯混合集群上，我们需要为不同芯片配置相应的加速策略。

国产化 AI 芯片由于生态还在不断完善过程中，调优策略仍然需要继续优化，所以导致算力没有充分发挥出来。

百度百舸基于在 GPU 上的长期投入与沉淀（比如，百度百舸平台的 A800 MFU 值达到 80。该值等于 A800 MFU 的理想值，即百度百舸完全发挥了 A800 的能力），在 AI 加速套件 AIAK-LLM 中构建了「Accelerator 抽象层」，使得国产化 AI 芯片充分发挥各自算力。（关于 MFU 的含义，请参考文末名词解释）

其中，「Accelerator 抽象层」面向应用，屏蔽底层芯片在硬件层面的差异，将芯片算子与上层策略解耦开来，芯片厂商仅需要进行各自芯片的算子调优。百度百舸在 GPU 上沉淀的上层各项策略（比如通信 overlap、显存 overlap 、张量并行、数据并行等）都会平滑地迁移到各种芯片上，确保了各种国产芯片在百度百舸上能达到一个非常高的运行效率。