低精度训练方法

Accelerate 提供了集成，可以使用指定的支持硬件通过 TransformersEngine 和 MS-AMP 包在较低精度的训练方法上进行训练。本文档将帮助您了解哪些硬件受支持、如何配置您的 加速器 以利用低精度方法以及训练时可以预期的结果。

FP8训练的含义

要深入了解使用 PyTorch 和 Accelerate 进行 FP8 训练的细节，请查看概念指南，了解为什么这可能很困难。但本质上，与其使用 BF16 进行训练，不如使用 8 位而不是 16 位来执行模型训练的某些（或所有）方面。挑战在于如何在不降低最终性能的情况下做到这一点。

此功能仅在特定的 NVIDIA 硬件上启用，具体来说：

• 3000 系列及以后的消费级显卡（例如 4090）
• 基于 Hopper 的 GPU 架构（例如 H100 和 H200）

这将带来一些内存使用量的提升（因为我们减少了一半某些训练部分所需的内存），并且对于可以将某些层替换为支持 FP8 的层的更大模型，也应该会看到吞吐量的提高。

配置加速器

目前支持两种不同的 FP8 后端（TransformersEngine 和 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 混合精度设置的其他部分），您可以使用utils.FP8RecipeKwargs 或在您的配置 yaml/在 accelerate launch 中进行澄清。

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

mixed_precision: fp8
fp8_config:
  amax_compute_algorithm: max
  amax_history_length: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: 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_algorithm: max
    amax_history_length: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: false
    use_autocast_during_eval: false

示例集

我们在 accelerate 存储库中提供了使用 accelerate 及其底层实现进行 FP8 训练的示例。目前，我们支持展示以下脚本：

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

了解更多信息，请访问此处

进一步阅读

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

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