Yi 系列模型训练优化最佳实践
在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 训练性能和方法论上面不断地探索创新，进一步挖掘和榨干新一代显卡的算力，达到更高的训练效率。...