近日，字节跳动联合北京大学的研究团队发表了一篇论文《MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs》，提出一个用于训练大语言模型的生产系统MegaScale，解决在万卡集群上训练大模型时面临的效率和稳定性挑战。
该论文介绍了MegaScale的设计、实现和部署。此外，文中还提到了万卡以上的集群规模遇到的问题及其解决方案。本文将简单翻译部分内容，文末附完整论文下载。
万卡集群的两大挑战
大模型时代，算力就是生产力。大模型的背后意味着巨大的计算资源，模型大小和训练数据大小是决定模型能力的关键因素。市场的主力玩家们利用数万个GPU构建大型人工智能集群，以训练LLM。但当GPU集群达到万卡规模，如何实现高效率、高稳定的训练？
第一个挑战是实现大规模的高效率训练。Model FLOPs utilization (MFU)是实际吞吐量与假设最大吞吐量之比，是评估模型训练效率的通用指标，可以直接反映端到端的训练速度。为了训练LLM，需要将模型分布为多个GPU上，并且GPU之间需进行大量通信以推动进展。除了通信之外，如操作符优化、数据预处理和GPU内存消耗等因素对MFU也有着显著影响。
第二个挑战是在大规模上实现训练的高稳定性，即在整个过程中保持高效率训练。在大模型训练中，稳定性十分重要。失败和延迟虽是大模型训练中的常态，但其故障成本非常高。如何缩短故障恢复时间至关重要，一个掉队者不仅会影响自己的工作，还会拖慢数万个GPU的整个作业。
为了应对这些挑战，字节跳动提出MegaScale系统，并已部署在自家的数据中心。那么，MegaScale是如何解决上述问题的呢？
如何实现大模型的高效训练？
想要在不损害模型准确性的情况下处理急剧增加的计算需求，需要采用最先进的算法优化、通信策略、数据流水线管理以及网络性能调优技术。下文深入探讨了用于优化大型模型训练的方法，以实现大规模的高效率训练。
算法优化
在算法层面进行了一些优化，在不影响准确性的前提下，提高训练效率。主要包括并行transformer块、滑动窗口注意力（SWA）和LAMB优化器。
并行transformer块：采用transformer块的并行版本，代替标准的序列化公式。这种方法使得注意力块和MLP块的计算可以并行执行，从而减少了计算时间。先前的研究表明，这种修改不会降低具有数千亿参数的模型的质量。
滑动窗口注意力（SWA）是一种稀疏注意力机制，它在输入序列中的每个标记周围使用固定大小的窗口，比全自注意力更高效。通过堆叠此类窗口注意力层，模型能够有效地捕获输入数据中广泛的上下文信息，同时创建大感受野，从而在不影响准确性的情况下加快训练速度。
LAMB优化器：大规模的高效训练通常受到批量大小限制的阻碍。特别是，增加批量大小可能会对模型收敛产生不利影响。LAMB优化器能够使BERT的训练批量大小扩展到64K，而不影响准确性。
3D并行中的通信重叠
3D并行指张量并行、流水线并行和数据并行。
在数据并行中有两个主要通信操作：all-gather操作和reduce-scatter操作。在3D并行中，单个设备可能承载多个模型块。重叠是基于模型块实现的，以最大化带宽利用。all-gather操作在模型块的前向传递之前触发，reduce-scatter操作在它的后向传递之后开始。这导致第一个all-gather操作和最后一个reduce-scatter操作无法隐藏。受到PyTorch FSDP的启发，初始的all-gather操作在每次迭代的开始时被预取，允许它与数据加载操作重叠，有效地将减少了通信时间。
在流水线并行中，MegaScale使用交错1F1B调度方法，以实现通信的重叠。在热身阶段，前向传递仅依赖于其先前的接收。我们解耦了通常一起实现的发送和接收，通过打破这种依赖关系，使得发送操作能够与计算重叠。在张量/序列并行中，介绍了融合通信和计算等优化策略，以及将GEMM内核分成小块并与通信进行流水线执行。
高效操作符
尽管在MegatronLM中已经对GEMM操作符进行了优化，但其他操作符中还有进一步增强的机会。注意力部分采用了FlashAttention-2，改进了不同线程块和warp之间的工作分配。LayerNorm和GeLU由先前实现中的细粒度内核组成。通过将这些内核融合在一起，减少了与启动多个内核相关的开销，并有助于优化内存访问模式，从而实现更好的性能。
数据流水线优化
数据预处理和加载经常被忽视。然而，这些操作在每个训练步骤开始时会产生不可忽视的GPU空闲时间。优化这些操作对于训练过程的效率至关重要。
异步数据预处理。数据预处理不在关键路径上。因此，当GPU工作器在每个训练步骤结束同步梯度时，可以开始后续步骤的数据预处理，这就隐藏了预处理的开销。
消除冗余数据加载器。在分布式训练的典型数据加载阶段，每个GPU工作器都配备了自己的数据加载器，负责将训练数据读入CPU内存，然后转发到GPU。这导致工作线程之间为争夺磁盘读取带宽，因此产生了瓶颈。我们观察到，在LLM训练设置中，同一台机器内的GPU工作器处于相同的张量并行组。因此，它们每次迭代的输入本质上是相同的。基于这一观察，我们采用了两层树状的方法，在每台机器上使用一个专用的数据加载器将训练数据读入共享内存。随后，每个GPU工作器负责将必要的数据复制到自己的GPU内存中。这就消除了冗余读取，并显著提高了数据传输的效率。
集体通信群初始化
在分布式训练中，初始化阶段涉及在GPU工作器之间建立NVIDIA集体通信库（NCCL）通信组。由于这种开销在小规模场景中相对较小，因此默认使用torch.distributed。随着GPU数量扩展到超过一万个，naive实现引入的开销变得无法忍受。
torch.distributed初始化时间过长有两个原因。第一个问题在于同步步骤，其中每个进程在初始化特定通信组结束时参与了一个屏障操作。这个屏障使用TCPStore，以单线程、阻塞的读写方式操作。可以用非阻塞和异步的Redis替换TCPStore。第二个问题与全局屏障的不慎使用有关。每个进程在初始化其相应的通信组后执行一个全局屏障。我们精心设计了通信组的初始化顺序，以最小化全局屏障的需求，降低了时间复杂度。
在未经优化的情况下，2048张GPU的集群初始化时间是1047秒，优化后可降至5秒以下；万卡GPU集群的初始化时间则可降至30秒以下。
网络性能调优
分析了3D并行中机器间的流量，并设计了技术方案来提高网络性能。包括网络拓扑设计、减少ECMP哈希冲突、拥塞控制和重传超时设置。
网络拓扑。我们的数据中心网络是基于Broadcom Tomahawk 4芯片构建的高性能交换机。每个Tomahawk芯片的总带宽为25.6Tbps，具有64×400Gbps端口。三层交换机以CLOS类似的拓扑连接，以连接超过10000个GPU。每层交换机的下行链路和上行链路的带宽比为1:1。也就是说，32个端口用于下行，32个端口用于上行。该网络以较小的直径提供了高带宽，每个节点都可以在有限的跳数内与其他节点通信。
减少ECMP哈希冲突。我们精心设计了网络拓扑，并调度网络流量以减少ECMP哈希冲突。首先，在机架ToR交换机上把上行与下行链路分开，一个 400G 下行链路端口通过特定的 AOC 电缆分为两个 200G 下行链路端口，有效降低冲突率。
拥塞控制。在分布式训练中大规模使用默认的DCQCN协议时，all-to-all通信可能会导致拥塞和PFC级别的提高。过度使用PFC可能会导致头部阻塞（HoL），从而降低网络吞吐量。为了缓解这些问题，我们开发了一个结合了Swift和DCQCN原理的算法，该算法将往返时间（RTT）的精确测量与显式拥塞通知（ECN）的快速拥塞响应能力相结合。这种方法显著提高了吞吐量，并最小化了与PFC相关的拥塞。
重传超时设置。NCCL中的参数可以设置以控制重传定时器和重试次数。我们调整这些参数以在链路抖动时快速恢复。为了进一步减少恢复时间，我们在NIC上启用了adap_retrans功能。此功能支持在较短的时间间隔内进行重传，当链路抖动周期较短时，有助于更快地恢复传输过程。
容错性
随着训练集群扩展到超过数万个GPU，软件和硬件故障几乎是不可避免的。我们为LLM训练设计了一个健壮的训练框架，实现了自动故障识别和快速恢复，在最小的人为干预和对正在进行的训练任务最小影响的情况下实现容错性。
如上图所示，在接收到训练任务后，驱动程序进程会与自定义的Kubernetes接口进行交互，以便分配计算资源并为每个执行器启动相应的Pod。一个执行器管理一个节点。执行器完成初始化任务后将在每个GPU上创建训练进程，并启动一个健壮的训练守护进程，定期向驱动程序发送heartbeat以便实时检测异常并预警。当检测到异常状态或在预定时间内未收到状态报告时，会触发故障恢复程序，将暂停所有正在进行的训练任务，并命令它们自我检查诊断。
一旦识别出问题节点，驱动程序将向Kubernetes提交要被封锁的节点的IP地址，以及在这些节点上运行的Pod信息，Kubernetes将驱逐故障节点，并用健康节点替换。此外，还有一个用户界面可以手动删除问题节点。恢复过程完成后，驱动程序会从最新的checkpoint恢复训练。我们优化了checkpoint和恢复过程，以最小化训练进度的损失。
为了增强对训练稳定性和性能的监控，开发了一个精度达到毫秒级的监控系统。采用不同级别的监控来跟踪各种指标。此外，文中还讲述了如何实现checkpoint快速恢复、训练故障排除，以及MegaScale部署和运营的经验，感兴趣的可下载论文查阅。
结论
本文深入探讨了MegaScale的设计、实现和部署。通过算法-系统协同设计，MegaScale优化了训练效率。在12288个GPU上训练一个175B LLM模型时，MegaScale实现了55.2%的MFU，比Megatron-LM提高了1.34倍。
我们强调在整个训练过程中需要容错，并实现了一个定制的健壮训练框架，以自动定位和修复故障。此外，还提供了一套全面的监控工具，用于深入观察系统组件和事件，便于复杂异常的根本原因识别。我们相信，我们的工作不仅为那些从事LLM训练的人提供了实用的见解，也为这个快速发展的领域的未来研究铺平了道路。
关注SDNLAB公众号，后台回复0306MegaScale，即可获取论文下载链接
