FP8 通过其独特的数值表示方式，能够在保持一定精度的同时，在大模型训练中提高训练速度、节省内存占用，最终降低训练成本。
AI 大模型开发系统 Colossal-AI 的混合精度训练再度升级，支持主流的 BF16(O2) + FP8(O1) 的新一代混合精度训练方案。
仅需一行代码，即可对主流 LLM 模型能够获得平均 30% 的加速效果，降低相应大模型开发成本，并保证训练收敛性。
无需引入额外的手写 CUDA 算子，避免了较长的 AOT 编译时间和复杂的编译环境配置。
FP8 混合精度训练
低精度计算一直是 GPU 硬件发展趋势。
从最早的 FP32，到目前通用的 FP16/BF16，再到 Hopper 系列芯片 (H100, H200, H800 等）支持的 FP8，低精度计算速度越来越快，所需的内存也越来越低，非常符合大模型时代对硬件的需求。
目前 FP8 混合精度训练影响训练结果的最大因素就是 scaling 方案，常见的方案有两种：
-
延迟 scaling
-
实时 scaling
延迟 scaling 采用之前一段时间窗口内的 scaling 值来估计当前 scaling，同时将 scaling 的更新和矩阵乘法 (gemm) 融合起来。这种计算方法效率较高，但由于是估算的 scaling，所以对收敛性影响较大。
实时 scaling 直接采用当前的张量值来计算 scaling，所以计算效率较低，但是对收敛性影响较小。根据英伟达的报告，这两种 scaling 方案的计算效率差距在 10% 以内。
Colossal-AI 采用了对训练收敛性影响较小的实时 scaling 方案，同时实现有着不输其他延迟 scaling 实现的性能。
在单卡 H100 上对矩阵乘法进行的测试，可以看到矩阵的维度越大，FP8 的加速效果越明显，而且 Colossal-AI 的实现与 Transformer Engine 的性能几乎一致，如图 1 所示。但 Transformer Engine 需要复杂的 AOT 编译环境配置和较长的编译时间。
△图 1. 单卡 GEMM 性能测试
为了实验结果更贴近现实，Colossal-AI 直接在主流 LLM 上进行了实际训练的测试。
首先在 H100 单卡上进行了测试，以下测试中 Transformer Engine (TE) 采用的其默认的延迟 scaling 方案。
同时进行了收敛性测试，可以看到 FP8 混合精度训练的 loss 曲线与 bf16 的基本一致，如图 4 所示：
△图 4. H100 单卡 LLaMA2-7B 混合精度训练 loss 曲线
Colossal-AI 还测试了 H800 多卡并行训练场景下的性能。在单机 8 卡 H800 上训练 LLaMA2-7B，Colossal-AI FP8 对比 Colossal-AI BF16 有 35% 的吞吐提升，对比 Torch FSDP BF16 有 94% 的吞吐提升。
在单机 8 卡 H800 上训练 LLaMA2-13B，Colossal-AI FP8 对比 Colossal-AI BF16 有 39% 的吞吐提升。
在 2 机 16 卡 H800 上训练 Cohere Command-R 35B，Colossal-AI FP8 对比 Colossal-AI BF16 有 10% 的吞吐提升，如图 7 所示：
根据英伟达的报告和测试经验，对 FP8 混合精度训练性能调优有一些初步的认识：
-
尽量少使用张量并行，用流水线并行代替张量并行
-
模型 hidden size 越大，加速效果越明显
-
矩阵乘法占比高的模型加速效果大
由于上述实验中 Command-R 35B 采用了张量并行，所以加速效果不太明显。
Colossal-AI 对 FP8 的支持较为广泛，各种并行方式都能和 FP8 混合精度训练兼容。使用时，仅需在初始化 plugin 时开启 FP8 即可：
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin
...
plugin = LowLevelZeroPlugin(..., use_fp8=True)
plugin = GeminiPlugin(..., use_fp8=True)
plugin = HybridParallelPlugin(..., use_fp8=True)
除此之外，无需多余的代码和 AOT 编译。
- 本文系量子位获授权刊载，观点仅为作者所有。
— 完 —
