一行代码训练费用再降30%，AI大模型混合精度训练再升级

FP8通过其独特的数值表示方法，可以提高训练速度，节省内存占用，最终降低训练成本，同时保持一定的精度。

AI大模型开发系统Colossal-AI的再次升级混合精度训练，BF16(O2)支持主流 新一代FP8(O1)混合精度训练计划。

只需一行代码，就可以获得主流LLM模型。加速效果平均值为30%，降低相应的大模型开发成本，保证训练的收敛。

无需引入额外的手写CUDA算法，避免了长时间的AOT编译和复杂的编译环境设备。

混合精度FP8训练

GPU硬件的发展趋势是低精度计算。

从最早的FP32，到目前通用的FP16/BF16，再到Hopper系列芯片(H100, H200, FP8支持H800等，低精度计算速度越来越快，所需内存也越来越低，非常适合大模型时代硬件的需要。

当前影响FP8混合精度训练结果的最大因素是scaling方案，常见的方案有两种：

延迟scaling

即时scaling

延迟scaling选择前一段时间窗口中的scaling值来估算当前的scaling值，同时更新scaling和矩阵乘法。(gemm)整合起来。这样的计算方法效率很高，但是因为它是一个估计的scaling，所以对收敛有很大的影响。

即时scaling当前的张数值直接用于计算scaling，因此计算效率较低，但对收敛的影响较小。据英伟达报道，这两种scaling策略的计算效率差在10%以内。

Colossal-AI采用了对训练收敛影响较小的实时scaling方案，同时实现了不输其它延迟scaling的性能。

通过检测单卡H100中矩阵乘法，可以看出矩阵的维度越大，FP8的加速效果就越明显，而且Colossal-AI的实现和Transformer 如图1所示，Engine的性能几乎相同。但是Transformer Engine需要复杂的AOT编译环境配置和较长的编译时间。

△

为使实验结论更加贴近现实，Colossal-AI直接测试了主流LLM的实际训练。

先在H100单卡上测试一下，Transformermer下面的测试。 Engine (TE)其默认延迟scaling方案的选择。

与此同时进行了收敛测试，可以看到FP8混合精度训练的loss曲线与bf16基本一致，如图4所示：

△

Colossal-在H800多卡并行训练场景中，AI还测试了性能。单机8卡H800训练LLaMA2-7B，Colossal-AI ColossalossFP8对比-AI 与Torch相比，BF16的吞吐量提高了35%。 FSDP 94%的BF16吞吐量增加。

训练LLaMA2-13B，单机8卡H800，Colossal-AI ColossalossFP8对比-AI 39%的BF16吞吐量增加。

Cohererer在2机16卡H800上训练 Command-R 35B，Colossal-AI ColossalossFP8对比-AI 如图7所示，BF16的吞吐量提高了10%：

基于英伟达的报告和测试经验，对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编译。

本文来自微信微信官方账号“量子位”，作者：关注前沿技术，36氪经授权发布。

本文仅代表作者观点，版权归原创者所有，如需转载请在文中注明来源及作者名字。

免责声明：本文系转载编辑文章，仅作分享之用。如分享内容、图片侵犯到您的版权或非授权发布，请及时与我们联系进行审核处理或删除，您可以发送材料至邮箱：service@tojoy.com