FlagScale 全面升级，支持跨厂商异构芯片合池训练

近日，FlagScale 开源大模型并行训练框架全面升级。智源团队与天数智芯团队合作，实现了“英伟达芯片+其它AI芯片”集群上单一大模型任务的异构混合训练，并在70B大模型上验证了不同架构芯片上进行异构混合训练的有效性。同时，为了加速多种AI芯片在大模型训练场景的使用，智源积极探索高效灵活的芯片适配方案，通过与硬件厂商的深入合作，FlagScale 已在6家不同厂商的多款AI芯片上适配Aquila2系列大模型的大规模训练。

FlagScale目标是在多芯片上实现大模型稳定高效训练与推理的并行框架，期待更多合作伙伴参与开源贡献。

FlagScale v0.3版本：

https://github.com/FlagOpen/FlagScale/tree/release/v0.3

FlagScale支持跨厂商异构芯片混合训练

FlagScale 已经实现同一厂商不同代际或型号的芯片合池训练，这次更进一步，实现了“英伟达芯片+其它AI芯片”之间异构混合训练。由于不同厂商的卡间互联协议不同，为了实现“英伟达芯片+其它AI芯片”高速互联，智源团队与天数智芯协作，优化天数智芯的iXCCL通信库，使其在通信元语操作上以及API接口上兼容英伟达NCCL，然后将框架编译链接到同一iXCCL通信库上，从而在用户和AI框架无感知的情况下实现异构算力芯片间高效通信，进而实现不同架构芯片混合训练。同时，双方还协作优化了流水线并行的分配方式，并针对不同芯片算力、内存带宽、内存容量的差异来为不同芯片配置不同的流水线并行策略，以使得训练过程中能充分发挥不同芯片的性能，最终率先实现了通用GPU异构大模型高效训练方案。

实测结论一：跨厂商异构芯片混合训练效率损耗较低

为了评估异构混合训练的效率，我们定义了异构混合训练效率的上限及下限，并在“英伟达+天数智芯”混合集群进行了异构混合训练实验，实验结果表明，混合集群训练效率损耗较低，其训练吞吐约为两个单一芯片集群上吞吐之和的 92.81%，说明“英伟达+天数智芯”的混合训练有收益。

表1: 混合训练试验

实验过程如下：

假设有两个不同类型硬件的集群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说明异构混合训练的并行策略和优化越好。

实测结论二：跨厂商异构芯片训练不影响模型性能

为了验证“英伟达芯片+其它AI芯片”集群上的异构混合方案是否会对大模型训练产生性能影响，我们在Aquila2-70B-Expr上进行了实际模型性能评测对比，实验结果如下表所示，可以看到模型很好地保持了与同构训练一致的性能，说明“英伟达芯片+其它AI芯片”集群上进行异构训练，对模型性能影响甚微。

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

实验过程如下：

1. 选择与NVIDIA同构机器上的Aquila2-70B-Expr作为基线实验；

2. 基线实验与对比实验都选择相同的70B模型checkpoint开始，保持模型、超参、数据集一致，同时训练到相同的迭代步停止；

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

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

FlagScale 已成功适配6家AI芯片

FlagScale已经成功支持Aquila2系列模型在昆仑芯、天数智芯等6家不同的AI芯片上的大规模训练，并作为FlagPerf大规模训练评测后端，完成了各芯片与英伟达机器在相同配置下的逐步loss对齐和整体loss趋势对齐。

注：FlagPerf是智源推出的开源AI硬件评测引擎，项目地址 https://github.com/FlagOpen/FlagPerf

为了加速多种AI芯片在大模型训练场景的使用，FlagScale 为硬件厂商提供两种适配模式：monkey-patching模式和in-place模式，欢迎各大芯片厂商一起共建开源生态。

• Monkey-patching模式：因为 Python 是一种动态类型语言，允许在运行时改变对象的结构，该模式就是利用这一特性，在不修改源代码的情况下动态地对FlagScale中的类或者模块进行替换或者修改，而用户对此修改无感知。

• In-place模式：硬件厂商可以通过全局唯一设备类型标识直接在FlagScale代码中添加对应修改逻辑，这种方式对代码进行侵入式修改，适用monkey-patching无法解决的特殊情形。

未来工作

智源团队将继续致力于研究在多种AI芯片、异构训练所需的并行策略和性能优化关键技术，并与更多合作伙伴共建支持超大模型和兼容不同AI芯片的开源分布式训练与推理框架。

后续 FlagScale 将进一步优化当前异构混合并行策略和完善异构数值精度对齐方法，提升异构混合训练效率，积极参与异构通信库及相关标准研制工作，争取在更多芯片上实现端到端的高效异构混合训练。同时，为了支持AI芯片产业和生态繁荣，FlagScale也将继续与厂商协作，完善适配机制和进行深度优化，实现更多AI芯片上的端到端大规模高效训练。

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