MegaScale：万级GPU集群中大模型训练

论文链接：https://arxiv.org/abs/2402.15627

MegaScale系统简介

MegaScale是一个专为在超过10,000个GPU上训练大型语言模型（LLMs）而设计的生产系统。该系统通过算法和系统组件的协同设计，解决了大规模训练中的效率和稳定性挑战，从而实现了高效的训练。

问题现状

随着模型参数规模和训练数据量的增加，大型语言模型训练所需的计算资源也日益增长。这导致了训练效率和稳定性的挑战成为关键问题。

• 第一个挑战是在规模上实现高训练效率。模型FLOPs利用率（MFU）是观察到的吞吐量与假设峰值FLOPs为100%时的理论最大吞吐量的比率[7]。这是一个标准的评估训练效率的指标，直接转化为端到端训练速度。LLM训练并不是令人尴尬的并行。为了训练一个LLM，模型被分割在GPU上，GPU之间进行大量通信以取得进展。除了通信，操作符优化、数据预处理和GPU内存消耗等因素也显著影响MFU。
• 第二个挑战是在规模上实现高训练稳定性，即在整个训练过程中保持高训练效率。从生产的角度来看，稳定性尤为重要，因为LLMs的训练时间很长。用一万亿个标记训练一个LLM可能需要数周时间。这个规模和时间比常规DNN训练作业的规模和时间大得多。对于LLM训练来说，故障和落后者是常态而不是例外。在这样的规模下，故障和落后者的后果是毁灭性的。故障非常昂贵，考虑到规模之大，减少恢复时间至关重要。一个落后者不仅影响自己的工作，还会拖慢涉及数万个GPU的整个作业

相关研究

目前已经有多种大型语言模型被开发和训练，如GPT-3、PaLM等。然而，大多数研究侧重于模型本身的性能提升，而对支撑这些模型训练的系统基础设施的具体细节探讨相对较少。

解决方案主要内容

• 算法优化：

• 并行Transformer模块：将Transformer架构的不同部分分布在多个GPU上进行并行处理。

• 滑动窗口注意力（Sliding Window Attention, SWA）：减少长序列处理时的计算复杂度。

• LAMB优化器：大规模训练受到批量大小限制。特别是，增加批量可能会影响模型收敛。LAMB可以将批量大小扩展到4倍，而不会损失准确性。针对大规模模型的优化算法，有效处理稀疏梯度更新。
  

• 通信重叠：通过在数据并行、流水线并行以及张量/序列并行中采用技术来隐藏通信开销。

• 高效操作符：

• FlashAttention-2：优化注意力机制的实现，提高计算效率。

• LayerNorm：优化层归一化的实现，提高计算速度。

• GeLU：优化激活函数的实现，加快运算速度。

• 数据流水线优化：采用异步数据预处理和数据加载器优化策略，确保数据流的连续性和高效性。

• 集体通信组初始化：优化NCCL通信组的初始化过程，减少初始化时间。

• 网络性能调优：包括设计自定义网络拓扑结构、减少ECMP哈希冲突、定制拥塞控制策略等措施。

• 健壮性训练框架：提供自动化故障定位和恢复功能，确保训练过程的高稳定性。
  

心跳检测与预警系统

心跳检测机制是实时监控集群健康状况的关键组件，通过持续监测各个节点的活动状态，能够及时发现并预警潜在的异常情况。系统采用毫秒级的监测频率，对IP地址、Pod名称、硬件状态、训练进程及RDMA流量进行全方位监控。一旦监测到任何异常，如硬件故障、训练停滞或网络流量突变，系统将立即触发警报，为管理员提供早期预警，以便迅速采取措施。

自诊断系统

自诊断系统具备轻量级软硬件故障诊断能力，能够在不影响正常运行的前提下，自动定位并识别系统中存在的问题。这一机制极大地提升了问题排查的效率，减少了故障响应时间，确保了集群的稳定运行。

主机内网测试

主机内网测试旨在评估RDMA技术在本地环境下的性能表现，主要包括：

1. RDMA网卡到主机内部端点（如内存节点和GPU）的带宽测试：确保数据传输的高速率和低延迟，为高效计算奠定基础。

2. 同一主机上不同RDMA网卡之间的连接与带宽测试：验证多网卡环境下数据通信的可靠性和效率。

3. NCCL测试：利用NVIDIA Collective Communications Library进行以下测试：

   • 服务器内GPU间的alltoall测试，评估数据交换效率。
   • 同一交换机下服务器间的allreduce测试，检验大规模分布式计算的同步性能。

故障恢复机制：双阶段Checkpoint策略

为了提高故障恢复的速度与效率，我们采用了创新的双阶段Checkpoint保存与调用机制：

快速Checkpoint保存

• 第一阶段：在模型训练过程中，将Checkpoint临时写入主机内存，实现快速保存，减少磁盘I/O带来的延迟。

• 第二阶段：在模型继续训练的同时，异步地将Checkpoint从内存中写入高容量的HDFS存储系统，确保数据持久化。

快速Checkpoint调用

• 第一阶段：当需要恢复训练时，单个GPU首先从HDFS读取Checkpoint数据，减少整体等待时间。

• 第二阶段：读取Checkpoint的GPU将数据广播给所有其他需要相同数据的GPU，实现快速同步，加快恢复速度。

这套双阶段Checkpoint机制不仅显著提高了故障恢复的效率，还保障了训练进程的连续性与数据的安全性，是构建高可用、高性能分布式训练平台的重要组成部分。

8. 深入的可观测性：开发诊断工具来监控系统组件和事件，以便及时发现问题并采取措施。
   
   

实验数据及效果

• MegaScale在使用12,288个GPU训练一个拥有1750亿参数的LLM时，实现了55.2%的模型FLOPs利用率（MFU），相比Megatron-LM提高了1.34倍。
  
  

• 在实际生产环境中，MegaScale成功训练了一个拥有数千亿参数的专有模型，在多周的训练过程中经历了多次故障，但系统能够自动检测并恢复训练过程，保证了训练的连续性和效率。