在Yi系列模型训练的过程中,提升训练吞吐和稳定性一直是我们持续优化的目标。通过这些训练相关的性能优化,我们在Yi模型训练的吞吐和稳定性得到了一个可观的效果,提升了整体的模型训练效率和迭代速度。我们具体的优化手段包括:

  • 在模型设计阶段,联合算法专家,通过大量的自动化试验寻找最优模型结构和超参设置。
  • 算子计算维度,针对通过编译优化、手写融合算子和引入开源高性能算子等手段提升模型的计算效率。
  • 在分布式通信维度,通过针对 TP/PP/DP/EP/CP 等不同分布式策略的通信计算并行,减少通信带来的开销。
  • 在集群调度维度,根据集群节点分布情况分配最佳的训练节点,并结合训练框架调整分布式并行策略的调度优先级,达到最佳的通信效率。
  • 在模型权重保存和加载维度,通过异步存储方式隐藏保存的开销,并根据检测数据并行维度是否变化,智能优化保存和加载过程中不必要的通信开销。
  • 在模型训练过程中,通过自动化的定位故障硬件,实现单点故障的隔离和任务自动重启,做到了无需人工干预。
  • 在故障恢复阶段,优化冗余的数据集加载过程,并优化通信启动开销,保证了快速的训练恢复。

通过一系列的优化,Yi 系列模型 BF16 精度的训练中达到了均值 58% 和峰值 63% 的 MFU 性能。但是在千亿模型 Yi-Large 的前期设计阶段,我们发现不同的模型架构对 MFU影响很大,例如 GPT 架构相对 OpenLLaMA 架构而言对于计算更加友好,MFU 也会更高。我们针对千亿模型模型架构设计了相同量级参数,但是不同超参的多组模型,发现瘦长的模型效果更好但是 MFU 相对较低。为了提高训练效率,我们开始在训练工程上下功夫。 训练千亿大模型是一个庞大的系统工程,需要强大的 GPU 服务器、GPU-Direct RDMA 网络高速互联和海量存储等支撑,每一环都至关重要。完整的几万亿 tokens 数据集预训练,需要数月的时间,训练成本极大,因此需要把训练的各个环节都优化到极致,达到最佳的训练吞吐。

Hopper 架构的 GPU 服务器为大模型训练提供了强大的 TensorCore 计算能力和 NVLink 互联,特别是 FP8 TensorCore 计算是这一代架构独特的能力,相比 BF16 能够提供更加强大的算力。因此我们决定在千亿模型预训练中采用 FP8 训练,并联合 NVIDIA 开发人员一起不断地攻克和解决 FP8 链路中遇到的种种难题,最终端到端的打通了 FP8 大模型训练和推理链路,并通过训练精度容错方案,在精度和评测性能保持对齐。 由于 FP8 的理论 TFLOPs 是 BF16 的两倍,训练初期切换到 FP8 的 MFU 比较低,经过在训练多个维度的性能优化,相比 BF16 训练千亿大模型训练的吞吐提升 30%,硬件 MFU 达到 30%+,提高了训练效率。未来我们也会在 FP8 训练性能和方法论上面不断地探索创新,进一步挖掘和榨干新一代显卡的算力,达到更高的训练效率。

除此之外,我们判断可以在如下技术方向继续深度优化,提升训练有效吞吐。

  • 训练局部重启,GPU 和 IB 网络等错误需要将数千卡任务整体重启,基于 NCCL 通信库和 Pytorch 训练框架的深度定制,可以实现训练任务的局部替换和重启,加速任务故障恢复速度。
  • 故障预测能力建设,对包括 GPU、光模块和光纤跳线等硬件设备做数字画像,基于他们的历史温度、性能和吞吐等特征建模,预测有可能损坏的硬件模块并提前更换。
  • 探索更多的 FP8 量化算法和策略,更加自适应的解决不同场景精度问题,减少人为干预和修正。并尝试更多模块的 FP8 计算与通信,提升训练整体FP8的优化占比,提升 FP8 TensorCore 的硬件利用率。
  • 算法和工程联合优化,在保证效果情况下,设计对硬件性能更友好的模型结构。
  • 面向 MoE 超大规模模型训练的集群调度、工程优化和算法设计,达到最高效的 MoE 训练效率。

延伸阅读