NSDI ’24: 21st USENIX Symposium on Networked Systems Design and Implementation, April 16–18, 2024, Santa Clara, CA, USA

https://www.usenix.org/system/files/nsdi24-jiang-ziheng.pdf

一、研究背景与动机：万卡训练时代的系统挑战

大语言模型（LLM）的训练规模已从数百块 GPU 跃升至数万块 GPU。以 GPT-3（1750 亿参数）和 PaLM（5400 亿参数）为代表的前沿模型，要求在数周甚至数月内持续占用整个集群进行训练。这与传统的深度学习训练任务形成了鲜明对比——过去一个 ResNet 训练任务只需要几十到几百块 GPU，而如今一个 LLM 训练任务就需要独占上万块 GPU。

在这种前所未有的规模下，系统面临两大核心挑战。第一是训练效率问题。模型 FLOPs 利用率（MFU）是衡量训练效率的标准指标，它反映了实际吞吐量与硬件理论峰值的比值。LLM 训练并非简单的数据并行，而是需要在数据并行、流水线并行和张量并行三个维度上同时分配计算和通信，任何一个环节的低效都会拉低整体 MFU。第二是训练稳定性问题。训练一个万亿 token 的模型可能持续数周，在这个时间跨度内，硬件故障和性能掉队节点（straggler）是常态而非例外。一次故障可能导致上万块 GPU 全部停滞，而回退到上一个检查点可能意味着数小时的训练进度丢失。

二、核心方法：算法与系统的全栈协同设计

MegaScale 采用了算法-系统协同设计的原则，在模型架构、通信策略、算子优化、数据管道和网络调优等层面进行了全方位优化。

在算法层面，MegaScale 引入了三项关键改进。第一是并行 Transformer 块，将传统串行执行的注意力模块和 MLP 模块改为并行执行，减少了计算时间。第二是滑动窗口注意力，将注意力计算的复杂度从 O(s×s) 降低到 O(s×w)，其中 w 远小于序列长度 s，在不损失模型精度的前提下显著提升了训练速度。第三是 LAMB 优化器，使得批量大小可以扩大 4 倍而不影响模型收敛，从而将流水线并行中的气泡（pipeline bubble）减少了 87.5%。

在通信层面，MegaScale 针对 3D 并行中的每种通信模式设计了专门的重叠策略。对于数据并行，将 AllGather 操作预取到迭代开始阶段，使其与数据加载重叠。对于流水线并行，将 Send 和 Receive 操作解耦，使发送操作能够与计算重叠。对于张量并行和序列并行，将 AllGather 和 ReduceScatter 操作融合到线性层的 GEMM 计算中，通过将 GEMM 核切分为小块并与通信进行流水线执行来隐藏通信开销。

在数据管道方面，MegaScale 采用异步数据预处理，在 GPU 同步梯度时预处理下一步的数据。同时，利用同一节点内的张量并行组输入相同的特点，采用树状数据加载方式，由单一加载器读取数据到共享内存，各 GPU 再自行拷贝，消除了冗余磁盘读取。

在网络层面，MegaScale 采用三层 CLOS 拓扑连接超过一万块 GPU，通过 multi-rail 方式将 8 块网卡连接到 8 个不同的 ToR 交换机以减少 ECMP 哈希冲突，并开发了融合 Swift 和 DCQCN 的拥塞控制算法来降低 PFC 引发的队头阻塞。

三、容错与诊断：深度可观测性驱动的稳定训练

MegaScale 构建了一套完整的容错训练框架。Driver 进程通过 Kubernetes 管理计算资源，每个 Executor 管理一个节点并定期发送心跳。心跳信息不仅包含基本状态，还封装了训练进程日志和 RDMA 流量指标。当 Driver 检测到异常或心跳超时时，会触发故障恢复流程：暂停所有训练任务，命令各节点运行轻量级自检诊断（包括环回带宽测试、RNIC 间连通性测试和 NCCL 通信测试），识别并隔离故障节点，由 Kubernetes 补充健康节点，最后从最近的检查点恢复训练。

在检查点优化方面，MegaScale 采用两阶段方法。第一阶段，每块 GPU 将状态写入主机内存（利用 PCIe 高带宽仅需几秒），然后立即继续训练。第二阶段，后台进程异步将状态传输到分布式文件系统。恢复时，利用同一数据并行组共享状态分片的特点，仅由一个 worker 从 HDFS 读取再广播给组内其他成员，线性降低了 HDFS 带宽压力。

对于难以通过自检发现的性能问题，MegaScale 开发了基于 CUDA Event 的性能分析工具。该工具以极低开销（无需 CUDA 同步）记录每个节点上关键代码段的执行时间，并提供两种可视化模式：热力图模式可以一眼识别出计算掉队节点（约占集群的 0.5%），时间线模式可以从分布式视角展示流水线执行的实际时序和数据依赖关系。此外，3D 并行训练可视化工具可以在通信超时时展示各 GPU 的逻辑拓扑位置和数据流方向，帮助快速定位阻塞源头。

四、实验评估：在 12288 块 GPU 上达到 55.2% MFU

MegaScale 在 ByteDance 内部部署的超过一万块 NVIDIA Ampere GPU 的 AI 集群上进行了全面评估。

在训练效率方面，MegaScale 在 12288 块 GPU 上训练 1750 亿参数模型时达到了 55.2% 的 MFU，相比 Megatron-LM 的 41.2% 提升了 1.34 倍。消融实验表明，各项优化的贡献分别为：并行 Transformer 块和滑动窗口注意力贡献 5.6%，3D 并行通信重叠贡献 6.2%，高效算子贡献 1.7%，数据管道等杂项优化贡献 1.1%，LAMB 优化器扩大批量贡献 3.0%。在弱扩展性实验中，MegaScale 在从 2240 到 11200 块 GPU 的范围内保持了近线性的扩展性能，MFU 稳定在 54% 左右。

在稳定性方面，论文展示了一次真实的生产训练过程：在超过一万块 GPU 上训练一个数千亿参数的模型，历时数周，loss 持续收敛。期间训练重启超过 100 次，其中 90% 以上的故障由容错框架自动识别和修复，从故障检测到恢复的平均时间不超过 10 分钟，有效训练时间率保持在 90% 以上。

五、经验教训：大规模训练中的典型问题

论文分享了若干有价值的实战经验。关于计算掉队节点，通过 CUDA Event 热力图发现约 0.5% 的机器在执行相同前向计算时耗时多出约 10%，排除这些节点后 MFU 提升了 0.7%。关于 MFU 逐步下降的问题，经排查发现根因是不规则的垃圾回收和某些 PyTorch 操作引入的性能波动，导致各 rank 发起 ReduceScatter 的时间逐渐偏移、越来越不同步，修复后 MFU 不再随步数下降。关于网络接口抖动问题，需要将 NCCL 超时阈值显式设置为较大值以避免网卡短暂断开时训练直接崩溃，同时从信号质量层面对网卡、AOC 线缆和交换机进行底层质量管控。

六、总结与意义

MegaScale 是目前公开发表的最大规模 LLM 训练系统之一，它从系统角度展示了万卡级训练的完整设计空间。其核心贡献不仅在于 55.2% MFU 这一数字本身，更在于它揭示了大规模训练中通信重叠、容错恢复和性能诊断这三个维度的系统性方法论。对于构建下一代 AI 基础设施——无论是更大规模的密集模型训练还是 MoE 模型训练——MegaScale 的全栈协同设计思路和深度可观测性原则都提供了极具价值的参考。