文章摘要
【关 键 词】 混合精度、FP8训练、加速效果、内存优化、并行训练
AI大模型开发系统Colossal-AI实现了混合精度训练的升级，支持BF16(O2)与FP8(O1)的新一代混合精度训练方案。这种升级通过FP8的独特数值表示方式，在保持精度的同时，提高了训练速度、节省了内存占用，并降低了训练成本。只需一行代码，即可在主流大型语言模型（LLM）上实现平均30%的加速效果，同时保证训练的收敛性。
低精度计算是GPU硬件发展的一个趋势，从FP32到FP16/BF16，再到Hopper系列芯片支持的FP8，计算速度越来越快，内存需求越来越低，符合大模型时代对硬件的需求。FP8混合精度训练中，scaling方案对训练结果影响最大，主要有两种方案：延迟scaling和实时scaling。延迟scaling使用过去一段时间内的scaling值来估计当前scaling，计算效率高，但可能影响收敛性。实时scaling直接使用当前张量值计算scaling，计算效率较低，但对收敛性影响小。Colossal-AI采用了实时scaling方案，实现了高性能且对训练收敛性影响较小。
在单卡H100上进行的矩阵乘法测试显示，随着矩阵维度的增大，FP8的加速效果越明显，且Colossal-AI的实现与Transformer Engine的性能几乎一致。在H100单卡上进行的LLaMA2-7B模型的混合精度训练测试中，FP8混合精度训练的loss曲线与BF16基本一致，显示出良好的收敛性。
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%的吞吐提升。
根据英伟达的报告和测试经验，FP8混合精度训练性能调优的建议包括：尽量少使用张量并行，用流水线并行代替张量并行；模型hidden size越大，加速效果越明显；矩阵乘法占比高的模型加速效果大。Colossal-AI对FP8的支持广泛，各种并行方式都能与FP8混合精度训练兼容，使用时只需在初始化plugin时开启FP8即可，无需多余的代码和AOT编译。
原文和模型
【原文链接】 阅读原文 [ 917字 | 4分钟 ]
【原文作者】 量子位
【摘要模型】 moonshot-v1-32k
【摘要评分】 ★★☆☆☆
