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

在2024年的2月末。字帖跳动发布了一篇论文，叫做《MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs》。标题中的MegaScale指的是是一个大语言模型的生产框架，类似于英伟达的开源框架Megatron-LM。

英伟达这套框架主要是通过数据并行（data parallelism）、张量并行（tensor parallelism）、流水线并行（pipeline parallelism）这三点来提高大规模GPU的效率。我们曾在硅星人的文章《Gemini背后，谷歌真正可怕之处并不在模型本身……》中科普到，谷歌提出了一个概念叫做MFU，全称为模型FLOPs利用率（Model FLOPs Utilization），数字越高代表训练速度越快。

MegaScale在处理1750亿参数规模的语言模型时，在12288个GPU上实现了55.2%的MFU，相比于Megatron-LM提高了大约1.34倍。为了保证训练过程中的高效率和稳定性，团队特别注重深入观测性，并开发了一系列诊断工具，用于监控系统组件深层事件，识别问题根源，从而实现容错性和减缓落后节点的影响。

当选择了一个事件时，其依赖关系就会显现出来

那么接下来我们就来仔细拆解一下，说说MegaScale主要的实现方法以及这个“诊断工具”。MegaScale使用了一种叫做“3D并行”的设计。我们刚才说了，Megatron-LM使用了三种并行手段，而MegaScale则是将这三种并行合在了一起。一句话概括一下这么做的好处：减少GPU之间的通信消耗，进而增加并行计算能力。

在数据并行方面，MegaScale将模型参数在数据并行度（data parallel rank）内的各个worker之间进行分发，每个worker持有模型的一份参数副本。在前向传播阶段，通过all-gather操作收集来自其他worker的最新模型参数。在反向传播阶段，使用reduce-scatter操作聚合各worker计算出的梯度。MegaScale借鉴了PyTorch FSDP的做法，通过提前预取初始的all-gather操作，并让它与数据加载操作重叠。这么做最大的好处是减少了每次迭代中的通信时间。

在张量并行性和序列并行性中实现与并行transformer块结合的重叠通信

所谓all-gather和reduce-scatter，其实是分布式计算和深度学习中常用的通信原语，主要用于在并行环境中实现数据的聚合和分发。在深度学习训练的背景下，all-gather常用于数据并行场景中，当模型在多个GPU或节点上并行计算时，每个GPU计算完各自的梯度或者更新后的模型参数后，需要通过all-gather操作把所有节点上的梯度或参数聚合到每个节点上，以便后续进行全局梯度平均或参数同步。

那reduce-scatter则是用于梯度的聚合。当各个节点完成了本地的梯度计算后，通过reduce-scatter操作将所有节点的梯度汇总，然后将汇总得到的全局梯度均等地分散到所有节点，这样每个节点都可以用同一个全局梯度来更新自己的模型参数。

可能你还会有点晕，不过别急，等我把这三种并行计算都梳理完，我会用很一个故事来解释。在张量并行中，MegaScale基于模型块的基础来实现一种通信重叠。单个设备可能托管多个模型片段，对每个片段独立进行前向传播和反向传播，并在此基础上安排通信操作。all-gather操作在模型片段的前向传播之前触发，而reduce-scatter操作在其反向传播之后开始。针对无法隐藏首次all-gather和最后一次reduce-scatter操作的情况，MegaScale采用优化手段尽可能减少这部分开销。

在流水线并行上，MegaScale采用交错1F1B调度方法，在热身阶段（warm-up phase），前向传播仅依赖于之前的receive操作，因此可以将send和receive解耦合，打破原本可能存在的阻塞关系，使得send操作可以在等待compute的同时进行，实现计算与通信的重叠。到了冷却阶段（cool-down phase）则是热身阶段的逆过程，同样可以应用类似的技术。

1F1B

接下来我换种简单的方式解释，想象一下，老师罚你抄课文，但是一个人抄太慢了，于是你把这本书分成许多页，让你的同学们每个人抄一部分，这就是数据并行的基本概念。如果你发现一本书里的某一页太大，一个人都抄不完，这时你可以选择横向拆分这张大纸，每人抄一部分，再拼接回去，这就是张量并行。然后老师看你抄课文这么快，于是罚你不仅要抄课文，还要画插图，而且不同环节的工作必须有序衔接。为此，你组织了一个流水线作业，每个人负责一个工序，比如一人打草稿，另一人描边，第三人上色，这样就构成了流水线并行。

而3D并行呢，就是你对此次抄课文事情的统筹规划工作，让系统能够有效平衡计算资源的利用和通信成本，解决单个GPU无法承载超大规模模型的问题，从而在保证训练效率的同时，大幅度提升了模型训练的可扩展性。

再说说MegaScale的诊断工具，它用来记录了每个机器排名在运行过程中关键代码段的执行时间。与之前的工具（如torch分析器或Megatron-LM计时器）不同，这个诊断工具的作用是基于GPU里面的CUDA事件方法计时事件。这种方法最大程度地减少了对CUDA同步的需求，从而防止了性能下降，且允许在生产训练作业中始终运行它。该工具提供了两种可视化模式，并可以从不同的角度分析收集到的数据。

颜色的深度表示在不同计算节点（rank）上代码段的运行时间

工具记录了在统一时间线上采集的事件流，当选择特定事件时，显示相关的依赖关系，尤其是在三D并行训练场景中，可以直观展现各GPU工作者间的逻辑拓扑结构、数据流向及其涉及的通信操作。若发生错误，可以直接查看选定GPU工作者的错误信息，从而辅助快速定位训练异常。而当某个GPU工作者因为执行NCCL通信操作故障导致整个集群陷入阻塞状态时，该工具通过让每个GPU在通信超时时记录自身的当前事件，据此构建基于3D并行设置逻辑拓扑的数据依赖可视化表示，迅速锁定问题节点。

通过CUDA事件监控器记录每个计算节点上关键代码段的执行时间，对比分析不同机器之间的性能差异。该工具利用CUDA事件方法精准测量，规避了过多CUDA同步造成的性能损失，可直接应用于生产环境中的训练任务。可视化模式如热力图清晰展示各设备间的计算阶段耗时差异，帮助找出拖慢训练速度的“落后者”。

是的，即便是相同的GPU，他们在不同的算法中表现是不同的。就好像同样接受9年制义务教育，有人是考上了清华北大，而我只记得我怎么被老师罚抄课文一样。尤其是在超过1万颗GPU并行的场合下，通过诊断工具能快速发现问题，进而提高整体并行计算的效率。

而且团队考虑的非常周全，甚至连交换机都要“压榨”。研究团队专门设计了一个网络拓扑结构，并安排网络流量以减少ECMP散列冲突。首先，在顶部交换机（ToR）级别，一个400G的下行端口被分成两个200G的下行端口，并使用特定的AOC电缆。由于每个上行链路的带宽是下行链路的两倍，冲突概率降低了。其次，服务器上的八个200G NIC以多线方式连接到八个不同的交换机。通过相同的ToR交换机集合连接的GPU服务器数量可以达到64个。我们策略性地安排了来自我们训练任务的数据密集节点在同一个顶部交换机（ToR）下运行。这种方法显著减少了通信所需的交换机跳数，并进一步减少了ECMP散列冲突的概率。

MegaScale带给我们怎样的启示？

从设计逻辑上来说，MegaScale的确是很有新意。同时他们也验证了12288个GPU的情况下，框架的运行效率远远超过2020年的Megatron-LM。但问题就在于，我们熟知的大语言模型，比如ChatGPT4，它至少得有20000个英伟达A100，而GPT5估计大约要使用50000个H100的算力。MegaScale如果想要获得市场认可，团队需要带着它走出实验室，感受实际环境的压力。

Megatron-LM在2024年3月4日仍在更新

此外Megatron-LM也好，DeepMind的Mesh-TensorFlow也好，他们都是开源工具，而且Megatron-LM甚至在我写这篇稿子的时候（2024年3月4日），还在更新，内容是给ChatGPT模型添加--rotary-base参数（这个参数通常用于指定基准旋转角度，可以指定旋转操作相对于特定的基准角度进行，而不是相对于默认的基准角度，通过这个参数，可以指定旋转操作相对于特定的基准角度进行，而不是相对于默认的基准角度）。MegaScale目前还只是停留在实验室闭源阶段，没有开源。当然，由于是MegaScale本身是由大学和企业共同开发的，这种涉及商业合作的项目往往需要知识产权保护，因此不会开源。