[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

FlagScale 实现的异构混合训练方案的优势：

• 提高资源利用率：通过异构混合训练，可以更有效地利用现有的硬件资源，包括不同型号和性能的设备。
• 提升训练效率：异构混合训练可以提高训练效率，尤其是在资源有限的情况下。
• 保持模型性能：通过精心设计的并行策略，即使在异构环境下，也能保持模型的性能。
• 灵活性和可扩展性：这种方案提供了更多的灵活性和可扩展性，可以根据具体的硬件配置和需求进行调整

FlagScale 异构混合训练方案-官方实验结论

异构训练效率接近同构方案：

• 实验表明，使用 FlagScale 进行异构混合训练的效率接近理论上的上限，即合池训练没有带来任何性能损失的情况。
• 在某些情况下，由于混合后资源变多，甚至可能解锁新的优化空间，从而带来更高的训练效率提升。

异构训练方案不影响模型性能：

• 实验结果显示，通过 FlagScale 异构方案训练的 Aquila2-70B-Expr 模型，很好地保持了与同构训练一致的性能。
• 在 NVIDIA 异构训练或国产芯片上异构训练，对模型性能的影响甚微。

国产芯片集群上的同构训练能够达到 NVIDIA 同构训练上的模型性能：

• 尽管国产机器训练可能需要对检查点（checkpoint）进行重切分，无法完全保持随机初始状态一致，但实验表明，国产芯片集群上的同构训练能够达到 NVIDIA 同构训练上的模型性能。

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

• 开源的 Aquila2-70B-Expr 模型虽然只经过了 1.2T 数据的训练，但依然表现出强大的潜力，为学界和业界提供了一个优秀的探索起点。

AquilaChat2-70B-Expr 在主观能力评测中表现优异：

• 在经过与 Aquila2-34B 相同的指令数据集进行微调训练后，AquilaChat2-70B-Expr 在主观能力评测中超过了经过 2T 数据训练的 AquilaChat2-34B，展现了出色的学习能力。

实测结论：

• 实测证明，不管是在 NVIDIA 芯片还是在国产芯片上进行异构混合训练，都有性能收益。
• 使用 FlagScale 进行混合异构训练效率接近上限，并且由于混合后资源变多，在某些情况能解锁新的优化空间。

源码分析

项目目标

优化大规模模型的计算资源使用，同时保持数值稳定性和模型效果。

功能亮点

• 提供 BAAI 大型模型训练的实际配置、优化方案和超参数设置。
• 帮助开发者快速建立包括训练、微调、推理和服务在内的基础且完整的语言模型（LLM）流程。
• 支持 Aquila 模型的异构混合训练。
• 支持模型权重转换到 Huggingface 格式。
• 与上游项目保持同步更新。

更新日志

• 2024.4.11 发布新版本（v0.3），支持 NVIDIA 和 Iluvatar 芯片的 Aquila2-70B-Expr 模型异构混合训练。
• 2023.11.30 发布新版本（v0.2），支持不同代际同构或兼容架构芯片的异构训练，包括 NVIDIA 和 Iluvatar CoreX 芯片。
• 2023.10.11 发布初始版本（v0.1），支持 Aquila 模型和 Aquila2-B、Aquila2-34B 的实际训练方案。

快速开始

• 建议开发者遵循 Megatron-LM 的使用方式。
• 提供了 Aquila 语言模型的安装和预训练指南。
• 使用 Hydra 进行配置管理，配置文件分为实验、系统、模型和数据四个主要部分。

异构训练

• 异构流水线并行：指定不同设备处理不同阶段的网络层。
• 异构数据并行：指定不同设备处理不同的数据实例。

转换到 HuggingFace

• 提供了将 FlagScale 模型转换为 Huggingface 格式的指南