大型语言模型以前未曾有过的语言创造和理解的能力已经得到展示，为逻辑、数学、物理和其他领域的进步铺平了道路。但大型语言模型的训练非常昂贵。例如，要训练一个540B模型，PaLM需要6144个TPUv4芯片，而GPT-3 175B的预训练需要数千个拍夫洛普/秒的计算。这凸显了降低大型语言模型训练成本的需求，尤其是为了扩展下一代极智能模型。其中最有前景的节约成本方法是低精度训练，它可以提供快速处理、少内存使用和最小的通信开销。目前大多数的训练系统，如Megatron-LM、MetaSeq和Colossal-AI，通常默认使用FP16/BF16混合精度或FP32完全精度来训练大型语言模型。
对于大型模型来说，这是为了获得完全的准确性而非必需的。随着Nvidia H100 GPU的到来，FP8正在成为下一代低精度表示的数据类型。与现有的16位和32位浮点混合精度训练相比，FP8理论上具有2倍的加速、50%至75%的内存成本降低以及50%至75%的通信节省。这些结果对于扩展下一代基础模型非常鼓舞人心。遗憾的是，FP8训练还需更多、更少的支持。Nvidia Transformer Engine是唯一可行的框架，然而它只使用FP8进行GEMM计算，并将主权重和梯度保持在极高精度（如FP16或FP32）中。因此，综合性能提高、内存节省和通信成本节约相对较小，使FP8的全部潜力仍然隐藏着。
Microsoft Azure和Microsoft Research的研究人员提供了一个高效的FP8混合精度框架，用于解决大型语言模型训练中的这一问题。其主要原理是在大型模型训练过程中利用低精度的FP8进行计算、存储和通信。这将大大降低系统需求，与之前的框架相比。更准确地说，他们创建了三个优化阶段，使用FP8简化分布式和混合精度训练。三个阶段逐步引入了优化器、分布式并行训练和8位集中通信。更高的优化级意味着在大型语言模型训练过程中使用了更多的FP8。此外，他们的系统还提供FP8低位平行性，包括张量、流水线和序列平行性。它可以实现规模化训练，例如使用数千个GPU训练的GPT-175B，打开了低精度并行训练的大门。
使用建议的FP8低精度框架对GPT风格模型的训练（包括有监督的精细调节和预训练），进行验证。将他们的FP8方法与广泛使用的BF16混合精度训练方法进行比较，实验结果显示了显著的改进，例如实际内存使用量减少了27%至42%，重要的是，重量梯度通信开销减少了63%至65%。无论是在预训练还是下游任务中，使用FP8训练的模型与利用BF16高精度的模型性能相当，而不需要对学习率和权重衰减等超参数进行任何调整。在GPT-175B模型的训练过程中，有一点值得注意，他们的FP8混合精度框架在H100 GPU平台上使用的内存比TE少21%，节省的训练时间比TE少17%。
图1：使用我们的FP8混合精度训练方法与更流行的BF16方法在一台拥有80G RAM的Nvidia H100 GPU集群上可以实现的最大模型尺寸的比较。
更重要的是，当模型规模增大时，如图1所示，使用低精度FP8所达到的成本节约可能进一步增强。为了更好地匹配预训练LLMs与最终任务和用户偏好，他们使用FP8混合精度进行指令微调和人工输入的强化学习。特别是，他们使用公开可用的用户共享的指令跟踪数据对预训练模型进行微调。在获得27%的训练速度增益的同时，使用他们的FP8混合精度调整的模型在AlpacaEval和MT-Bench基准测试中表现与使用半精度BF16的模型相似。此外，FP8混合精度在需要加载许多训练模型的RLHF过程中显示出显著的优势。
通过在训练中使用FP8，流行的RLHF框架AlpacaFarm可以在模型权重上实现46%的减少，并在优化器状态的内存使用上实现62%的减少。这更加显示了他们的FP8低精度训练架构的灵活性和适应性。以下是他们为未来一代LLMs进一步发展FP8低精度训练所做的贡献。
• 一种逐渐解锁8位权重、梯度、优化器和分布式训练的易于使用的FP8混合精度训练新框架。通过仅更改超参数和训练凭据，可以轻松将当前16/32位混合精度等价物与这个8位框架互换。他们还提供了一个PyTorch的实现，只需几行代码即可进行8位低精度训练。
• 一系列新的以FP8训练的GPT风格模型。通过将其应用于GPT预训练和微调，他们展示了所提出的FP8方案在从7B到175B参数的一系列模型尺寸上的能力。他们为流行的并行计算范式提供了FP8支持（张量、流水线和序列并行），使FP8可用于训练大型基础模型。他们基于Megatron-LM实现的第一个FP8 GPT训练代码库已公开提供。他们期待引入他们的FP8框架将为未来一代面向大基础模型的低精度训练系统提供新的标准。