超越英伟达！字节跳动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 本身是由大学和企业共同开发的，这种涉及商业合作的项目往往需要知识产权保护，因此不会开源。
