Accelerate FP8 低精度训练完整指南：TransformersEngine、torchao 与 MS-AMP 三后端配置实战

人工智能深度学习分布式训练【免费下载链接】accelerate A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support项目地址https://gitcode.com/gh_mirrors/ac/accelerate点击查看免费下载本指南围绕 Hugging Face Accelerate 的 FP8 低精度训练集成展开讲解在支持 FP8 的 NVIDIA 硬件上如何通过Accelerator统一配置TransformersEngine、torchao以及遗留的 MS-AMP三种后端覆盖代码、config.yaml、accelerate config与accelerate launch四种配置途径。读完本文你将掌握每个后端参数的含义与默认值、与 FSDP2/torch.compile 的联动方式以及如何在单卡、DDP、FSDP 与 DeepSpeed 场景下落地 FP8 训练。FP8 训练意味着什么在传统混合精度训练如 BF16中模型训练的各个环节通常以 16 位浮点数表示。FP8 训练的核心思想是将训练的某些或全部环节改用 8 位浮点数完成从而在不明显牺牲最终精度的前提下降低显存占用并提升吞吐。FP8 并非通用能力它依赖专门硬件对 FP8 运算的原生支持。根据 低精度训练使用指南 的说明目前仅在以下 NVIDIA 硬件上启用3000 系列之后的家用消费级显卡如 4090基于 Hopper 架构的 GPU如 H100、H200。带来的收益主要包括两方面一是显存占用下降——部分训练环节的位宽从 16 位减半到 8 位二是对可以将足够多层替换为 FP8 层的较大模型吞吐量通常会提升。需要说明的是收益并非无条件出现。从 概念指南 的实证经验看只有模型中以可替换层nn.Linear、nn.LayerNorm等为主体的较大模型FP8 的性能收益才会显现通常参数规模要达到数十亿量级小模型的收益有限。关于 FP8 训练的更深层原理精度缩放、amax 追踪、E4M3/E5M2 格式取舍等可以阅读仓库内的 FP8 训练概念指南。快速开始统一的 FP8 配置入口Accelerate 对 FP8 的设计是统一入口、可插拔后端。无论选用哪个后端核心 API 完全一致只需要给Accelerator传入mixed_precisionfp8。它同样适用于accelerate config的混合精度提问环节以及config.yaml中的mixed_precision键。from accelerate import Accelerator accelerator Accelerator(mixed_precisionfp8)当只设置mixed_precisionfp8而不指定后端时Accelerate 会自动探测本机已安装的 FP8 库。从 accelerator.py 的初始化逻辑 可以确认自动选择顺序为若安装了torchao优先选用AO后端否则若安装了transformer-engine选用TE后端否则若安装了msamp选用MSAMP后端三者均未安装则抛出ImportError提示有效的后端为torchao、transformer-engine和msamp。要显式指定后端并精细控制 FP8 混合精度的行为使用RecipeKwargs一族的数据类AORecipeKwargstorchao、TERecipeKwargsTransformersEngine、MSAMPRecipeKwargsMS-AMP。它们的类定义集中在 utils/dataclasses.pyfrom accelerate import Accelerator from accelerate.utils import TERecipeKwargs, AORecipeKwargs # 使用 TransformersEngine kwargs [TERecipeKwargs()] # 或者使用 torchao # kwargs [AORecipeKwargs()] accelerator Accelerator(mixed_precisionfp8, kwarg_handlerskwargs)后端的选择也会体现在Accelerator.fp8_backend属性上其取值来自FP8BackendType枚举NO/TE/MSAMP/AO定义见 utils/dataclasses.py。在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: falsefp8_config中backend支持TE、MSAMP、AO三种取值。仓库还提供了可直接参考的完整模板 examples/config_yaml_templates/fp8.yaml其中同时注释了 TE 与 MS-AMP 两种后端的配置写法。各后端在训练环节中的位宽对比不同后端对哪些环节保持高位宽、哪些环节降为 8 位的策略差异很大。下表源自 低精度训练概念指南 引用的对比数据直观展示各方案的精度分布| 优化级别 | 计算GEMM | 通信 | 权重 | 主权重 | 权重梯度 | 优化器状态 | | -- | -- | -- | -- | -- | -- | -- | | FP16 AMP | FP16 | FP32 | FP32 | N/A | FP32 | FP32FP32 | | NVIDIA TE | FP8 | FP32 | FP32 | N/A | FP32 | FP32FP32 | | MS-AMP O1 | FP8 | FP8 | FP16 | N/A | FP8 | FP32FP32 | | MS-AMP O2 | FP8 | FP8 | FP16 | N/A | FP8 | FP8FP16 | | MS-AMP O3 | FP8 | FP8 | FP8 | FP16 | FP8 | FP8FP16 |可见 TE 只把计算GEMM降为 FP8其余保持 FP32因此显存节省最少但对最终精度的影响也最小MS-AMP 则逐级把通信、权重、梯度、优化器状态都压到 FP8/FP16显存与通信带宽收益更大。配置 TransformersEngine逐参数深入TransformersEngineTE由 NVIDIA 提供通过即插即用替换层的方式工作。Accelerate 在准备模型时会把以下层递归替换为 TE 版本见 utils/transformer_engine.pynn.LayerNorm→te.LayerNormnn.Linear→te.Linear替换过程发生在accelerator.prepare()内部源码注释指出te.LayerNorm相对nn.LayerNorm在显存与速度上并无明显收益因此实际收益主要来自nn.Linear→te.Linear的替换。从 accelerator.py 的_prepare_te实现看Accelerate 会先以torch.no_grad()完成层替换随后把优化器参数组中的旧参数映射为新参数保证优化器状态不丢失。TE 后端提供了丰富的可调参数Accelerate 在 TERecipeKwargs 的 docstring 中给出了完整说明与默认值| 参数 | 类型/取值范围 | 默认值 | 含义 | | -- | -- | -- | -- | |use_autocast_during_eval| bool |False| 评估阶段是否启用 FP8 autocast通常关闭可获得更真实的评估指标 | |margin| int |0| 梯度缩放使用的 margin | |interval| int |1| 缩放因子重计算的间隔 | |fp8_format|HYBRID/E4M3/E5M2|HYBRID| FP8 配方格式训练一般用HYBRID评估可用E4M3或E5M2| |amax_history_len| int |1024| 缩放因子计算使用的历史长度 | |amax_compute_algo|max/most_recent|most_recent| 缩放因子统计算法 | |override_linear_precision| (bool, bool, bool) |(False, False, False)| 是否让 fprop / dgrad / wgrad 三种 GEMM 以更高精度执行 | |use_mxfp8_block_scaling| bool |False| 是否启用 MXFP8 块缩放受硬件与 TE 版本支持情况限制 |注意fp8_format与amax_compute_algo在代码中会被强制upper()/lower()归一化并做严格取值校验非法值会直接抛出ValueError。这些参数同时也支持通过ACCELERATE_FP8_前缀的环境变量覆盖如ACCELERATE_FP8_FORMAT、ACCELERATE_FP8_MARGIN默认值解析逻辑见 dataclasses.py 中 TERecipeKwargs.post_init。在代码中使用 TE 后端from accelerate import Accelerator from accelerate.utils import TERecipeKwargs kwargs [TERecipeKwargs(fp8_formatHYBRID, amax_compute_algomax)] accelerator Accelerator(mixed_precisionfp8, kwarg_handlerskwargs)等价地通过accelerate launch命令行指定accelerate launch --mixed_precisionfp8 --fp8_backendte --fp8_formatHYBRID --fp8_amax_compute_algomax train.py使用accelerate launch --fp8_backendte -h可以查看该后端相关的全部可用参数。在config.yaml中配置 TE 后端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: falseTE 的 FP8 autocast 包装机制TE 后端与前两个后端的一个实现差异在于Accelerate 会把模型的前向传播包装进fp8_autocast上下文且该包装是上下文感知的——默认在eval模式下自动关闭 FP8 autocast以保证评估指标更真实。相关实现位于 utils/transformer_engine.py 的 contextual_fp8_autocast而 apply_fp8_autowrap 负责根据 handler 参数构造 TE 的DelayedScaling或启用 MXFP8 时的MXFP8BlockScaling配方并完成包装。从源码结构看在多 GPUMULTI_GPU与 FSDP 分布式类型下Accelerate 会启用延迟 FP8 autocastdelayed_fp8_autocast将包装动作推迟到分布式包装完成之后执行见 accelerator.py 与 accelerator.py单卡场景则提前包装accelerator.py。配置 torchaoFSDP2 与 torch.compile 的联动torchao是 PyTorch 官方驱动的、可二次开发的 FP8 后端比前两者更易上手。它的一个关键差异在于数值稳定性策略实践表明把模型的第一个和最后一个线性层保持在原始精度FP32 或 BF16其余层再量化到 FP8通常更有利于稳定收敛。这一点从源码与测试中都可以得到印证convert_model_to_fp8_ao会先通过find_first_last_linear_layers找出模型中第一个和最后一个nn.Linear再用filter_linear_layers把它们从 FP8 转换中排除utils/ao.py 与 utils/ao.py。测试用例 tests/test_fp8.py 专门验证了跳过首尾线性层这一行为embed_proj与lm_head保持为普通nn.Linear只有中间层block[0]被转换为Float8Linear。提示torchao的 FP8 API 目前仍处于实验阶段可能随版本变化。AORecipeKwargs的核心配置项定义见 dataclasses.pyconfig一个torchao.float8.Float8LinearConfig实例。不传时使用内置默认配置pad_inner_dimTrue把矩阵维度补齐到 16 的倍数这是torch._scaled_mm运算的硬性要求避免运行时对齐错误enable_fsdp_float8_all_gatherTrue为 FSDP2 启用 FP8 all-gather在聚合前把参数先转为 FP8相比 BF16 节省约 50% 的通信带宽。module_filter_func一个接收模块与层名的过滤函数返回布尔值决定该层是否转 FP8默认逻辑为accelerate.utils.ao.filter_linear_layers。torchao后端常与 FSDP2、torch.compileinductor组合使用。仓库文档给出的完整示例from accelerate import Accelerator from accelerate.utils import AORecipeKwargs, TorchDynamoPlugin, FullyShardedDataParallelPlugin from torchao.float8 import Float8LinearConfig fsdp2_plugin FullyShardedDataParallelPlugin( fsdp_version2, cpu_ram_efficient_loadingFalse, # CPU RAM 高效加载无法与 fp8 torchao 兼容 fsdp_auto_wrap_policyTRANSFORMER_BASED_WRAP, ) dynamo_plugin TorchDynamoPlugin( backendinductor, use_regional_compilationTrue, ) fp8_config Float8LinearConfig( enable_fsdp_float8_all_gatherTrue, # 在 FSDP2 中使用 FP8 all_gather pad_inner_dimTrue, ) kwargs [AORecipeKwargs( configfp8_config )] accelerator Accelerator( mixed_precisionfp8, fsdp_pluginfsdp2_plugin, dynamo_plugindynamo_plugin, kwarg_handlerskwargs, )对应的config.yaml写法mixed_precision: fp8 fsdp_config: fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP fsdp_cpu_ram_efficient_loading: false fsdp_version: 2 fp8_config: backend: AO pad_inner_dim: true enable_fsdp_float8_all_gather: true dynamo_config: dynamo_backend: INDUCTOR dynamo_use_regional_compilation: true命令行等价写法accelerate launch --fp8_backendao ...用accelerate launch --fp8_backendao -h查看相关参数。关于Float8LinearConfig各参数的进一步细节需要参考torchao包的官方文档。从源码看torchao后端还有两个值得注意的实现细节FSDP2 预计算缩放当使用 FSDP2 且启用enable_fsdp_float8_all_gather时Accelerate 会给优化器注册一个 step 后的钩子调用precompute_float8_dynamic_scale_for_fsdp预计算动态缩放因子见 accelerator.py 的 _prepare_ao环境变量覆盖ACCELERATE_FP8_PAD_INNER_DIM与ACCELERATE_FP8_ENABLE_FSDP_FLOAT8_ALL_GATHER可分别覆盖上述两个默认值。配置 MS-AMP遗留后端与弃用警告⚠️ 已弃用 / 不再维护MS-AMP 已不再由微软积极维护。根据 低精度训练使用指南 中的警告MS-AMP 仓库自 2023 年以来未再更新并存在已知兼容性问题需要 CUDA 11.x不支持 CUDA 12.x 及更高版本需要与新版 PyTorch 不兼容的旧版 NCCL不支持新版 PyTorch2.2。所有新建及现有 FP8 训练工作流都强烈建议改用 TransformersEngine 或 torchao两者均积极维护并支持现代 CUDA/PyTorch 版本其中基于torchao的原生 PyTorch FP8 支持尤其有希望成为厂商中立的方案。MS-AMP 后端保留在 Accelerate 中仅用于遗留兼容未来版本可能移除。Accelerate 的 MS-AMP 集成只有一个配置参数优化级别。当前支持O1与O2两个级别注意是字母 o不是数字 0O1将权重梯度与all_reduce通信转换为 8 位执行其余环节保持 16 位。主要降低 GPU 显存占用并提升通信带宽O2在 O1 基础上将一阶优化器状态也转换为 8 位二阶状态保持 FP16目前仅支持Adam优化器。在尽量不损失最终精度的前提下节省最多的显存。同时从概念指南可知MS-AMP 还存在一个 O3 级别将内存中的模型权重整体降为 FP8、主权重存为 FP16主要用于 DeepSpeed 等 DDP 场景但仅支持到 DeepSpeed 0.9.2因此未包含在 Accelerate 集成中。在代码中指定优化级别from accelerate import Accelerator from accelerate.utils import MSAMPRecipeKwargs kwargs [MSAMPRecipeKwargs(opt_levelO2)] accelerator Accelerator(mixed_precisionfp8, kwarg_handlerskwargs)通过accelerate launch指定accelerate launch --fp8_backendmsamp --fp8_opt_levelO2 train.py在config.yaml中mixed_precision: fp8 fp8_config: backend: MSAMP opt_level: O2从源码看MS-AMP 的初始化调用在 accelerator.py 的 _prepare_msamp要求模型与优化器必须同时传给prepare()且只支持单一模型、单一优化器最终调用msamp.initialize(model, optimizer, opt_levelopt_level)。此外当 MS-AMP 与 DeepSpeed 组合时Accelerate 会改用msamp.deepspeed.initialize打补丁的初始化入口accelerator.py。这段代码在触发时会发出FutureWarning提示用户迁移到 TE 或 torchao。说明FP8RecipeKwargs是早期版本遗留的统称类同时继承 TE 与 MS-AMP 参数含backend字段现已标记弃用并将在 Accelerate v2.0.0 移除新代码应改用TERecipeKwargs或MSAMPRecipeKwargs详见 dataclasses.py 中的弃用注释。accelerate launch 的 FP8 命令行参数全表launch.py 中定义了一个独立的 FP8 Arguments 参数组全部参数要求--mixed_precisionfp8才能生效| 参数 | 取值 | 默认值 | 适用后端 | 说明 | | -- | -- | -- | -- | -- | |--fp8_backend|ao/te/msamp| 自动探测 | 全部 | 选择 FP8 后端 | |--fp8_use_autocast_during_eval| 标志位 |False| te | 评估时是否启用 FP8 autocast | |--fp8_margin| int |0| te | 梯度缩放 margin | |--fp8_interval| int |1| te | 缩放因子重计算间隔 | |--fp8_format|HYBRID/E4M3/E5M2|HYBRID| te | FP8 配方格式 | |--fp8_amax_history_len| int |1024| te | amax 缩放历史长度 | |--fp8_amax_compute_algo|max/most_recent|most_recent| te | amax 统计算法 | |--fp8_override_linear_precision| 逗号分隔布尔 |(False, False, False)| te | fprop/dgrad/wgrad 是否高位宽 | |--fp8_opt_level|O1/O2|O2| msamp | MS-AMP 优化级别 | |--fp8_enable_fsdp_float8_all_gather| bool |true| ao | FSDP2 是否启用 FP8 all-gather | |--fp8_pad_inner_dim| bool |true| ao | 是否补齐矩阵内维以满足_scaled_mm对齐 |此外accelerate launch的--mixed_precision参数合法取值即为no/fp16/bf16/fp8launch.py。通过 accelerate config 交互式配置 FP8在accelerate config交互流程中当选择混合精度为fp8时cluster.py 会依次询问选择 FP8 后端ao/te/msamp三选一若选TE是否在评估模式使用 FP8 autocast默认 NOmargin默认 0interval默认 1权重格式HYBRID/E4M3/E5M2默认 HYBRIDamax 缩放因子计算的历史长度默认 1024amax 统计算法max/most_recent是否让 fprop/dgrad/wgrad GEMM 以更高精度执行默认 NO若选是还会逐项询问三个布尔若选MSAMP询问优化级别O1/O2默认 O2若选AO询问是否启用 FSDP2 FP8 all-gather默认 YES、是否补齐内维默认 YES。整个流程中Accelerate 会先检查对应后端库是否已安装is_fp8_available、is_transformer_engine_available、is_msamp_available、is_torchao_available未安装会直接报错提示。官方示例与基准验证仓库在 benchmarks/fp8 下为三种后端分别提供了可运行的基准脚本覆盖四种训练形态各自独立的脚本单卡non_distributed.pyDDP 多卡ddp.pyFSDP 全分片fsdp.pytorchao 与 transformer_engine 目录下有DeepSpeed ZeRO 1–3distrib_deepspeed.py。以 benchmarks/fp8/torchao/non_distributed.py 为例该脚本用 BERTbert-base-cased在 GLUE MRPC 任务上对比两条路径一条是直接使用torchao原生convert_to_float8_training的基线另一条是使用Accelerator(mixed_precisionfp8, kwarg_handlers[AORecipeKwargs()])的集成路径最后断言两条路径训练前后的 accuracy 与 F1 完全一致从而验证 Accelerate 集成与原生产物等价。每个后端目录还附带Dockerfile建议使用 Docker 镜像而非手动安装 FP8 库尤其是torchao。基准 READMEbenchmarks/fp8/torchao/README.md给出的运行方式非常简洁单卡直接用python运行其余形态用accelerate launch启动python non_distributed.py # 或 accelerate launch ddp.py # ddp.py / fsdp.py / distrib_deepspeed.pybenchmarks/fp8目录下的fp8_utils.py集中封装了模型、优化器、数据加载等训练工具README.md有进一步说明。测试保障FP8 集成的行为约束测试套件 tests/test_fp8.py 对三种后端分别做了分层验证其中值得关注的行为约束包括TE 与 AO 后端的单卡 / 多卡 / DeepSpeed 场景均可通过accelerate launch端到端拉起use_mxfp8_block_scaling的 MXFP8 块缩放路径有独立测试需硬件与 TE 支持单独的模块过滤测试验证了convert_model_to_fp8_ao会跳过模型首尾线性层tests/test_fp8.py。从这些测试可以推断Accelerate 对 FP8 的承诺是无论选择哪个后端、哪种分布式策略用户代码只感知mixed_precisionfp8RecipeKwargs这一层抽象底层差异由prepare()内部的_prepare_te/_prepare_ao/_prepare_msamp分别消化。总结与选型建议面对三种后端可以按以下原则选型新项目优先TransformersEngine追求 NVIDIA 硬件上的最优算子性能与torchaoPyTorch 官方、厂商中立、可二次开发且与 FSDP2 inductor 编译链路协同良好MS-AMP 仅遗留兼容存在 CUDA 11 / NCCL / PyTorch 版本兼容性风险且代码路径会发出弃用警告不建议用于新工作组合使用的可能性概念指南中提到 MS-AMP 与 TE 组合理论上可以兼顾 TE 的 FP8 算子与 MS-AMP 的显存削减但同样受 MS-AMP 停止维护的影响不建议新项目采用。配置入口方面无论是Accelerator(mixed_precisionfp8)代码级配置、config.yaml的fp8_config段落、accelerate config的交互式引导还是accelerate launch --fp8_backend...命令行参数四种方式等价且可互相转换。硬件前提是 3000 系列之后的 NVIDIA 消费卡或 Hopper 架构H100/H200GPU且需在环境中预先安装对应的 FP8 库。如需进一步深入推荐继续阅读仓库内的 FP8 训练概念指南讲解 TE/torchao/MS-AMP 的精度机理与层替换细节、kwargs handlers 参考RecipeKwargs系列类的完整 API以及 FP8 参考文档convert_model、has_transformer_engine_layers、contextual_fp8_autocast、apply_fp8_autowrap等底层工具函数。赞分享人工智能深度学习分布式训练【免费下载链接】accelerate A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support项目地址https://gitcode.com/gh_mirrors/ac/accelerate点击查看免费下载相关推荐Accelerate 低精度FP8训练指南TransformerEngine、torchao 与 MS-AMP 原理与实践Accelerate 低精度FP8训练指南TransformerEngine、torchao 与 MS AMP 原理与实践 FP8 低精度训练是新一代 N人工智能深度学习分布式训练slime 低精度训练与 Rollout 实践BF16 训练 FP8/INT4 推理的完整配置指南slime 低精度训练与 Rollout 实践BF16 训练 FP8/INT4 推理的完整配置指南 导读 本篇文章以 slime 官方文档《Low Pr人工智能大模型强化学习RLHF分布式训练Colossal-AI 混合精度训练完全指南基于 Booster 的 AMP 配置与实战Torch/Apex/Naive AMPColossal AI 混合精度训练完全指南基于 Booster 的 AMP 配置与实战Torch/Apex/Naive AMP 导读 本文以 Colos人工智能大模型分布式训练深度学习模型优化高性能计算上一篇TIS平台入门指南30分钟快速搭建你的第一个数据同步任务下一篇ik_llama.cpp PR 269 深度剖析修复 ggml_compute_forward_dup_q 的量化张量复制/转置回归创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考