大型语言模型（LLM）在语言理解和生成方面展现出前所未有的能力，在推理、数学、科学等领域取得了突破性进展。然而，训练 LLM 极其昂贵。例如，PaLM 使用 6,144 个 TPUv4 芯片来训练一个 5400 亿参数的模型，而 GPT-3 1750 亿参数模型的预训练则消耗了数千个 petaflop/s-days 的计算量。这促使人们需要降低 LLM 的训练成本，特别是对于下一代超智能模型的扩展而言。
低精度训练是降低成本最有希望的方向之一，因为它可以提供高速、小内存占用和低通信开销。大多数现有的训练系统，例如 Megatron-LM、MetaSeq 和 Colossal-AI，默认情况下使用 FP32 全精度或 FP16/BF16 混合精度训练 LLM。然而，对于大型模型而言，这并非实现完全精度的必要条件。随着英伟达 H100 GPU 的发布，FP8 正在成为下一代低精度表示的数据类型。理论上，与当前的 16 位和 32 位浮点混合精度训练相比，FP8 可以实现 2 倍的加速、50% – 75% 的内存成本节省和 50% – 75% 的通信节省，这对于扩展下一代基础模型非常有前景。
不幸的是，目前对 FP8 训练的支持很少且有限。唯一可用的框架是英伟达 Transformer Engine (TE)，但它仅将 FP8 用于 GEMM 计算，而仍然使用高精度（例如 FP16 或 FP32）来保留主权重和梯度。因此，端到端的加速、内存和通信成本节省非常有限，无法完全展现 FP8 的威力。为了解决这个问题，我们提出了一种针对 LLM 训练的极度优化的 FP8 混合精度框架。核心思想是将 FP8 计算、存储和通信渗透到大型模型训练的整个过程中，使前向和反向传播都使用低精度 FP8，从而与之前的框架相比大大降低系统工作量。具体来说，我们设计了三个优化级别，利用 FP8 来简化混合精度和分布式训练。这三个级别以增量的方式逐步整合 8 位集体通信、优化器和分布式并行训练。更高的优化级别表示在 LLM 训练期间使用更多 FP8。此外，对于大规模训练，例如在数千个 GPU 上训练的 GPT-175B，我们的框架提供了 FP8 低位并行性，包括张量、流水线和序列并行性，为下一代低精度并行训练铺平了道路。
使用 FP8 训练 LLM 并非易事。挑战源于诸如数据下溢或上溢之类的问题，以及 FP8 数据格式固有的较窄动态范围和降低精度所产生的量化误差。这些挑战会导致训练过程中的数值不稳定和不可逆的偏差。为了解决这些问题，我们提出了两种技术：精度解耦和自动缩放，以防止关键信息的丢失。前者涉及解耦数据精度对权重、梯度、优化器状态等参数的影响，并将降低精度分配给对精度不敏感的组件。后者是通过动态调整张量缩放因子来保留 FP8 数据格式表示范围内的梯度值，从而减轻所有减少通信期间发生的下溢和上溢。
为了验证所提出的 FP8 低精度框架，我们将它应用于 GPT 风格的模型训练，包括预训练和监督微调 (SFT)。实验结果证明了我们 FP8 方法的有效性，产生了巨大的效益，包括实际内存使用量减少 29% 到 39%（例如，GPT-7B 减少 29%，而 GPT-175B 减少 39%）以及与流行的 BF16 混合精度训练方法相比，权重相关的通信开销显着减少 63% 到 65%。在不改变任何超参数（例如学习率和权重衰减）的情况下，使用 FP8 训练的模型在预训练和下游任务中都表现出与使用 BF16 高精度的模型相当的性能。值得注意的是，在 GPT-175B 模型的训练过程中，我们的 FP8 混合精度框架将训练时间缩短了 37%，而 H100 GPU 平台上的内存消耗减少了 42%。更重要的是，随着模型规模的不断扩大，通过使用低精度 FP8 实现的成本降低可以进一步提高，如图 1 所示。
对于微调，我们使用 FP8 混合精度进行指令微调和人类反馈强化学习 (RLHF)，以更好地将预训练的 LLM 与最终任务和用户偏好对齐。具体来说，我们在公开的用户共享指令遵循数据（ShareGPT）上微调预训练模型。使用我们的 FP8 混合精度微调的模型在 AlpacaEval 和 MT-Bench 基准测试中表现出与使用半精度 BF16 的模型相当的性能，同时训练速度提高了 27%。此外，FP8 混合精度在 RLHF 中展现出巨大的潜力，RLHF 是一个在训练过程中需要加载多个模型的过程。通过在训练中使用 FP8，流行的 RLHF 框架 AlpacaFarm 可以使模型权重减少 32%，优化器状态的内存消耗减少 62%。这进一步证明了我们 FP8 低精度训练框架的多功能性和适应性。
我们为推动下一代 FP8 低精度 LLM 训练的设计做出了以下贡献：
- 一个新的 FP8 混合精度训练框架。 它以附加方式逐步解锁 8 位权重、梯度、优化器和分布式训练，使用方便。这个 8 位框架可以作为现有 16/32 位混合精度对应物的简单替代，无需对超参数和训练收据进行任何更改。此外，我们提供了一个 Pytorch 实现，可以在几行代码中实现 8 位低精度训练。
- 一个新的 FP8 训练的 GPT 风格模型系列。 我们将提出的 FP8 方案应用于 GPT 预训练和微调（即 SFT 和 RLHF），并在从 70 亿到 1750 亿参数的各种模型规模上证明了它的潜力。我们为流行的并行计算范式配备了 FP8 支持，包括张量、流水线和序列并行性，使 FP8 可以用于训练大型基础模型。我们开源了第一个基于 Megatron-LM 实现的 FP8 GPT 训练代码库。
我们预计，我们的 FP8 框架的发布将为下一代低精度训练系统建立一个新的范式，该系统专门用于大型基础模型。
FP8 LLM
混合精度已广泛应用于 LLM 训练，以提高计算和内存效率。最流行的混合精度方案是 FP16-FP32 和 BF16-FP32。由于 FP16 的数值范围有限，FP16-FP32 方案已被证明在训练大型模型时存在不稳定性。因此，社区现在通常采用 BF16-FP32 来训练 LLM，例如 Megatron-Turing NLG-530B、Bloom-175B 和 Gopher。其根本原因是 BF16 具有较宽的动态范围，可以保持数值稳定性，同时匹配全精度 FP32 的性能。此外，与 FP32 相比，BF16 使用的位数只有一半，因此可以显著减少内存占用，同时提高计算效率。
FP8 是从 16 位数据格式自然演进而来的，可以进一步降低计算成本。然而，使用低精度 FP8 训练 LLM 带来新的挑战。FP81 的动态范围和表示精度远低于 BF16 和 FP16，这不可避免地会导致更多训练崩溃，例如损失激增甚至 NaN。为了解决这些问题，提出了张量缩放技术。核心思想是在将更高精度值转换为 FP8 之前，先用缩放因子将其相乘，以便将它们移动到与 FP8 格式的表示范围更好地重叠的范围内。这种每张量缩放技术可以减少数据量化误差，同时提高数值稳定性和精度，从而可以使用低精度 FP8 来训练大型模型。
不幸的是，目前对 FP8 低精度训练的支持有限。英伟达 TE 仅支持 Transformer 中线性层的 FP8 计算，而其他所有操作（例如权重更新和梯度同步）仍然使用更高精度。在这项工作中，我们提出了一种针对 LLM 训练的极度优化的 FP8 混合精度策略。新的 FP8 优化包括三个关键方面：FP8 通信、FP8 优化器和 FP8 分布式训练。通过整合这些方面，LLM 的训练（例如 1750 亿参数的 GPT-3 模型）可以充分利用 FP8 低精度的优势，提高训练效率。
2.1 FP8 梯度和所有减少通信
现有的混合精度训练方法通常使用 16 位或 32 位数据类型来计算和存储梯度，这导致在整个训练过程中对集体通信的带宽要求很高。我们发现，将 FP8 直接应用于梯度会导致精度下降。根本问题在于低位所有减少操作导致的下溢和上溢问题。具体来说，在所有减少过程中，跨 GPU 的梯度累加会导致数值范围超出 FP8 的表示范围，从而导致梯度信息丢失，并最终导致模型训练崩溃。
为了解决这个问题，我们提出了一种新的自动缩放技术，用于在 FP8 所有减少过程中动态调整张量缩放因子。该技术通过在所有减少之前和之后应用缩放因子来确保 FP8 梯度在所有减少过程中保持在表示范围之内。这种动态缩放方法可以有效地防止下溢和上溢，从而提高 FP8 梯度通信的精度和稳定性。
2.2 FP8 优化器
除了梯度通信之外，优化器状态（例如动量和 RMSProp）也需要大量的内存，并且通常使用高精度（例如 FP32）进行存储。为了进一步降低内存占用，我们提出使用 FP8 来存储优化器状态。然而，直接将 FP8 应用于优化器状态会导致训练不稳定，因为优化器状态的精度对模型训练的稳定性至关重要。
为了解决这个问题，我们提出了一种精度解耦技术，它将优化器状态的精度与模型权重和梯度的精度分离。具体来说，我们使用 FP8 来存储优化器状态，同时使用更高的精度（例如 FP16 或 BF16）来存储模型权重和梯度。这种精度解耦方法可以有效地降低内存占用，同时保持训练的稳定性。
2.3 FP8 分布式训练
在分布式训练中，模型权重需要在多个 GPU 之间进行同步，这会导致大量的通信开销。为了降低通信开销，我们提出使用 FP8 来存储模型权重。然而，直接将 FP8 应用于模型权重会导致精度下降，因为模型权重的精度对模型的性能至关重要。
为了解决这个问题，我们提出了一种新的 FP8 分布式训练方法，它将模型权重的精度与梯度的精度分离。具体来说，我们使用 FP8 来存储模型权重，同时使用更高的精度（例如 FP16 或 BF16）来存储梯度。这种精度解耦方法可以有效地降低通信开销，同时保持模型的性能。
3 FP8 训练框架
我们的 FP8 混合精度训练框架旨在通过逐步引入 FP8 计算、存储和通信来提高 LLM 训练的效率。该框架包括三个优化级别，每个级别都针对 LLM 训练的不同方面进行优化。
3.1 级别 1：FP8 梯度通信
级别 1 框架使用 FP8 来存储和传输梯度，同时使用更高的精度（例如 FP16 或 BF16）来存储模型权重和优化器状态。这种方法可以有效地降低梯度通信的带宽要求，同时保持训练的稳定性。
3.2 级别 2：FP8 优化器
级别 2 框架使用 FP8 来存储和传输梯度和优化器状态，同时使用更高的精度（例如 FP16 或 BF16）来存储模型权重。这种方法可以进一步降低内存占用，同时保持训练的稳定性。
3.3 级别 3：FP8 分布式训练
级别 3 框架使用 FP8 来存储和传输梯度、优化器状态和模型权重。这种方法可以最大限度地降低内存占用和通信开销，同时保持训练的稳定性。
4 实验结果
为了评估我们 FP8 混合精度训练框架的有效性，我们在 GPT 风格的模型上进行了实验，包括预训练和监督微调 (SFT)。实验结果表明，我们的 FP8 框架可以显著降低内存占用和训练时间，同时保持模型的性能。
4.1 预训练
我们在 GPT-7B 和 GPT-175B 模型上进行了预训练实验，并与流行的 BF16 混合精度训练方法进行了比较。实验结果表明，我们的 FP8 框架可以将实际内存使用量减少 29% 到 39%，同时将权重相关的通信开销减少 63% 到 65%。此外，我们的 FP8 框架可以将 GPT-175B 模型的训练时间缩短 37%，同时将 H100 GPU 平台上的内存消耗减少 42%。
4.2 微调
我们在 AlpacaEval 和 MT-Bench 基准测试上对 GPT-7B 模型进行了指令微调实验，并与使用半精度 BF16 的方法进行了比较。实验结果表明，我们的 FP8 框架可以实现与 BF16 相当的性能，同时训练速度提高了 27%。此外，我们还评估了 FP8 在 RLHF 中的潜力，发现它可以使模型权重减少 32%，优化器状态的内存消耗减少 62%。
5 结论
我们提出了一种新的 FP8 混合精度训练框架，它可以有效地降低 LLM 训练的成本，同时保持模型的性能。该框架通过逐步引入 FP8 计算、存储和通信来实现优化，并提供三个优化级别，以满足不同的训练需求。实验结果表明，我们的 FP8 框架可以显著降低内存占用和训练时间，同时保持模型的性能。我们相信，我们的 FP8 框架将为下一代低精度训练系统建立一个新的范式，该系统专门用于大型基础模型。
参考文献
- Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., … & Amodei, D. (2020). Language models are few-shot learners. arXiv preprint arXiv:2005.14165.
- Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, N., Roberts, A., … & Tay, Y. (2022). PaLM: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311.
- Li, B., Zhang, M., Hu, J., Li, R., Ni, B., Xiong, Y., … & Liu, Z. (2023a). Colossal-AI: A unified deep learning system for large-scale parallel training. arXiv preprint arXiv:2302.09060.
- Li, J., Chen, M., Wang, S., Zhou, D., & Xiong, L. (2023b). AlpacaEval: A comprehensive evaluation benchmark for instruction-following language models. arXiv preprint arXiv:2304.05799.
- Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., … & Wu, Y. (2017). Mixed precision training. arXiv preprint arXiv:1710.03740.
- Micikevicius, P., Alben, J., Wu, Y., Teerapittayanon, S., & Gomez, A. (2022). fp8 training for deep learning. arXiv preprint arXiv:2203.14784.
- Nvidia. (2022a). Nvidia H100 Tensor Core GPU. https://www.nvidia.com/en-us/data-center/h100/
- Nvidia. (2022b). Nvidia Transformer Engine. https://developer.nvidia.com/nvidia-transformer-engine
- OpenAI. (2023). GPT-4. https://openai.com/product/gpt-4
- Rae, J. W., Borgeaud, S., Cai, T., Olah, C., Leike, J., Allen, D., … & Hinton, G. (2021). Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11406.
- Scao, T., Lachaux, A., Lhuillier, M., Schick, T., Bordes, A., Févry, T., … & Jaffré, J. (2022). Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2203.15556.
- ShareGPT. (2023). ShareGPT. https://github.com/nomic-ai/gpt-4-11k
论文综述和评论：FP8-LM: Training FP8 Large Language Models
《FP8-LM: Training FP8大型语言模型》是一篇关于使用FP8低比特数据格式训练大型语言模型（LLMs）的论文。该论文探索了在LLMs训练中使用FP8低精度数据格式的有效性和效率。作者提出了一个新的FP8自动混合精度训练框架，该框架通过逐步引入8位梯度、优化器状态和分布式学习，实现了LLMs的混合精度和分布式并行训练。论文通过实验结果展示，使用FP8混合精度训练框架，可以显著降低实际内存使用量，并且训练速度比广泛采用的BF16框架快75%。此外，该框架还具有通用性，可以应用于其他任务，如LLMs指令调优和强化学习。作者在论文中开源了FP8低精度训练框架。
这篇论文的主要贡献在于提出了一种新的FP8混合精度训练框架，该框架充分利用了FP8低精度数据格式的优势，降低了LLMs训练的成本。通过逐步引入FP8梯度、优化器状态和分布式学习，论文展示了在训练大型语言模型时，使用FP8混合精度训练框架可以同时降低内存使用和训练时间，从而降低了大型基础模型的训练成本。此外，该框架的通用性使其可以无缝应用于其他任务，为微调费用节约提供了可能。作者通过开源FP8低精度训练框架，为下一代低精度训练系统的设计奠定了基础。
论文的创新点在于提出了使用FP8低精度数据格式训练LLMs的新方法，并通过实验证明了其有效性。该方法充分利用了FP8的优势，降低了训练成本，并在保持模型精度的同时提高了训练效率。此外，论文还介绍了针对FP8训练的挑战和解决方案，如数据下溢和上溢问题以及张量缩放技术。
总的来说，这篇论文对于推动下一代FP8低精度训练系统的设计具有重要意义。通过引入新的FP8混合精度训练框架，论文提供了一种降低大型语言模型训练成本的有效方法，并为其他任务的低精度训练提供了一种可行的解决方案。该论文的实验结果表明，FP8混合精度训练框架在减少内存使用和提高训练效率方面具有显著的优势。尽管还存在一些挑战，如数据表示范围的限制和量化误差，但这些问题可以通过提出的解决方案得到有效解决。
微软自动混合精度 (MS-AMP) 简介
MS-AMP 是微软开发的用于深度学习的自动混合精度软件包。它利用 FP8（8 位浮点数）格式来加速训练过程，并减少内存占用，从而使训练更大的模型成为可能。
主要特点
MS-AMP 提供了三种优化级别，以满足不同的训练需求：
- O1 优化: 将 FP8 应用于权重和权重梯度，并在通信中支持 FP8。
- O2 优化: 支持 FP8 用于 Adam 和 AdamW 两种优化器。
- O3 优化: 支持 FP8 用于分布式并行训练和 ZeRO 优化器，这对于训练大型模型至关重要。
此外，MS-AMP 还提供了四个训练示例，涵盖了 Swin-Transformer、DeiT、RoBERTa 和 GPT-3 等流行模型。
与 Transformer Engine 的优势
与 NVIDIA 的 Transformer Engine 相比，MS-AMP 具有以下优势：
- 加速内存受限操作: 通过访问单个字节，而不是半精度或单精度，从而加速内存受限操作。
- 减少内存需求: 降低训练模型的内存需求，从而支持训练更大的模型。
- 加速分布式模型通信: 通过传输低精度梯度来加速分布式模型的通信。
- 缩短大型语言模型的训练时间: 允许使用更大的 mini-batch，从而缩短大型语言模型的训练时间。
性能
MS-AMP 在保持高精度的情况下，仅使用一小部分内存，并在各种任务上取得了出色的性能，包括 GPT-3、DeiT 和 Swin Transformer。
模型性能:
MS-AMP 在 GPT-3、Swin-Transformer、DeiT 和 RoBERTa 等模型上进行了评估，结果表明，使用 MS-AMP 训练的模型与使用 FP16/BF16 AMP 训练的模型性能相当。
系统性能:
在 NVIDIA H100 平台上训练 GPT-175B 模型时，MS-AMP 与 BF16 混合精度方法相比，实际内存使用量减少了 39%，训练时间减少了 37%。对于较小的模型，MS-AMP 的 O2 模式与 FP16 AMP 相比，可以为 Swin-1.0B 模型节省 44% 的内存，为 ViT-1.2B 模型节省 26% 的内存。
总结
MS-AMP 是一个强大的工具，可以显著提高大型语言模型的训练效率。它通过利用 FP8 格式来减少内存占用和通信开销，从而加速训练过程并支持训练更大的模型。
参考文献: