千亿参数LLM的训练效率优化

模型训练硬件利用率影响因素

• Llama 3.1 模型训练案例：最大 405B 模型，使用 16K H100 训练集群，54 天预训练，466 次任务中断，展示了模型训练中硬件利用率的挑战。
• MFU (Model FLOPS Utilization)：模型算力利用率，衡量模型计算资源的使用效率。
• Goodput：衡量 AI 系统效率的指标，包括调度 Goodput、运行时 Goodput 和程序 Goodput。

分布式训练效率提升

• 数据并行 (Data Parallelism)：DDP (Distributed Data Parallel)，适用于模型较小、数据量较大的场景。
• 张量并行 (Tensor Parallelism)：Megatron-LM-1，均匀拆分模型，通信量大。
• 流水线并行 (Pipeline Parallelism)：GPipe (Google)、PipeDream (Microsoft)、Megatron-LM-2，通信量低，但引入 bubble 问题。
• 专家并行 (Expert Parallelism)：Switch Transformers (Google)、Megatron Expert Parallelism。
• 上下文并行 (Context Parallelism)：RingAttention，通过在线 softmax 计算块间注意力，但存在负载不均衡问题。
• 优化方法：
  • RingAttention：负载均衡版本改进基本原理。
  • SWA+CP：Sliding Window Attention + Context Parallel，解决序列长度问题，复用高性能 Attention 算子。
  • 模基共建：与 FullAttention+CP 混合使用。
  • 其它优化：通信计算并行。

FP8 训练经验分享

• FP8 训练简介：
  • E4M3：1 符号位、4 指数位、3 尾数位，范围 +/-448 和 nan。
  • E5M2：1 符号位、5 指数位、2 尾数位，范围 +/-57344 和 +/- inf。
  • 部分计算采用 FP8，前向用 E4M3，反向用 E5M2。
  • 缩放方式：Just-in-time scaling 和 Delayed scaling。
• FP8 训练实践：未详细展开。

MoE 训练经验分享

• MoE 结构：
  • Switch Transformers (Google)：细粒度专家和共享专家。
  • DeepSeekMoE：通信计算并行。
• Top-P 路由实现：Dynamic Routing in MoE Models。
• 其它优化：算子融合、细粒度重计算、自定义流水线并行层数划分等。

Goodput 优化

• Goodput 定义回顾：衡量 AI 系统效率的指标。
• 自研集群训推任务调度系统：
  • 拓扑亲和调度。
  • 故障监控与定位。
• FastCkpt：
  • 基本原理：GPU 内存到 CPU 内存同步，持续优化，分布式优化器切分。
  • 效果：每次保存 ckpt 的额外耗时趋近于 0，Goodput 提升 15%。
  • CPU 内存到磁盘异步，适配多种分布式策略。

总结与展望

• MFU：模型算力利用率。
• Goodput：
  • 分布式训练效率提升：DP/TP/PP/EP/CP 及其优化。
  • 训推任务调度系统：拓扑亲和调度、故障监控与定位。
• FP8 训练经验分享：
  • FP8 训练介绍/实践。
  • FastCkpt：基本原理、持续优化、提升效果。
• MoE 训练经验分享：
  • MoE 介绍。
  • MoE 系列优化：基本原理、持续优化、提升效果。