【FlagScale】异构算力混合训练方案

背景以及必要性

算力需求的高峰：随着人工智能（AI）和生成内容（AIGC）的发展，对计算资源的需求急剧增加。尤其是参数规模达到数百亿的大模型训练，需要大量的计算资源。
算力市场供应紧张：目前，算力市场供应紧张，获取大量相同型号的AI训练加速卡（如GPU）来训练大模型变得困难。
资源墙问题：企业在不同阶段购买了不同代际或不同厂商的AI加速硬件，这些硬件在数据中心中形成了多个“资源墙”，难以合池使用，限制了大规模异构混合训练的实现。
技术挑战：使用不同型号、不同架构的AI芯片进行大模型训练面临许多新的技术挑战，包括软硬件栈不兼容、通信效率低、负载均衡困难等。
成本问题：高昂的试错成本和技术适配难度使得许多企业难以负担大模型训练的费用。
开源社区的需求：为了促进AI大模型技术的发展，需要一个开源的、可扩展的并行框架来支持不同硬件上的大模型训练。
国产硬件的支持：随着国产AI硬件的发展，需要一个框架来支持这些硬件，促进国产AI硬件产业的繁荣。

开源项目

https://github.com/FlagOpen/FlagScale

以FlagScale训练的开源模型

https://github.com/FlagAI-Open/Aquila2
https://model.baai.ac.cn/models
https://huggingface.co/BAAI

FlagScale 有何特点

异构混合训练：FlagScale 支持在不同型号、不同架构的AI芯片上进行混合训练，打破了传统同构硬件训练的限制。
开源：FlagScale 是一个开源项目，允许社区贡献代码和优化，促进了技术的共享和进步。
高效性能：通过异构流水线并行和异构数据并行两种模式，FlagScale 能够实现高效的训练效率，接近同构训练的性能。
模型性能保持：FlagScale 训练的模型能够保持与同构训练一致的性能，确保了模型的质量和可靠性。
灵活性：FlagScale 支持多种并行策略的混合使用，包括数据并行、张量并行、流水线并行以及序列并行，提供了灵活的训练方案。
支持国产硬件：FlagScale 积极适配国产AI硬件，支持国产芯片在大模型训练场景的应用。
降低成本：通过异构训练，可以更有效地利用现有硬件资源，减少对昂贵同构硬件集群的依赖。
易于使用：FlagScale 提供了清晰的文档和示例，使得开发者能够快速上手，进行大模型训练。
社区支持：作为一个开源项目，FlagScale 得到了社区的广泛支持，包括代码贡献和问题解答。
持续更新：FlagScale 持续更新，以适应新的硬件和软件环境，确保框架的先进性和适用性。
多厂商支持：FlagScale 是首个支持多厂商异构算力合池训练的框架，有助于构建更加开放和多样化的AI生态系统。
性能优化：FlagScale 通过优化算法和超参数配置，提供了性能优化的参考基线

何为异构训练

同构训练

定义：同构训练指的是使用相同型号和相同架构的硬件设备进行模型训练。
特点：
- 硬件一致性：所有设备具有相同的性能和内存规格，便于管理和优化。
- 简化的编程模型：由于硬件一致，编程和调试相对简单。
- 通信效率：同构设备间通信通常更高效，因为硬件和网络接口标准化。
- 负载均衡：易于实现负载均衡，因为所有设备的处理能力相同。
挑战：
- 资源限制：需要大量相同型号的硬件，这可能导致资源浪费或成本过高。
- 扩展性问题：当模型规模超出单一硬件集群的处理能力时，扩展性受限。

异构训练

定义：异构训练是指使用不同型号、不同架构的硬件设备进行模型训练。
特点：
- 资源利用最大化：可以利用现有不同型号和性能的硬件资源，提高资源利用率。
- 灵活性：能够适应不同的硬件环境，提供更多的训练配置选项。
- 成本效益：通过使用现有的多样化硬件，可以降低购买昂贵同构硬件集群的成本。
挑战：
- 软硬件不兼容：不同硬件可能有不同的软件栈和性能特性，需要额外的适配工作。
- 通信效率：不同硬件间的通信可能存在效率问题，需要优化通信策略。
- 负载均衡：由于硬件性能差异，实现有效的负载均衡更加复杂

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

异构流水线并行（Heterogeneous Pipeline Parallelism）

在这种模式下，不同类型的设备负责处理神经网络中不同的层级。具体来说：

设备分配：根据设备的内存和计算能力，将网络的不同层分配给不同类型的设备。
微批次大小：所有设备使用单一的微批次大小（microbatch size）。
流水线处理：数据在设备间流动，每个设备完成其负责的层的计算后，将结果传递给下一个设备。

这种模式适合于内存需求大的设备放在流水线的前面，而算力大的设备放在后面，从而实现负载均衡
在这里插入图片描述
假设整个神经网络有6层，设备有3种类型，采用微批次大小为1：
a）类型A的设备处在流水线并行第1个阶段，负责计算网络层1；
b）类型B的设备处在流水线并行第2阶段，负责计算网络层2和层3；
c）类型C的设备处在流水线并行第3阶段，负责计算网络层4、5和6

异构数据并行（Heterogeneous Data Parallelism）

在这种模式下，不同类型的硬件处理不同的数据并行实例，但每个设备处理完整的神经网络层：

数据分配：不同类型的硬件处理不同的数据实例。
微批次大小：不同设备可能使用不同的微批次大小，根据设备的算力和内存进行调整。
完整层处理：每个设备处理的数据实例包括网络的所有层。

这种模式允许算力和内存较大的设备处理较大的微批次，而算力和内存较小的设备处理较小的微批次，实现负载均衡。

在这里插入图片描述
假设整个神经网络有6层，设备有3种类型：
a）类型A的设备处理数据并行第1个实例，负责计算网络层1~6层，采用微批次大小为3；
b）类型B的设备处理数据并行第2个实例，负责计算网络层1~6层，采用微批次大小为2；
c）类型C的设备处理数据并行第3个实例，负责计算网络层1~6层，采用微批次大小为1