打破异构算力束缚，FlagScale首次实现大模型的异构混合训练，开源70B模型实验版

AIGC浪潮带来了算力需求的高峰，而算力市场供应紧张的局面时日已久，数百亿参数规模的大模型训练需要庞大的计算资源。当前的训练方法仅支持同一种训练芯片构成集群来完成大模型训练，意味着数百亿参数的大模型训练需要上千张同一型号的AI训练加速卡（例如英伟达的GPU或某些国产训练加速卡），这在很多情况下难以满足。在长期的算力积累过程中，企业往往在不同阶段购买同一厂商不同代际的AI加速硬件、或者采购不同厂商的硬件，以满足不同工作负载的需求，同时尽可能地发挥各类设备的最大优势，节省使用成本。

但同一个数据中心中，这些不同型号的AI训练卡集群之间往往形成了多个“资源墙”（resource wall）。当我们需要训练大模型时，难以把同一个数据中心中不同型号的训练卡集群合池，以获得更大的训练资源。因此，如何打破“资源墙”的限制，实现大规模异构混合训练，是业界一直迫切需要解决的难题。

聚合异构算力进行大模型训练是解决算力束缚的重要途径和迫切需求，但也面临了巨大挑战。

• 大模型训练试错成本高：使用不同型号、不同架构的AI芯片进行大模型训练将面临许多新的、未知的技术挑战，不是每个企业都有能力负担高昂的试错成本。

• 芯片适配难度大：许多国产AI硬件新势力正在积极寻求与大模型训练场景的适配，但受限于不完善的软件生态，往往要求用户企业投入很多开发成本进行适配，尤其是模型训练。

智源团队通过升级FlagScale大模型并行框架，率先实现了在英伟达异构集群的混合训练，并通过与天数智芯团队合作，也实现了在天数智芯异构集群的混合训练；验证了异构并行训练方案在70B大模型训练上、在多种不同资源混合场景下的有效性，并通过FlagScale项目开源了所有实现代码、优化算法、超参数配置等，为业界提供可靠的参考基线。FlagScale目标是实现大模型稳定高效训练、微调与推理的并行框架。这一目标的实现离不开AI硬件企业的共同努力，当前天数智芯、昆仑芯已一同加入FlagScale开源社区进行代码贡献，期待更多合作伙伴参与开源贡献。同时，智源团队开源了使用异构混合方法训练得到的 Aquila2-70B-Expr（异构实验版）系列模型[1]，补齐70B参数量级中英双语大模型的缺口。该模型虽然仅经过1.2T数据训练，依然表现出强大的潜力，为学界和业界提供一个优秀的探索起点，可用于继续预训练及微调。

↓ 相关下载 

FlagScale开源项目

• https://github.com/FlagOpen/FlagScale

Aquila2-70B-Expr模型

• https://github.com/FlagAI-Open/Aquila2

• https://model.baai.ac.cn/models

• https://huggingface.co/BAAI

打破异构算力的束缚

为业界提供效率&模型性能基线

具体而言，当前异构混合分布式训练存在如下挑战：

• 不同架构设备的软硬件栈不兼容，数值精度也可能存在差异；

• 不同架构设备之间很难高效通信；

• 不同设备算力和内存不同，很难进行负载均衡切分。

这些挑战很难一次性解决，目前智源团队已经尝试在相同架构不同代际设备或者在兼容架构的不同设备上进行异构训练，未来将探索不同架构设备上的异构训练。

FlagScale 是首个支持多厂商异构算力合池训练的框架，当前实现了异构流水线并行及异构数据并行两种模式。经过实测验证，FlagScale v0.2 版本中的两种异构方案的实现均不影响70B模型的训练效率及模型性能。在实际训练过程中，FlagScale 实现的异构训练方案还有诸多“惊喜”：

1. 训练效率方面：不管是在NVIDIA芯片还是在国产芯片上进行异构混合训练，都有不错的性能收益，而且混合训练后有可能释放新的优化空间，带来更高的训练效率提升。

2. 模型性能方面：经过FlagScale异构方案训练的 Aquila2-70B-Expr模型，很好的保持了与同构训练一致的性能。此外，FlagScale还验证了在国产同构训练也能够达到NVIDIA同构训练上的模型性能，这对于国产芯片在大模型训练场景落地有重要的参考价值。

FlagScale异构混合训练方案‍‍‍‍‍

• 异构流水线并行：不同类型设备处理不同的网络层，同一类型设备处理相同网络层，所有类型只使用单一的微批次大小（microbatch size），具体流程如图1所示。实际训练时，异构流水线并行模式可以跟数据并行、张量并行以及序列并行进行混合来实现高效训练。根据反向传播算法内存使用特点，该模式适合将内存比较大的设备放在流水线并行靠前的阶段，内存小的设备放在流水线并行靠后的阶段，然后根据再设备的算力来分配不同的网络层来实现负载均衡。

• 异构数据并行模式：不同类型硬件处理不同数据并行实例，同一类型设备处理相同且完整的神经网络层，但不同类型设备将使用不同微批次大小，参考类似的工作[2]，具体流程如图2所示。实际训练时，异构数据并行模式可以跟张量并行、流水线并行以及序列并行进行混合来实现大规模高效训练。在该模式下，算力和内存都比较大的设备将处理较大的微批次大小，而算力和内存都比较小的设备将处理较小的微批次大小，从而实现不同设备上的负载均衡。

图1  异构流水线并行方案

假设整个神经网络有6层，设备有3种类型，采用微批次大小为1：

a）类型A的设备处在流水线并行第1个阶段，负责计算网络层1；

b）类型B的设备处在流水线并行第2阶段，负责计算网络层2和层3；

c）类型C的设备处在流水线并行第3阶段，负责计算网络层4、5和6。

图2  异构数据并行方案假设整个神经网络有6层，设备有3种类型：

a）类型A的设备处理数据并行第1个实例，负责计算网络层1~6层，采用微批次大小为3；

b）类型B的设备处理数据并行第2个实例，负责计算网络层1~6层，采用微批次大小为2；

c）类型C的设备处理数据并行第3个实例，负责计算网络层1~6层，采用微批次大小为1。

实测结论一：异构训练效率接近同构方案

为了评估异构混合训练的效率，我们定义了训练效率计算的方式、上限及下限。实测证明，不管是在NVIDIA芯片还是在国产芯片上进行异构混合训练，都有性能收益；使用FlagScale进行混合异构训练效率接近上限，并且由于混合后资源变多，在某些情况能解锁新的优化空间。

实验过程如下：

假设有两个不同类型硬件的集群A和B，f(A)表示模型在集群A上实测的每秒处理总token数（tokens/cluster/s），f(B)表示同一模型在集群B上实测的每秒处理总token数（tokens/cluster/s）。我们定义用集群A和集群B进行异构混合训练的性能下限为 lower bound=max(f(A), f(B)) ，即由于不能合池训练，因此只能使用其中最大能力的单一集群；性能上限为upper bound=f(A)+f(B)，即合池训练没有带来任何性能损失。这些定义也适用于扩展到多个集群，为了简单起见，这里用两个集群为例。

• 当异构混合训练实测性能超过lower bound，说明异构混合训练比用任意单一集群训练都有收益，超过lower bound越多说明通过异构混合训练聚合算力越有价值；
• 当混合训练实测性接近upper bound，说明异构混合训练后性能相比各个集群上单测性能总和来说，效率损耗较小，越接近upper bound说明异构混合训练的并行策略和优化越好。

我们分别在英伟达和天数智芯集群进行了三组异构混合训练实验，实验结果如下表所示。

表1  混合训练实验

• 配置1 结果显示：异构混合训练收益非常大，而且超过了upper bound，这是因为混合集群后解锁了新的优化，原来小集群下需要进行重计算策略才能训练起来模型，但在混合后的大集群就不需要进行重计算了。

• 配置2和3结果显示：异构混合训练有收益，而且离upper bound较近，说明异构混合训练的效率损耗较低。

实测结论二：异构训练方案不影响模型性能

为了验证英伟达以及国产芯片集群上的异构混合方案是否会对大模型训练产生性能影响，我们在Aquila2-70B-Expr上进行了实际模型性能评测对比，包括：1）NVIDIA A800+A100异构训练；2）天数BI-V100同构训练；3）天数BI-V100+BI-V150异构训练。

实验有以下结论：

• FlagScale v0.2 提供的异构流水线并行与异构数据并行实现，都可以很好地保持模型性能；

• 在NVIDIA异构训练或国产芯片上异构训练，对模型性能影响甚微；
• 国产芯片集群上的同构训练能够达到NVIDIA同构训练上的模型性能（尽管国产机器训练会对ckpt进行重切分，无法完全保持随机初始状态一致）；

具体评测方法：

1. 都选择与NVIDIA同构机器上的Aquila2-70B-Expr作为基线实验；
2. 基线实验与对比实验都选择相同的70B模型checkpoint开始，保持模型、超参、数据集一致，同时训练到相同的迭代步停止；

3. 将基线实验和对比实验送到真实数据集上进行评测，计算差异。

下表是规范化后的结果，其中每个数据集 diff%=(对比实验结果-基线实验结果)/基线实验结果*100%，所有数据集平均 diff%=sum(每个数据集diff%)/所有数据集个数。实验起始ckpt来自于NVIDIA同构机器训练。

表2  混合训练实验-模型性能评测

同步开源Aquila2-70B系列模型异构实验版

开源中英双语大模型矩阵“小巨人”入列

当前全球开源的70B大模型非常稀少，大多仅支持英文，而国内支持中英双语的开源大模型参数量也集中在6B至30B。

开源Aquila2-70B-Expr（异构实验版）系列模型旨在扩充70B参数量级中英双语大模型阵列，为学界和业界提供一个较为优秀的探索起点。虽然只经过1.2T数据训练，但Aquila2-70B-Expr基座模型仍然表现了出色的潜力。基于FlagScale项目及Aquila2-70B-Expr模型，开发者可以在同构或异构芯片集群上面，进行继续预训练或微调。这意味着，开发者在训练算力资源上拥有更多选择空间。

表3  Aquila2-70B base 模型评测结果

由于Aquila2-70B-Expr目前已经训练的英文数据量仅为LLama2-70B的大约三分之一，所以在MMLU上还依然落后于LLama2-70B。我们曾经对Aquila2-70B-Expr进行以MMLU训练集进行增广的数据进行了一小段持续训练，Aquila2-70B-Expr能迅速在MMLU的总体评测上提升至80.7分，可见该基础模型极为出色的学习能力。但我们认为，Aquila2-70B-Expr的更大价值在于给后续使用者一个泛化能力强大、后续学习能力强大的基座模型，因此依然保持模型在当前的状态进行开源。

为了展现对下游任务的学习能力，我们对Aquila2-70B-Expr基础模型采用了与最新发布的Aquila2-34B正式版本同样的指令数据集进行微调训练。AquilaChat2-70B-Expr 在主观能力评测中，超过了经过2T数据训练的 AquilaChat2-34B，也超过了其它模型在CLCC v2.0主观能力评测的得分，同样展现了极为出色的学习能力。

表4  AquilaChat2-70B SFT模型主观评测结果

CLCC v2.0 为智源自建的大语言模型主观评测集，详细介绍见文章

未来工作

智源团队将继续致力于研究在多种国产硬件、异构训练所需的并行策略和性能优化关键技术，形成一套面向大模型训练与推理的基础框架。

后续FlagScale将进一步优化当前异构训练实现，尝试将现在的两种异构模式进行混合，提供更灵活高效的异构训练方案；同时，为了支持国产AI硬件产业繁荣，FlagScale也将积极适配更多的国产硬件，并与FlagPerf进行联动，作为FlagPerf多机多卡评测的官方基准程序，积极探索不同国产硬件上的异构训练方案，积极参与异构通信库及相关标准研制工作。

FlagScale是智源FlagOpen（飞智）大模型开源技术体系的重要组成部分。FlagOpen旨在打造全面支撑大模型技术发展的开源算法体系和一站式基础软件平台，支持协同创新和开放竞争，共建共享大模型时代的“新Linux”开源开放生态。FlagOpen 包含大模型算法、模型、数据、工具、评测等重要组成部分，国内外开发者基于FlagOpen可以快速开启各种大模型的尝试、开发和研究工作，企业可以低门槛进行大模型研发。同时，FlagOpen大模型基础软件开源体系正逐步实现对多种深度学习框架、多种AI芯片的完整支持，支撑AI大模型软硬件生态的百花齐放。

备注:‍‍‍‍‍‍‍‍‍‍

[1] Aquila2-7B-Expr的训练过程是：于同构资源上进行了一段时间训练后，再迁移到异构资源上完成后续训练。 

[2] X. Jia et al., “Whale: Efficient Giant Model Training over Heterogeneous GPUs.”