低精度训练方法

新类型硬件的发布催生了能够更好地利用这些硬件的新训练范式。目前，这体现为使用诸如 TransformersEngine (TE) 或 torchao（原生 PyTorch FP8）等软件包以 8 位精度进行训练。

对于今天讨论的主题，我们建议先回顾一下低精度使用指南，因为本文档将经常引用该指南。

快速图表

下图快速展示了训练过程中每种方案的不同位精度：

TransformersEngine

TransformersEngine 是首个尝试以 8 位浮点数进行训练的方案。它通过为模型中的某些层提供可直接替换的层来实现，这些层利用其 FP8 引擎来减少位数（例如从 32 位降至 8 位），同时不会降低模型的最终精度。

具体而言，Accelerate 会查找并将以下层替换为 TransformersEngine 版本：

• 将 nn.LayerNorm 替换为 te.LayerNorm
• 将 nn.Linear 替换为 te.Linear

因此，我们最终得到的模型中，大部分层使用 BF16，而部分层使用 FP8，从而降低了部分内存占用。

据我们的经验，在使用 TransformerEngine 时，性能提升通常不会立即显现，除非模型中绝大多数层都由这两种可替换的层组成。因此，只有参数量达到数十亿或更大的模型才表现出性能提升。

TransformerEngine 可以接收许多不同的参数，用于自定义其执行 FP8 计算的方式以及计算结果的处理。完整的参数列表如下：

• margin：用于梯度缩放的边距。
• interval：用于指定重新计算缩放因子的频率间隔。
• fp8_format``: The format to use for the FP8 recipe. Must be one ofHYBRIDorE4M3. (GenerallyHYBRIDfor training,E4M3` for evaluation)
• amax_history_len：用于计算缩放因子的历史长度
• amax_compute_algo：用于计算缩放因子的算法。必须是 max 或 most_recent 之一。
• override_linear_precision：是否以更高精度执行 fprop、dgrad 和 wgrad GEMM。

您可以自定义 utils.TERecipeKwargs 中的每一项，以优化模型的性能。

如果我们回顾前文提到的图表，TE 仅将计算层转换为 FP8，而其余部分均保持为 FP32。因此，这种方式占用的内存最多，但其优势在于能够保证训练过程中最终精度的损失最小。