MegaScale：将大语言模型训练扩展至超过10,000个GPU

作者/机构: Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, Yulu Jia, Sun He, Hongmin Chen, Zhihao Bai, Qi Hou, Shipeng Yan, Ding Zhou, Yiyao Sheng, Zhuo Jiang, Haohan Xu, Haoran Wei, Zhang Zhang, Pengfei Nie, Leqi Zou, Sida Zhao, Liang Xiang, Zherui Liu, Zhe Li, Xiaoying Jia, Jianxi Ye, Xin Jin, Xin Liu
(1ByteDance, 2Peking University)

A1. 主要贡献

本文介绍了MegaScale的设计、实现和工程经验，这是一个用于在超过10,000个GPU规模上训练大语言模型（LLM）的生产系统。在如此大的规模下训练LLM给训练效率和稳定性带来了前所未有的挑战。

  • 研究目标与核心挑战:

    1. 提升大规模训练效率: 在上万GPU规模下，通信开销、算子优化、数据预处理和GPU内存消耗等因素严重影响模型FLOPs利用率（MFU）。目标是最大化MFU以加速端到端训练。
    2. 保障大规模训练稳定性: LLM训练周期长达数周，大规模集群中的硬件和软件故障成为常态。故障和掉队者（straggler）不仅影响自身，还会拖慢整个万卡级别的训练任务，因此必须减少恢复时间并缓解掉队者问题。

  • 核心创新与贡献:
    MegaScale采用了算法-系统协同设计和深度可观测性两大系统原则来应对上述挑战。

    1. 全栈算法-系统协同设计:
      • 算法层面: 采用了并行Transformer块、滑动窗口注意力和LAMB优化器等技术，在不牺牲模型精度的前提下提升效率。
      • 并行与通信: 结合了数据并行、流水线并行、张量并行和序列并行，并针对每种并行策略设计了定制化的通信与计算重叠技术。
      • 数据流水线: 通过预取和树状加载优化数据流水线。
      • 通信初始化: 利用非阻塞异步操作和消除全局屏障来加速大规模集体通信组的初始化。
      • 网络性能: 设计了定制的网络拓扑，减少ECMP哈希冲突，定制拥塞控制算法，并调整重传超时参数以获得高性能网络。
    2. 深度可观测性与稳定性保障:
      • 诊断工具集: 开发了一套诊断工具，通过监控系统堆栈深处的组件和事件，以多维度视角揭示性能问题，从而定位稳定性问题的根本原因。
      • 鲁棒的训练框架: 自动化故障定位与恢复。通过封装了多种信息的心跳消息进行实时异常检测和预警，并设计了一套诊断测试来识别故障节点。
      • 性能分析工具: 开发了用于记录细粒度CUDA事件的性能分析工具，可生成系统级的热力图和分布式视图的时间线追踪，以及一个3D并行训练可视化工具，用于展示数据依赖关系，以诊断复杂的掉队者问题。
      • 优化检查点与恢复: 优化了检查点和恢复流程，以减少中断时间。

通过这些优化，MegaScale在12,288个GPU上训练175B模型时，实现了55.2%的MFU，相比Megatron-LM提升了1.34倍。

A3. 背景知识

LLM的训练因其巨大的模型架构和海量数据集而计算密集。并行策略将训练过程分布到多个设备上。

数据并行。该策略在多个设备上复制模型和优化器状态，并将数据均匀分配给所有设备。每个模型副本并行执行前向和后向传播计算。每次迭代完成后，所有模型副本同步更新模型。Zero Redundancy Optimizer (ZeRO)【索引11，Zero: Memory optimizations toward training trillion parameter models. ArXiv, 2020】不是复制模型状态（如优化器状态、梯度和参数），而是在所有数据并行进程中分片这些状态。因此，传统的聚合梯度的all-reduce操作被分解为独立的reduce-scatter和all-gather操作。这是因为每个数据并行进程只保留总状态的一部分。ZeRO分为三个递增的优化阶段。其中，第二阶段常被用于分片优化器状态和梯度，同时确保不引入额外的通信开销。
图片描述
Figure 1: 使用ZeRO2的数据并行训练。

流水线并行。该策略将模型层分布在多个设备上，每个设备拥有一部分模型。同时，每个训练批次被细分为多个微批次（micro-batches）进行流水线执行。为了减少流水线气泡，提出了各种流水线调度策略，例如GPipe【索引12，GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism , NeurIPS, 2019】、PipeDream 1F1B【索引13，Pipedream: Generalized pipeline parallelism for dnn training, ACM SOSP, 2019】等。Megatron-LM【索引7，Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】采用了交错式1F1B调度（interleaved 1F1B）。每个工作节点（worker）上的流水线阶段被细分为多个虚拟阶段，每个虚拟阶段代表一个层的子集，称为模型块（model chunk）。最初，工作节点进入预热（warm-up）阶段，为有限数量的在途微批次执行前向传播。预热后，每个工作节点进入稳定（steady）阶段，执行一次前向传播紧接着一次后向传播，常简写为1F1B。在一个批次结束时，工作节点在冷却（cool-down）阶段完成任何剩余在途微批次的后向传播。
图片描述
Figure 2: 交错式1F1B流水线。图中展示了一个三阶段流水线，每个阶段又被分为两个虚拟阶段。

张量并行。该策略将单个算子分布在多个设备上，每个设备并行执行一部分计算。根据具体的分区策略及其与模型中前后算子的关系，分区可能需要在参与的GPU之间进行通信以分割输入然后合并输出。例如，我们可以在多个GPU之间分割MLP和自注意力块中的GEMM，以利用更多的计算单元。其他一些操作如LayerNorm和Dropout计算量较小，但需要大量激活内存。另一种称为序列并行的张量并行形式被提出来，沿着序列维度分布这些算子，以有效减少激活内存占用。

并行策略的组合。这些并行策略可以组合成3D并行，以在众多GPU上扩展LLM的训练【索引10，Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020】。鉴于张量并行相关的高通信开销，最好将此类通信限制在单个集群节点内。相反，数据并行和流水线并行更适合节点间通信。在这种情况下，我们选择优先构建数据并行组而非流水线并行组，这可以减轻数据并行的跨微型舱（cross-minipod）通信。

A2. 方法细节

3. 高效的大规模训练

3.1 算法优化

并行Transformer块。我们采用并行版本的Transformer块来替代标准的串行公式。具体来说，标准的Transformer块公式可以从
$y = x + \text{Attention}(\text{LayerNorm}(x))$
$z = y + \text{MLP}(\text{LayerNorm}(y))$
重构为
$z = x + \text{Attention}(\text{LayerNorm}(x)) + \text{MLP}(\text{LayerNorm}(x))$
通过这种方法，注意力块和MLP块的计算可以并行执行，从而减少计算时间。先前的工作【索引5，PaLM: Scaling language modeling with pathways, 2022】表明，这种修改不会降低拥有数千亿参数的模型的质量。我们采用的并行版本源于【索引14，GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https:// github.com/kingoflolz/mesh-transformer-jax, 2021】。

滑动窗口注意力 (SWA)。滑动窗口注意力【索引8，Longformer: The long-document transformer, 2020】是一种稀疏注意力机制，它围绕输入序列中的每个token使用一个固定大小的窗口。其计算复杂度为 $O(s \times w)$，其中 $s$ 是输入序列长度， $w$ 是固定的窗口大小。考虑到 $w \ll s$，滑动窗口注意力比全自注意力（计算复杂度为 $O(s \times s)$）更高效。过去的工作【索引8】和我们的微基准测试（§6.2）已经表明，通过堆叠多层这样的窗口化注意力，可以形成一个大的感受野，从而保留整个输入的信息。这使得训练可以更快，而不会损害准确性。

LAMB优化器。大规模高效训练常常受到批次大小的限制，因为增加批次大小可能会对模型收敛产生不利影响。LAMB优化器【索引9，Large batch optimization for deep learning: Training bert in 76 minutes, International Conference on Learning Representations, 2020】已被证明可以在不牺牲准确性的情况下，将BERT的训练批次大小扩展到64K。在LLM的场景中，我们的实验发现LAMB可以将批次大小扩展到4倍而没有准确性损失。在交错式流水线并行中，原始调度在使用1倍批次大小训练四步时包含 $4(p-1)$ 个流水线气泡【索引7】，而使用4倍批次大小训练一步的流水线气泡是 $\frac{1}{v}(p-1) \frac{4}{m}$。因此，MegaScale通过LAMB优化器减少了87.5%的流水线气泡。

3.2 3D并行中的通信重叠

数据并行中的重叠。如图1所示，对于数据并行，存在两个主要的通信操作。一个是all-gather操作，它在前向传播过程中从其他数据并行等级（rank）的工作节点获取最新的模型参数。另一个是reduce-scatter操作，它在后向传播中收集梯度。在3D并行中，单个设备可能托管多个模型块。重叠是以模型块为基础实现的，以最大化带宽利用。all-gather操作在一个模型块的前向传播之前触发，而reduce-scatter操作在其后向传播之后开始。这导致一个挑战，即第一个all-gather操作和最后一个reduce-scatter操作无法被隐藏。受PyTorch FSDP【索引15，Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023】的启发，初始的all-gather操作在每次迭代开始时被预取，使其能够与数据加载操作重叠，从而将通信时间有效减少 $1/(2 \times vpp_size)$。我们还优先启动高优先级的通信，通信算子的优先级由依赖该通信结果的相应计算算子的顺序决定。

流水线并行中的重叠。流水线并行以点对点（point-to-point）的发送/接收（send/receive）通信为特点。MegaScale使用第2节中提到的交错式1F1B调度方法。我们注意到，在预热阶段，前向传播只依赖于其之前的接收操作。因此，我们解耦了通常一起实现且可能被较慢一方阻塞的发送和接收操作。通过打破这种依赖关系，我们使发送操作能够与计算重叠，如Figure 4左侧所示。冷却阶段可以看作是预热阶段的逆过程，允许反向应用相同的技术。至于稳定阶段，前向和后向计算都独立于相邻的通信操作。以反向传播为例，如Figure 4右侧所示，其之前的接收是为下一次前向计算服务的，而发送是为前一阶段的后向计算服务的。因此，发送和接收操作可以被异步启动以与计算重叠。
图片描述
Figure 4: 流水线并行中的通信重叠。

张量/序列并行中的重叠。张量并行通常用于划分计算密集型操作中的权重，而像LayerNorm和Dropout这样的操作则沿序列维度进行划分以节省GPU内存。这需要all-gather和reduce-scatter操作来收集输入和重新分配输出。Figure 3a展示了在并行Transformer块架构中的这种通信模式，其中两个通信算子位于关键路径上。为了消除这个开销，我们选择将all-gather和reduce-scatter与FFN路径上的并行线性层（parallel Linears）进行融合（Figure 3b）。由于FFN路径上的GEMM核更大，通信可以被更好地隐藏。我们将GEMM核分解成小块，并与通信进行流水线执行（Figure 3c）。这个策略同样可以应用于后向传播。
图片描述
Figure 3: 在并行Transformer块（PTB）中重叠张量并行（TP）和序列并行（SP）的通信。

3.3 高效算子

算子融合与优化。尽管Megatron-LM中对GEMM算子进行了优化，我们发现在其他算子中仍有进一步提升的机会。对于注意力部分，我们采用了FlashAttention-2【索引16，Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】，它改进了不同线程块和warp之间的工作划分。对于LayerNorm和GeLU，我们观察到在以前的实现中它们是由细粒度的核（kernel）组成的。通过将这些核融合在一起，我们减少了启动多个核相关的开销，并有助于优化内存访问模式，从而实现更好的性能。

3.4 数据流水线

异步数据预处理。数据预处理不在关键路径上。因此，当GPU工作节点在每个训练步骤结束时同步梯度时，可以开始下一步的数据预处理，从而隐藏了预处理的开销。

消除冗余数据加载器。在分布式训练的典型数据加载阶段，每个GPU工作节点都配备有自己的数据加载器，负责将训练数据读入CPU内存，然后再转发到GPU。这导致工作节点之间对磁盘读取带宽的竞争，从而产生瓶颈。我们观察到，在LLM训练设置中，同一台机器内的GPU工作节点属于同一个张量并行组，因此它们每次迭代的输入是相同的。基于此观察，我们采用了一种两层树状方法。我们在每台机器上使用一个专用的数据加载器将训练数据读入一块共享内存中。随后，每个GPU工作节点负责将所需数据复制到自己的GPU内存中。这消除了冗余读取，并显著提高了数据传输的效率。

3.5 集体通信组初始化

加速通信组初始化。在分布式训练中，初始化阶段涉及在GPU工作节点之间建立NVIDIA集体通信库（NCCL）通信组。由于在小规模场景中这个开销相对可以忽略，默认使用torch.distributed。当GPU数量扩展到上万个时，朴素实现引入的开销变得无法容忍。我们在§6的同一AI集群上进行的实验测量表明，Megatron-LM在2048个NVIDIA Ampere GPU上的初始化时间约为1047秒。这不仅对例行测试和迭代开发（如超参数调整和调试中的微小代码调整）构成了重大障碍，也妨碍了快速重启和恢复机制的实现。

定位并解决初始化瓶颈。为了解决这个问题，我们对torch.distributed【索引17，Pytorch distributed: Experiences on accelerating data parallel training, 2020】进行了详细的性能分析，并确定了导致初始化时间过长的两个主要原因。第一个问题在于同步步骤，每个进程在初始化特定通信组结束时都会参与一个屏障（barrier）操作。这个屏障使用TCPStore，这是Pytorch内部一个单线程、阻塞读写的分布式键值存储实现。我们用非阻塞和异步的Redis替换了TCPStore，这将2048个GPU上的初始化时间减少到361秒。第二个问题与滥用全局屏障有关。每个进程在初始化其对应的通信组后都会执行一个全局屏障。我们精心设计了通信组的初始化顺序，以最小化对全局屏障的需求。这种方法将全局屏障的时间复杂度从$O(n^2)$降低到$O(n)$。通过这些优化，2048个GPU上的初始化时间减少到5秒以下，而在超过10,000个GPU上则减少到30秒以下。

3.6 网络性能调优

网络拓扑。我们的数据中心网络采用基于Broadcom Tomahawk 4芯片的高性能交换机构建。每个Tomahawk芯片的总带宽为25.6Tbps，拥有64个400Gbps端口。三层交换机以类CLOS拓扑连接，以连接超过10,000个GPU。在每一层的交换机中，下行和上行链路的带宽比例为1:1，即32个端口用作下行，32个用作上行。该网络提供了高带宽和小直径的特性，任何节点都可以在有限的跳数内与其他节点通信。

减少ECMP哈希冲突。我们精心设计了网络拓扑并调度网络流量以减少ECMP哈希冲突。首先，在机架顶（ToR）交换机层面，一个400G的下行端口通过特定的AOC线缆被分成两个200G的下行端口。由于每个上行链路的带宽是下行链路的两倍，冲突概率降低了。其次，服务器上的八个200G网卡以多轨（multi-rail）方式连接到八个不同的交换机。由同一组ToR交换机连接的GPU服务器数量可以达到64个。我们策略性地将我们训练任务中数据密集型的节点调度到同一个ToR交换机下运行，这种方法显著减少了通信所需的交换机跳数，并进一步降低了ECMP哈希冲突的概率。

拥塞控制。在分布式训练中，all-to-all通信可能会导致拥塞，并且在大规模使用默认的DCQCN【索引19，Congestion Control for Large-scale RDMA Deployments, ACM SIGCOMM Computer Communication Review, 2015】协议时会产生高水平的优先级流量控制（PFC）【索引18，Ieee 802.1 qbb - priority-based flow control, 2009】。过度使用PFC会导致队头阻塞（Head-of-Line blocking）【索引19】，从而降低网络吞吐量。为了缓解这些问题，我们开发了一种结合了Swift【索引20，Swift: Delay is Simple and Effective for Congestion Control in the Datacenter, SIGCOMM, 2020】和DCQCN原理的算法，它将往返时延（RTT）的精确测量与显式拥塞通知（ECN）的快速拥塞响应能力相结合。这种方法显著提高了吞吐量并最小化了与PFC相关的拥塞。

重传超时设置。可以通过设置NCCL中的参数来控制重传计时器和重试次数。我们调整了这些参数以在链路抖动（link flapping）下实现快速恢复。为了进一步减少恢复时间，我们在网卡上启用了adap_retrans特性。这个特性使得重传能够在更短的时间间隔内进行，并在链路抖动周期较短时帮助更快地恢复传输。

4. 容错

4.1 鲁棒的训练工作流

自动化故障识别与恢复流程。如图5所示，驱动进程（driver process）在收到提交的训练任务后，与一个定制的Kubernetes交互以分配计算资源，并为每个执行器（executor）启动相应的Pod。一个执行器管理一个节点。执行器完成一系列初始化任务后，会在每个GPU上创建训练进程，并启动一个鲁棒的训练守护进程，该进程定期向驱动进程发送心跳。这些心跳封装了各种信息，以实现实时异常检测和发出早期预警（§4.2）。当驱动进程检测到某个训练进程状态异常，或在预定时间窗口内未能收到来自执行器的心跳时，它会触发故障恢复程序。驱动进程将暂停所有执行器上的当前训练任务，并命令它们运行一系列自检诊断（§4.3）。一旦确定了有问题的节点，驱动进程会将这些节点的IP地址及其上运行的Pod信息提交给Kubernetes，后者会驱逐故障节点并用同等数量的健康节点补充集群。此外，我们提供一个用户界面，允许手动驱逐节点，特别是那些通过手动分析（如§5）确定的节点。恢复过程完成后，驱动进程从最新的检查点恢复训练。
图片描述
Figure 5: 鲁棒的训练工作流。

4.2 数据收集与分析

多维度心跳与监控。心跳消息包括执行器的基本信息，如IP地址、Pod名称和硬件信息等。此外，还会报告训练进程的当前状态，使驱动进程能够迅速检测到任何明显的异常。训练进程的stdout/stderr日志也包含在内，它们会被实时聚合、过滤和分析。如果检测到特定的警告或错误关键词，驱动进程将报告实时的诊断信息。此外，RDMA流量指标也被包括在内，作为网络利用率和效率的指标。某些训练过程中的异常可能不会表现为明确的错误，在这种情况下，RDMA流量指标成为一个关键指标。由于训练任务的周期性，每一步的网络流量特征应该表现出相似的模式。因此，RDMA流量的任何显著下降或异常波动都是潜在异常的信号。

分级监控系统。为了加强对训练稳定性和性能的监控，我们开发了一个精度达到毫秒级的监控系统。采用不同级别的监控来跟踪各种指标。秒级监控通常用于评估整体健康状况并排除常见的配置对训练的影响，例如ECN/PFC/QoS配置、链路抖动或任何其他网卡问题。毫秒级监控则用于确定网络是否拥塞，以及数据并行和流水线并行的数据传输速度是否达到了其物理极限。

4.3 诊断测试

轻量级诊断测试套件。在自检诊断中，执行时间与准确性之间存在权衡。我们通过迭代实验和优化，部署了一套轻量级诊断测试，有效覆盖了实际训练过程中遇到的大部分硬件和软件故障。
- 主机内网络测试: 我们使用内部开发的工具测试两方面。环回测试（Loopback test）测量从所有RDMA网卡（RNIC）到各种主机内端点（包括内存节点和GPU）的环回带宽，它在主机内进行全网状测试，覆盖所有可能的链路组合。RNIC到RNIC测试检查同一主机上不同RNIC之间的连接性和带宽性能。
- NCCL测试: 为了识别GPU通信中的潜在故障，我们在单个节点内的GPU之间运行一个all-to-all测试，以观察带宽是否符合预期基准。一旦主机内通信测试通过，每个节点还会与同一ToR交换机下的相邻机器进行all-reduce测试，以评估节点间的GPU通信。

4.4 快速检查点与恢复

两阶段快速检查点。为了在增加检查点频率的同时减少其引入的延迟，我们引入了一种优化的两阶段方法。第一阶段，每个GPU工作节点将其片上状态写入主机内存，然后继续训练过程。通过优化Pytorch的序列化机制和使用固定内存（pinned memory），得益于高PCIe带宽，这个过程可以缩短到几秒钟。第二阶段，一个后台进程接管，异步地将状态从主机内存传输到分布式文件系统（在我们的部署中是HDFS）进行集中维护。这种解耦使得GPU工作节点在转储状态后能几乎立即恢复训练。

优化的检查点恢复。在从检查点恢复的背景下，这是一个关键路径，因为没有上一个检查点就无法开始训练。瓶颈是HDFS的带宽。为了缓解这个瓶颈，我们提出了一种优化的数据检索策略。我们认识到多个GPU工作节点（例如，同一数据并行组中的工作节点）通常共享相同的状态分区。因此，我们指定组中的单个工作节点从HDFS读取共享的状态分区，从而线性地减少了负载。然后，这个工作节点将状态分区广播给共享相同数据的所有其他GPU工作节点。这种方法有效缓解了HDFS的带宽限制，从而大幅减少了恢复时间。

5. 训练故障排查

5.1 使用CUDA事件监视器进行性能诊断

发现性能不一致性。在数万GPU的规模下，我们观察到不同运行的计算效率不同，即使配置完全相同，这种不一致性依然存在（如图6所示）。我们还观察到训练任务的性能并非一直稳定，MFU会随时间逐渐下降。为了诊断这些性能问题，我们开发了一个性能分析工具，该工具记录每次运行中每个机器等级上关键代码段的执行时间。与之前的工具（如torch profiler或MegatronLM timer）不同，我们的工具基于CUDA事件方法计时，这最大限度地减少了对CUDA同步的需求，从而避免了性能下降，使我们能够在生产训练作业中持续运行它。
图片描述
Figure 6: 在大规模训练中观察到的不一致的MFU。不同的颜色表示同一训练任务的不同执行。

热力图可视化。第一种模式使用热力图来显示不同维度下机器之间的时间消耗差异，如图7所示。我们收集设备上计算阶段（前向和后向）的延迟数据，并在所有步骤中取平均值。聚合后的数据通过热力图可视化。热力图显示，一小部分机器（约0.5%）在训练期间表现出明显较慢的性能，从而阻碍了整体训练进度。这些掉队者（stragglers）决定了训练效率，导致不同运行中的效率不一致。排除这些异常机器后，各次运行的峰值MFU变得一致。
图片描述
Figure 7: 性能热力图。颜色表示一个等级上代码段的运行时间。该图还展示了3D可视化功能，其中选择了等级20，并且跨不同并行维度的依赖关系变得可见。

分布式时间线追踪。另一种模式以追踪格式显示机器上的事件时间线，并提供不同的分布式视图（数据并行、流水线并行、张量并行）。通过将不同等级的追踪跨度聚合到单一时间线上，我们获得了全面的视角，揭示了整体执行顺序、流水线气泡以及数据并行等级之间的同步特性。Figure 8展示了我们的分布式追踪器如何可视化流水线并行的实际执行情况，通过整合流水线并行组中的事件数据，详细说明了不同流水线阶段之间的数据依赖关系。
图片描述
Figure 8: 该追踪显示了在统一时间线上收集的流水线组中的事件。当选择一个事件时，依赖关系变得可见。

实时数据收集。来自CUDA事件计时器的每一条数据都存储在远程分析数据库中。计时器数据以逐行格式写入本地文件，一个独立的流处理器（streamer）进程将此日志文件实时同步到Kafka队列。分析数据库通过消费此Kafka队列中的数据保持更新，从而实现了在不中断训练作业的情况下进行即时分析。

5.2 3D并行训练可视化

定位复杂依赖下的故障。在3D并行和我们的优化技术（§3）下，数据流和任务序列的格局极其错综复杂。当单个GPU工作节点发生故障时，整个节点集群可能会在NCCL通信操作中停滞，最终导致系统范围的超时。从外部看，这种情况表现为一般的阻塞，但根本原因往往被大量的超时信息所掩盖。为了快速定位问题节点，我们让每个GPU工作节点在通信超时时记录其正在进行的事件。然后，这些日志被用来根据3D并行设置中的逻辑拓扑构建数据依赖关系的可视化表示。

可视化诊断方法。如图7所示，3D并行训练中的集群在逻辑上可以分为三个维度：张量并行、流水线并行和数据并行。当我们选择一个特定的GPU工作节点时，工具会显示其在逻辑拓扑中的位置、数据流向以及它所涉及的不同通信操作。重要的是，在发生错误时，该工具能直接访问工作节点的错误信息（如果有的话）。例如，在有缺陷的GPU概率性地导致NCCL通信操作阻塞的情况下，这种阻塞会挂起整台机器，导致其他依赖节点级联超时，最终使整个训练过程瘫痪。为了迅速识别这些故障节点，我们利用了3D并行训练可视化工具。因等待故障节点而超时的节点在退出时会记录其正在进行的操作。相比之下，带有故障GPU的节点被挂起，不会记录任何此类信息。因此，通过检查日志和可视化中的数据流，可以轻松地定位这些问题节点。

A4. 实验环境

  • 模型架构:

    • 实验使用了两种模型规模：175B和530B参数。
    • 所有情况下，序列长度为2048，词汇表大小为64,000。
    • 175B模型配置：96层，隐藏层维度12,288，96个注意力头。
    • 530B模型配置：105层，隐藏层维度20,480，128个注意力头。
    • 具体配置见Table 1。

    Table 1: 模型配置。
    图片描述

  • 硬件配置:

    • GPU: 最大的生产AI集群包含超过10,000个NVIDIA Ampere GPU。同时正在构建基于NVIDIA Hopper GPU的大型集群。
    • 网络: 数据中心网络基于Broadcom Tomahawk 4芯片的高性能交换机构建。
    • 服务器: 服务器配备8个200G网卡，以多轨方式连接到不同的ToR交换机。

  • 软件配置:

    • 代码实现: MegaScale构建于Megatron-LM【索引7】之上，实验对比的基线版本是Megatron-LM的GitHub提交285068c8【索引21】。
    • 依赖库: PyTorch。

A5. 实验结果

  • 训练性能:

    • 弱扩展性 (530B模型): 在批次大小与GPU数量成比例增加的设置下，MegaScale的MFU比Megatron-LM高出多达6.1%。随着规模扩大，Megatron-LM的MFU因掉队者和通信增多而下降1.6%，而MegaScale由于3D并行通信重叠，实现了近线性的扩展性 (Figure 9)。
      图片描述
      Figure 9: 530B模型的弱扩展性训练性能。

    • 强扩展性 (175B模型): 在固定批次大小（6144）下，MegaScale在所有GPU规模设置上均实现了高达1.34倍的加速。在12,288个GPU的最大规模下，MegaScale的MFU为55.2%，比Megatron-LM高出14个百分点。即使在最大规模下，由于计算通信比下降，MFU有所降低，但性能依然领先 (Table 2)。
      Table 2: 175B模型的强扩展性训练性能。括号中的数字表示MegaScale相比Megatron-LM的加速比。
      图片描述

    • 消融研究 (175B模型, 256 GPUs): 实验评估了各项优化的效果 (Table 3)。

      • 基线 (Megatron-LM) MFU为47.7%。
      • 算法优化（并行Transformer块、滑动窗口注意力）提升了5.6% MFU。
      • 3D并行通信重叠提升了6.2% MFU。
      • 高效算子提升了1.7% MFU。
      • 其他优化（数据流水线等）提升了1.1% MFU。
      • LAMB优化器（扩大批次大小）提升了3.0% MFU。
      • 最终，MegaScale的MFU达到65.3%，比基线高出17.6个百分点。
        Table 3: 训练175B模型（256 GPU，批次大小256）的MFU提升分解。
        图片描述

  • 模型收敛与稳定性:

    • 收敛性微基准测试: 在13B模型上的实验验证，并行Transformer块、滑动窗口注意力等算法优化，以及使用4倍批次大小的LAMB优化器，与基线相比均能达到可比较的损失，不会影响模型收​​敛 (Figure 10a, 10b)。
      (a) MegaScale（含算法优化）与Megatron-LM的训练损失曲线对比。
      图片描述
      Figure 10: 微基准实验中的训练损失曲线。
      (b) ADAM优化器与使用四倍批次大小的LAMB优化器的训练损失曲线。
      图片描述

    • 生产环境收敛与稳定性: 在一个超过10,000个GPU、持续数周的真实生产任务中（训练一个数千亿参数模型），模型损失持续收敛。期间，MegaScale在出现故障时成功修复和恢复了训练过程超过100次。超过90%的软硬件故障由鲁棒训练框架自动识别和修复，有效训练时间率保持在90%以上 (Figure 11)。
      图片描述
      Figure 11: 真实生产运行的归一化训练损失曲线。不同颜色表示训练重启。

  • 发现并修复的问题:

    • 计算掉队者: 使用CUDA事件计时器发现，特定主机在执行相同前向计算时耗时多约10%。隔离这些主机后，MFU提高了约0.7%。
    • MFU随时间下降: 通过分析发现，虽然计算时间稳定，但通信时间随训练步数增加而变长。根因是某些等级的梯度reduce-scatter启动较晚，这又追溯到前向计算阶段由不规律的垃圾回收和某些PyTorch操作引入的时间偏差。修改或移除这些有问题的代码段后，MFU下降问题得到解决 (Figure 12)。
      图片描述
      Figure 12: 解决了掉队者和有问题的代码段后，MFU变得稳定。
    • 频繁的网络接口抖动问题: 该问题由网卡、AOC线缆和交换机之间的链路质量不佳引起。解决方法是：1）将NCCL超时阈值设置得更大，以容忍短暂的接口down-up过程；2）对网卡信号强度、AOC线缆质量和交换机侧信号强度进行更低层次的质量控制。

方法细节中的参考文献引用汇总

  1. 【索引5，PaLM: Scaling language modeling with pathways, 2022】

    • 引用位置: 3.1节 算法优化 - 并行Transformer块
    • 引用内容: 引用其工作证明，将Transformer块的注意力和MLP并行化，不会降低拥有数千亿参数的模型的质量。

  2. 【索引8，Longformer: The long-document transformer, 2020】

    • 引用位置: 3.1节 算法优化 - 滑动窗口注意力 (SWA)
    • 引用内容: 引用其提出的滑动窗口注意力机制，并指出其工作证明了通过堆叠层可以创建大的感受野，从而在不牺牲精度的情况下提高效率。

  3. 【索引9，Large batch optimization for deep learning: Training bert in 76 minutes, International Conference on Learning Representations, 2020】

    • 引用位置: 3.1节 算法优化 - LAMB优化器
    • 引用内容: 引用其提出的LAMB优化器，该优化器被证明可以支持极大的批次大小而不损害模型准确性。

  4. 【索引14，GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model, https:// github.com/kingoflolz/mesh-transformer-jax, 2021】

    • 引用位置: 3.1节 算法优化 - 并行Transformer块
    • 引用内容: MegaScale采用的并行Transformer块源于此工作。

  5. 【索引15，Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023】

    • 引用位置: 3.2节 3D并行中的通信重叠 - 数据并行中的重叠
    • 引用内容: 受其启发，在每次迭代开始时预取初始的all-gather操作，使其与数据加载重叠。

  6. 【索引16，Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】

    • 引用位置: 3.3节 高效算子
    • 引用内容: 采用了此工作中提出的FlashAttention-2来优化注意力计算。

  7. 【索引17，Pytorch distributed: Experiences on accelerating data parallel training, 2020】

    • 引用位置: 3.5节 集体通信组初始化
    • 引用内容: 对torch.distributed进行了性能分析，以找出其在大规模场景下初始化缓慢的原因。

  8. 【索引18，Ieee 802.1 qbb - priority-based flow control, 2009】

    • 引用位置: 3.6节 网络性能调优 - 拥塞控制
    • 引用内容: 提及大规模all-to-all通信会导致高水平的优先级流量控制（PFC）。

  9. 【索引19，Congestion Control for Large-scale RDMA Deployments, ACM SIGCOMM Computer Communication Review, 2015】

    • 引用位置: 3.6节 网络性能调优 - 拥塞控制
    • 引用内容: 提及默认的DCQCN协议，以及过度使用PFC会导致的队头阻塞问题。MegaScale的拥塞控制算法结合了DCQCN的原理。

  10. 【索引20，Swift: Delay is Simple and Effective for Congestion Control in the Datacenter, SIGCOMM, 2020】

    • 引用位置: 3.6节 网络性能调优 - 拥塞控制
    • 引用内容: MegaScale开发的拥塞控制算法结合了Swift的原理。