Accelerate 文档
低精度训练方法
并获得增强的文档体验
开始使用
低精度训练方法
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 混合精度设置的其他部分),您可以使用 utils.AORecipeKwargs
、utils.TERecipeKwargs
或 utils.MSAMPRecipeKwargs
等 RecipeKwargs
数据类之一;您也可以在配置 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)
- 完全分片数据并行 (FSDP)
- DeepSpeed ZeRO 1 到 3
在此处了解更多信息 此处
延伸阅读
要了解有关 FP8 训练的更多信息,请查看以下资源
- 我们的概念指南,详细介绍了 TransformersEngine 和 MS-AMP
transformers-engine
文档MS-AMP
文档torchao
文档