低精度训练方法

Accelerate 提供了与 TransformersEngine、MS-AMP 和 torchao 包的集成，以便在指定的支持硬件上使用低精度方法进行训练。本文档将指导你了解支持哪些硬件，如何配置你的 Accelerator 以利用低精度方法，以及你在训练时可以期待什么。

FP8 训练的含义

要了解更多关于使用 PyTorch 和 Accelerate 进行 FP8 训练的细节，请查看关于为什么这可能很困难的概念指南。但本质上，与在 BF16 中训练不同，训练模型的某些（或全部）方面可以使用 8 位而不是 16 位来执行。挑战在于这样做而不会降低最终性能。

这仅在特定的 NVIDIA 硬件上启用，即

• 3000 系列消费级显卡之后的任何产品（例如 4090）
• 基于 Hopper 的 GPU 架构（例如 H100 和 H200）

这将导致使用的内存有所减少（因为我们已经将训练某些部分所需的内存减少了一半），并且对于可以用支持 FP8 的层替换某些层的较大模型，吞吐量也应该有所增加。

配置 Accelerator

目前支持三种不同的 FP8 后端（TransformersEngine、torchao 和 MS-AMP），每种后端都有不同的功能和配置。

要使用其中任何一个，都使用相同的核心 API。只需将 mixed_precision="fp8" 传递给 Accelerator，或在 accelerate config 提示混合精度时输入，或作为 config.yaml 文件中 mixed_precision 键的一部分。

from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")

默认情况下，如果你的环境中可用 MS-AMP，Accelerate 将自动使用它作为后端。要自己指定它（并自定义 FP8 混合精度设置的其他部分），你可以使用 RecipeKwargs 数据类之一，例如 utils.AORecipeKwargs、utils.TERecipeKwargs 或 utils.MSAMPRecipeKwargs；你也可以在你的配置 yaml 中/在 accelerate launch 期间明确指定它。

from accelerate import Accelerator
from accelerate.utils import MSAMPRecipeKwargs
kwargs = [MSAMPRecipeKwargs()]
# Or to specify the backend as `TransformersEngine` even if MS-AMP is installed
# kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

mixed_precision: fp8
fp8_config:
  amax_compute_algo: max
  amax_history_len: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: (false, false, false)
  use_autocast_during_eval: false

配置 MS-AMP

在这两者中，MS-AMP 通常更容易配置，因为它只有一个参数：优化级别。

目前，Accelerate 集成支持两种优化级别，"O1" 和 "O2"（使用字母“o”，而不是零）。

• "O1" 将把权重梯度和 all_reduce 通信转换为 8 位进行，而其余部分则以 16 位完成。这减少了一般的 GPU 内存使用并加快了通信带宽。
• "O2" 还会将一阶优化器状态转换为 8 位，而二阶状态则为 FP16。（目前只支持 Adam 优化器）。这会尽力最小化最终的精度下降，并能节省最大的内存。

要指定优化级别，请通过设置 optimization_level 参数将其传递给 FP8KwargsHandler

from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs
kwargs = [FP8RecipeKwargs(backend="msamp", optimization_level="O2")]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

或者在 accelerate launch 期间通过 --fp8_backend=msamp --fp8_opt_level=O2

同样，这也可以在你的 config.yaml 中设置

mixed_precision: fp8
fp8_config:
    backend: MSAMP
    opt_level: O2

配置 TransformersEngine

TransformersEngine 有许多用于自定义 FP8 计算如何以及进行哪些计算的选项。支持的参数及其含义的完整列表可在 NVIDIA 的文档中找到，但为了方便起见，它们也作为 FP8KwargsHandler 的文档字符串重新陈述。

Accelerate 试图设置合理的默认值，但自己探索和调整各种参数可能会带来更好的性能。

要使用它，请指定 backend="te" 并修改你想要的任何参数，作为 kwarg 处理器的一部分。

from accelerate import Accelerator
from accelerate.utils import FP8RecipeKwargs
kwargs = [FP8RecipeKwargs(backend="te", ...)]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

或者在 accelerate launch 期间通过 --fp8_backend=te ...。使用 accelerate launch --fp8_backend=te -h 查看相关参数。

同样，这也可以在你的 config.yaml 中设置

mixed_precision: fp8
fp8_config:
    amax_compute_algo: max
    amax_history_len: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: (false, false, false)
    use_autocast_during_eval: false

配置 torchao

torchao 是一个由 PyTorch 驱动的、可修改的 FP8 后端，旨在比前两个引擎更易于使用。与前两个相比，ao 的一个核心区别是，为了数值稳定性，通常最好将模型的第一层和最后一层保持在常规精度（无论是 FP32 还是 BF16），然后将其他层量化到 FP8。因此，ao 的配置看起来有点不同。

注意：此 API 是实验性的，可能会发生变化。

from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs
kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

要了解有关要使用的具体参数的更多信息，请参阅官方 torchao 仓库。

示例集合

我们有示例展示了使用 accelerate 及其底层实现进行 FP8 训练，这些示例可在 accelerate 仓库中找到。目前我们支持的脚本展示了

• 单 GPU
• 分布式数据并行（多 GPU）
• 完全分片数据并行
• DeepSpeed ZeRO 1 到 3

在此处了解更多信息。

进一步阅读

要了解有关在 FP8 中训练的更多信息，请查看以下资源

• 我们的概念指南，详细介绍了 TransformersEngine 和 MS-AMP
• transformers-engine 文档
• MS-AMP 文档
• torchao 文档