超越英伟达！字节跳动MegaScale如何实现大规模GPU的高效利用？

1,221 0 0

模型信息

【模型公司】 月之暗面
【模型名称】 moonshot-v1-32k
【摘要评分】 ★★★★★

文章摘要

在2024年2月，字帖跳动发布了一篇关于大语言模型训练框架的论文《MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs》。MegaScale是一个类似于英伟达Megatron-LM的开源框架，旨在提高大规模GPU的效率。该框架通过数据并行、张量并行和流水线并行来实现这一目标。在处理1750亿参数规模的语言模型时，MegaScale在12288个GPU上实现了55.2%的模型FLOPs利用率（MFU），比Megatron-LM提高了约1.34倍。为了确保训练的高效率和稳定性，团队开发了一系列诊断工具，用于监控系统组件深层事件，识别问题根源，实现容错性和减缓落后节点的影响。

MegaScale采用了“3D并行”设计，结合了数据并行、张量并行和流水线并行，减少了GPU之间的通信消耗，增加了并行计算能力。在数据并行方面，模型参数在各个worker之间分发，每个worker持有模型的一份参数副本。在前向传播阶段，通过all-gather操作收集来自其他worker的最新模型参数。在反向传播阶段，使用reduce-scatter操作聚合各worker计算出的梯度。MegaScale还借鉴了PyTorch FSDP的做法，通过提前预取初始的all-gather操作，并让它与数据加载操作重叠，减少了每次迭代中的通信时间。

在张量并行和序列并行性中，MegaScale实现了与并行transformer块结合的重叠通信。all-gather和reduce-scatter是分布式计算和深度学习中常用的通信原语，用于在并行环境中实现数据的聚合和分发。MegaScale在张量并行中基于模型块实现通信重叠，优化了首次all-gather和最后一次reduce-scatter操作的开销。在流水线并行上，采用交错1F1B调度方法，实现了计算与通信的重叠。

MegaScale的诊断工具记录了每个机器排名在运行过程中关键代码段的执行时间，基于GPU里面的CUDA事件方法计时事件，最大程度地减少了对CUDA同步的需求，防止了性能下降。该工具提供了两种可视化模式，并可以从不同的角度分析收集到的数据。颜色的深度表示在不同计算节点（rank）上代码段的运行时间。工具记录了在统一时间线上采集的事件流，当选择特定事件时，显示相关的依赖关系，尤其是在三D并行训练场景中，可以直观展现各GPU工作者间的逻辑拓扑结构、数据流向及其涉及的通信操作。若发生错误，可以直接查看选定GPU工作者的错误信息，从而辅助快速定位训练异常。

MegaScale的研究团队还设计了一个网络拓扑结构，并安排网络流量以减少ECMP散列冲突。在顶部交换机（ToR）级别，一个400G的下行端口被分成两个200G的下行端口，并使用特定的AOC电缆。服务器上的八个200G NIC以多线方式连接到八个不同的交换机。这种方法显著减少了通信所需的交换机跳数，并进一步减少了ECMP散列冲突的概率。

尽管MegaScale在设计逻辑上具有创新性，并且在12288个GPU的情况下运行效率超过Megatron-LM，但要获得市场认可，团队需要将其从实验室带到实际环境中，感受实际压力。Megatron-LM和DeepMind的Mesh-TensorFlow都是开源工具，而MegaScale目前还处于闭源阶段，没有开源。由于涉及商业合作，MegaScale可能不会开源，但这种合作项目通常需要知识产权保护。