FA2 + FP8训练性能 | AI生成和翻译
问题：能否解释一下在训练中结合使用FA2和FP8的性能对比？
答案：
本报告展示了在模型训练任务中，将FA2（FlashAttention 2）与FP8（8位浮点精度）相结合对性能的影响，并与其他配置进行了对比。以下是详细解析：
关键指标说明
| 指标 | 含义 |
|---|---|
| 步进时间 | 处理一个训练步骤所需的时间（数值越小越好）。 |
| 吞吐量 | 每秒处理的token数量（数值越大越好）。 |
| MFU | 多算子融合利用率，衡量GPU计算效率的指标（数值越大越好）。 |
| 内存 | 训练期间占用的GPU内存。 |
对比的配置方案
-
bf16 + SDPA（基准配置）
- 步进时间：7.7秒
- 吞吐量：68K tokens/秒
- MFU：27.5%
- 内存：105 GB
-
FP8 + SDPA
- 步进时间：6.4秒
- 吞吐量：82K tokens/秒
- MFU：33.1%
- 内存：92 GB
- 使用FP8可减少内存占用，并比bf16更快地加速训练。
-
FP8 + FA2（SSSL）
- 步进时间：5.27秒
- 吞吐量：99.5K tokens/秒
- MFU：36.4%
- 内存：92 GB
- 该组合将FP8精度与FA2（一种优化的注意力核，SSSL代表单次Softmax及其他改进）结合，在不增加内存使用的前提下实现了最快的速度和最高的GPU利用率。
加速效果分析
-
与 bf16基准配置 相比：
- 速度提升 1.46倍
- MFU提升 +8.9%
- 预估训练时间从62小时降至约42.5小时
-
与 仅使用FP8 相比：
- 速度提升 1.21倍
- MFU提升 +3.3%
- 预估训练时间从52小时降至约42.5小时
总结
- FA2 + FP8是一项突破性组合：它能缩短步进时间、提高吞吐量、提升GPU效率，同时保持较低的内存占用。
- 最大的收益来自FA2优化的注意力计算与FP8较低精度的结合，这种做法在加速训练的同时，仍能维持可接受的模型精度。
参考来源：