跳到主要内容 大语言模型 (LLM) 分布式高效训练技术综述:背景、并行、计算、内存、通信、容错、展望 | 极客日志
Python AI 算法
大语言模型 (LLM) 分布式高效训练技术综述:背景、并行、计算、内存、通信、容错、展望 大规模语言模型(LLM)分布式训练面临可扩展性、效率和可靠性三大挑战。综述了 LLM 训练系统的最新进展,涵盖 AI 加速器、网络、存储及调度基础设施。重点分析了混合并行、自动并行及异构并行策略,探讨了计算优化(如算子优化、混合精度训练)、内存优化(如激活重计算、ZeRO、卸载)及通信优化(如集体通信算法、调度、网络内聚合)。此外,文章还讨论了容错机制,包括故障分析、异常检测、检查点恢复及无检查点恢复技术。最后展望了光电计算等未来方向。
FrontendX 发布于 2025/2/7 更新于 2026/4/20 1 浏览大规模语言模型(LLMs),如 GPT 和 LLaMA,以其先进的能力正在彻底改变人工智能行业。训练这些模型需要大量的 GPU 集群和显著的计算时间,在可扩展性、效率和可靠性方面带来了主要挑战。本调查探讨了 LLMs 训练系统的最近进展,包括在 AI 加速器、网络、存储和调度方面的训练基础设施的创新。此外,调查涵盖了并行策略,以及在分布式 LLM 训练中的计算、通信和内存优化。它还包括在长时间训练期间维护系统可靠性的方法。通过检查当前的创新和未来的方向,本调查旨在为改进 LLM 训练系统和应对持续挑战提供有价值的见解。此外,传统的基于数字电路的计算系统在满足 LLMs 的计算需求方面面临重大限制,凸显了需要创新解决方案,如光子计算和光子网络。
1 引言
大规模语言模型(LLMs)正在改变人工智能行业,展示了在广泛的任务和应用中的卓越能力,包括个人助理 [1]、代码副驾驶 [2]、芯片设计 [3] 和科学发现 [4]。这一革命的成功建立在基于 transformer 的 LLMs 前所未有的规模上,如 GPT [5]、LLaMA [6]、Gemini [7] 等。此外,有证据表明 LLMs 的扩展尚未达到平台期 [8]。这一趋势显著改变了底层训练系统和基础设施的设计,因为 LLM 通常遵循相对固定的架构,其训练专门占据了大量的 GPU 集群超过延长的时间段。例如,在 Meta 的生产集群上,LLaMA-3 的预训练大约需要 54 天,使用 16K H100-80GB GPU [9]。LLM 训练突显了当今训练系统和基础设施在'SER',即可扩展性、效率和可靠性方面的重大挑战。可扩展性要求基础设施和系统无缝适应数万个 GPU 或 AI 加速器的大规模集群,同时保持训练的正确性和模型的准确性。这需要在硬件配置、网络和训练框架方面的创新解决方案。效率侧重于在整个集群中最大化资源利用率,通常通过模型 FLOPs 利用率(MFU)来衡量。实现高 MFU 涉及优化计算、最小化通信开销,并以前所未有的规模有效管理内存。可靠性对于 LLM 训练的持续时间至关重要,通常持续数周到数月。系统必须保持一致的性能,并对各种类型的故障具有弹性,包括硬件故障、网络问题和软件错误。它应该能够快速检测并从这些故障中恢复,而不会显著损失进度或训练质量。这些相互关联的挑战需要系统和基础设施设计的全面方法,推动大规模分布式计算和高性能机器学习系统的界限,并为研究和创新开辟新途径。这篇综述论文旨在提供 LLM 训练系统和基础设施进展的全面概述,解决上述挑战。
2 背景
2.1 基于 Transformer 的 LLMs
当前最先进的 LLMs 主要是基于 Transformer 的。它们的核心架构围绕注意力机制构建 [10],该机制允许模型动态地权衡句子中不同单词的重要性。图 2 描绘了一个典型的 Transformer 层的架构 [10],该层可以多次堆叠以构建一个 LLM。输入文本首先被标记化为单个标记,然后通过嵌入层转换为标记向量 X。为了保持文本的顺序特性,将位置信息嵌入到标记向量中。然后,得到的标记向量被送入 Transformer 层,该层由一个注意力块和一个前馈神经网络(FFN)块组成。假设输入的标记向量是 X = [x1, x2, …, xn]。这些标记首先通过线性变换转换为查询 Q、键 K 和值 V 张量。注意力机制如下计算注意力输出:
其中 d 是键张量的维度。这个公式通过计算加权和来确保 LLM 可以专注于输入序列的相关部分,其中权重来源于查询和键之间的相似性。在注意力层之后,输出被传递到 FFN 进行进一步处理。如今,LLMs 通常遵循原始的仅解码器 Transformer 架构,但对注意力机制和 FFN 进行了修改以提高效率和性能。原始的注意力机制,称为多头注意力(MHA)[10],由于对键 - 值缓存的二次计算复杂度和高内存消耗而受到限制。为了解决这些问题,提出了几种变体,如多查询注意力(MQA)[11]、组查询注意力(GQA)[12] 和多潜在注意力(MLA)[13]。FFN 组件的一个显著进步是混合专家(MoE)[14], [15] 架构,它采用稀疏激活的 FFN。在 MoE 中,每个输入只有一部分 FFN 层(或专家)被激活,显著减少了计算负载,同时保持了高模型容量。
相关免费在线工具 加密/解密文本 使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
RSA密钥对生成器 生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
Mermaid 预览与可视化编辑 基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
curl 转代码 解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online
Base64 字符串编码/解码 将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具,Base64 字符串编码/解码在线工具,online
Base64 文件转换器 将字符串、文件或图像转换为其 Base64 表示形式。 在线工具,Base64 文件转换器在线工具,online
LLM 训练工作负载的特性与传统的深度学习工作负载显著不同,主要是因为它们的复杂性和规模。这些独特的特性影响训练系统设计、性能、可扩展性和资源利用率。在这里,我们强调 LLMs 的关键差异和要求。
(1) 同质模型架构。 与以往采用不同模型架构(例如,LSTM [16],CNN [17])执行不同任务的 DL 工作负载不同,LLMs 主要使用 Transformer 架构 [10]。像 GPT [5]、LLaMA [6]、InternLM [18] 和 MOSS [19] 这样的模型都共享这一共同基础。这种架构的统一性为针对特定模型架构优化系统性能提供了重要潜力。
(2) 前所未有的规模和训练时长。 LLM 训练以前所未有的规模运行,通常使用具有数千亿参数的模型和 TB 级训练数据集进行更新。这种规模需要在大型 GPU 集群上进行分布式训练,并在保持高效率方面提出挑战。此外,LLMs 的训练可以持续数周或数月,要求强大的容错机制和高效的检查点策略,以防止数据丢失并促进中断训练会话的恢复。
(3) 专门的软件优化。 为了适应 LLMs 的巨大模型大小,专门的系统实现了先进的技术来优化执行。例如,Megatron [20] 和 Alpa [21] 通过混合并行性加速训练。DeepSpeed [22] 通过集成状态分片优化器减少内存消耗。
(4) 训练范式的转变。 传统的 DL 工作负载遵循特定于任务的范式,在特定领域的数据上训练模型以执行特定任务,例如翻译。相比之下,LLMs 采用自监督训练方法,在大量数据集上进行训练以创建基础模型,然后适应各种下游任务。这一范式转变代表了模型开发流程的实质性变化,包括预训练和对齐阶段,并且与以前的 DL 工作负载相比,产生了不同的工作负载特性。从数据中心的角度来看,LLM 开发涉及许多与预训练相关的小规模工作负载,包括对齐(即微调)和定期评估工作负载 [23]。
LLM 训练工作负载的独特特性带来了在开发高效训练系统和基础设施方面的重大挑战。这些挑战主要表现在三个关键领域:可扩展性、效率和可靠性。每个挑战都直接源于 LLMs 的大规模和其训练过程的复杂性,需要创新解决方案,推动分布式计算和机器学习系统的界限。下面,我们详细讨论这些挑战及其对 LLM 训练的影响:
(1) 可扩展性。 LLMs 的成功在很大程度上归因于它们的规模,随着 LLMs 变得更大,性能通常会提高 [8]。然而,模型规模的扩展引入了可观的可扩展性挑战,因为训练 LLMs 需要越来越大的 GPU 或专门的 AI 加速器集群。首先,构建可扩展的基础设施,提供巨大的计算和内存容量是必要的。这涉及设计和部署大量的 GPU 或专门的 AI 加速器、连接这些设备的高性能网络,以及能够处理庞大数据集和模型检查点的分布式存储系统。挑战在于确保这些组件在大规模上高效协同工作,管理大规模部署中的热散失、功耗和硬件故障。其次,设计可扩展的训练系统,可以有效地并行使用大量加速器至关重要。这包括设计并行化策略和通信算法,以实现在数千个加速器上近乎线性的可扩展性,同时保持一致的 LLM 准确性。
(2) 效率。 LLM 训练的巨大计算需求转化为高昂的训练成本,这使得最大限度地提高硬件和软件系统的效率变得至关重要。效率可以通过 MFU(模型 FLOPs 利用率)来衡量,它量化了系统使用可用计算资源的有效性。然而,在大规模上实现高效率仍然是一个重大挑战。例如,LLaMA3 在 16K GPU 上仅实现了 38% 到 41% 的 MFU [9],突显了随着系统规模扩大,保持高利用率的难度。最大化效率要求在并行性、计算、通信和内存方面进行优化。首先,分布式 LLM 训练的并行性需要精心设计,以最小化通信需求。其次,优化的计算操作符和较低精度的算术是实现高 GPU FLOPS 利用率的关键。第三,需要最小化通信开销,以减少 GPU 空闲时间。最后,需要有效的内存优化,以在现有硬件中容纳 LLMs 并减少重计算的 FLOPs 浪费。
(3) 可靠性。 确保 LLM 训练在长时间内的可靠性至关重要。由于训练作业可以在大型 GPU 集群上持续数周到数月,训练失败的可能性增加,需要快速的故障检测和恢复机制,以实现弹性的 LLM 训练。首先,LLM 训练作业可能因各种错误而崩溃,使得在数万个 GPU 上快速确定确切的故障原因变得困难。其次,LLM 训练作业的挂起导致所有 GPU 由于训练的同步性质而变得空闲,导致显著的浪费。此外,一些复杂异常,如冗余链路故障或落后者,可能不会导致立即崩溃,但可能导致训练减速。这种不稳定性可能导致训练效率降低。为了应对这些挑战,需要强大的异常检测系统,能够检测出灾难性故障和性能下降。此外,实施能够无缝处理节点故障和网络问题的容错训练框架至关重要。
这项工作侧重于基于 Transformer 的 LLMs 的高效训练系统和基础设施,包括分布式基础设施的设计、并行范式、计算和通信优化、内存管理的有效性和训练系统的弹性。我们还研究了新兴工作负载(如 MoE,一种有前景的高效 LLM 变体)和微调的高效训练系统,微调是调整 LLMs 能力的必要阶段。然而,这项工作没有涵盖有前途的 LLM 架构的演变 [24], [25] 以及训练 [26]、指令调整 [27] 和对齐 [28] 的算法,以实现强大和安全的 LLMs。虽然以前的工作 [29]–[31] 讨论了 LLM 训练系统的一些方面,但它们的主要焦点不是高效训练系统和基础设施的设计。Wan 等人 [29] 旨在提供模型和数据为中心方法中高效 LLM 进展的全面视图。Liu 等人 [30] 涵盖了 LLMs 的训练和推理部署技术。Xu 等人 [31] 针对 LLM 开发讨论了资源高效策略,包括算法和系统方面。这项工作还讨论了量化 LLM 训练和高效 LLM 微调的方法,但我们专注于系统方法。Zhu 等人 [32] 和 Han 等人 [33] 讨论了压缩和微调 LLMs 的算法方法。这项工作讨论范围不包括高级优化算法 [34] 和分布式 DNN 训练系统 [35]。虽然 Liang 等人 [36] 广泛回顾了自动并行化方法,它们的重点是一般的 DNN,而不是特定的 LLMs。
在本节中,我们探讨了用于训练 LLMs 的基础设施设计,包括加速器、网络和调度系统(见图 3)。
LLMs 的快速发展在很大程度上得益于 GPU 和 AI 加速器的进步,这些加速器对提高模型训练性能至关重要。
NVIDIA GPU 已成为分布式 LLM 训练的重要组成部分,因为它们具有处理并行计算的卓越能力。这些处理器构建有众多紧凑、高效的内核,能够同时执行众多任务。GPU 的设计非常适合 LLM 训练中的矩阵和向量运算。它们支持各种数值精度格式,如 FP32、TF32、FP16、BF16、FP8、INT8,甚至是 FP4。这使得研究人员能够很好地平衡训练速度和准确性,使 LLM 训练更加高效 [110]。NVIDIA 的 GPU 编程语言(即 CUDA)使研究人员更容易管理在 GPU 上并行拆分和处理任务的方式。这有助于研究人员充分利用 GPU 来训练先进的 LLMs。一个典型的 GPU 由一系列流式多处理器(SM)组成,每个 SM 包含多个核心,它们共享一个指令单元,但能够并行执行不同的线程。每个 SM 中的共享内存允许线程之间进行有效的数据交换和同步,这对于优化 LLM 计算所需的内存访问模式至关重要。此外,GPU 配备了高带宽内存(HBM),这加速了数据传输并缓解了计算密集型任务中的内存访问瓶颈。最新的 GPU 架构,如 NVIDIA 的 Ampere [37]、Hopper [38] 和 Blackwell [39],正在不断推动 LLM 计算的边界。它们提供了增强的内存带宽和容量、每秒浮点运算次数(FLOPS)的增加,以及专门的混合精度计算单元,如 Tensor Cores。值得注意的是,NVIDIA 的 Hopper 架构通过引入 Transformer Engine [111],这是一个利用混合 FP8 和 FP16 精度来加速基于 Transformer 的 LLMs 训练的功能。
在 AMD GPU 上进行分布式 LLM 训练已成为现实,特别是在 Frontier [112] 上,这是世界上第一台 E 级超级计算机。每个 Frontier 节点都配备了 8 个 MI250X [40] AMD GPU,每个 GPU 有 64GB 的 HBM,理论 FP16 峰值性能为 191.5 TFLOPS。这种配置为高效训练万亿参数模型提供了无与伦比的机会。解锁这一潜力的关键在于适应现有的 CUDA 基础工具和框架到 ROCm 平台 [113], [114]。值得注意的是,已经开发了 ROCm 支持的 FlashAttention [115] 和 FlashAttention2 [116] 版本,允许高效执行注意力。各种具有强大计算能力和软件优化的 AI 加速器已经被开发出来训练 LLMs。GAUDI [41] 提供了一个异构计算架构,包括两个矩阵乘法引擎和一群完全可编程的张量处理核心,能够高效处理 LLM 训练操作。这个处理器可以使用 384 个 GAUDI2 卡训练具有 1750 亿参数的 GPT-3 模型 [117]。Google TPUv4 [42] 超级计算机有 4096 个芯片,支持 LLM 训练,平均达到峰值 FLOPS 的约 60%。Graphcore Bow Pod64 [43],一个包含 64 个 Bowclass IPUs 的机架设置,实现了 22 petaFLOPS。它支持使用 256 个 IPUs 训练 GPT-3 模型。Cerebras CS-2 [44] 是一个晶圆级深度学习加速器,包含 850,000 个处理核心,每个核心提供 48KB 的专用 SRAM 内存。它被用来训练 Cerebras-GPT,一个开放计算优化的语言模型家族 [118]。
通信开销是扩展 LLM 训练的主要障碍 [119], [120]。例如,在训练期间减少模型梯度可能导致超过 90% 的训练时间花在通信上 [121]。为了解决这个问题,研究社区专注于改善 LLM 训练的通信基础设施。
芯片到芯片通信对于节点内 AI 加速器之间的数据传输至关重要,这显著影响了 LLM 训练的效率。传统上,这种通信依赖于 PCI Express(PCIe)[122],它采用树形拓扑结构——一个分层结构,多个设备连接到单个根复杂体。多年来,PCIe 已经提高了其带宽:PCIe 3.0 每条线路提供大约 1 GB/s 的带宽,对于 16 条线路的配置总计约 16 GB/s;PCIe 4.0 将带宽翻倍至每条线路 2 GB/s,而 PCIe 5.0 进一步将其增加到每条线路 4 GB/s。尽管这些增强,PCIe 在带宽、延迟和可扩展性方面的固有限制使其对 LLM 训练不太理想 [123]。为了解决这些限制,像 NVLink [45] 这样的专用芯片到芯片互连越来越多地被用于 LLM 训练。与传统 PCIe 相比,这些高级互连通过使用各种拓扑结构:立方体网格、全连接和 3D-环面,提供了显著更高的带宽和更低的延迟。此外,共享内存模型、专门的通信协议和同步机制也起着至关重要的作用。
立方体网格拓扑。 NVLink-1.0 [45] 为 GPU 提供了直接和高速的连接,每个链路提供 160 GB/s 的双向带宽。这种架构能够在四个 GPU 之间形成平面网格结构,在八个 GPU 之间形成立方体网格拓扑,可以配置成 DGX-1 服务器。这种立方体网格配置,尽管不是全到全连接,但显著提高了 GPU 上的数据通信效率和训练性能。
全连接拓扑。 许多互连使用基于交换机的或基于 P2P 的全连接拓扑来提高芯片到芯片通信性能。NVIDIA 使用 NVSwitch [46] 实现 GPU 之间的基于交换机的全到全互连。在 DGX-2 [124] 系统中,六个 NVSwitch 完全连接了十六个 GPU,提供了任何两个 GPU 之间 300 GB/s 的双向带宽。这个带宽在 NVSwitch 2.0 中增加到 600 GB/s,在 NVSwitch 3.0 中进一步增加到 900 GB/s。Intel、AMD 和华为 Ascend 使用基于 P2P 的全连接拓扑为他们的加速器,其中每个芯片使用以太网或无限 Fabric [47] 直接连接到同一节点中的每个其他芯片。与传统的基于交换机的拓扑相比,基于 P2P 拓扑中两个 GPU 之间的带宽受到直接连接链路带宽的限制。
2D/3D-环面拓扑。 Google 的 TPU 系统使用环面网络拓扑 [59] 进行芯片到芯片通信。它通过将每个 TPU 芯片连接到其四个相邻邻居在一个网格中,边缘包裹,形成一个环形结构。这种架构设计确保了由于芯片之间存在多个直接路径,因此具有低延迟和高带宽。具体来说,TPUv2 [48] 超级计算机采用了 16x16 的 2D 环面配置,包含 256 个芯片,通过高速芯片间互连(ICI)链路相互连接。TPUv3 [49] 超级计算机使用 32x32 的 2D 环面,包含 1024 个芯片。从 2D 环面设计发展而来,TPUv4 [42] 超级计算机将计算资源组织成多机器立方体,具有 3D 环面拓扑。每个 TPU 机器包含四个芯片,以 2x2x1 网格排列,通过 ICI 链路相互连接。十六个这样的 TPU 机器组合成一个数据中心机架,其中机架内的 ICI 链路相互连接,形成一个 4x4x4 网格,从而形成一个 3D 环面结构。这种高级配置显著提高了通信效率和可扩展性,特别有利于 LLM 训练。
远程直接内存访问(RDMA)[54] 实现了节点间高速、低延迟的数据传输。RDMA 允许直接从一个计算机的内存访问另一个计算机的内存,而不涉及任一节点的操作系统。GPUDirect-RDMA [50] 通过启用跨不同节点的 GPU 之间的直接通信来增强这一过程,完全绕过 CPU。这项技术对 LLM 训练特别有益,因为它加速了模型参数和梯度的同步。两种最流行的 RDMA 技术是 InfiniBand [51] 和基于汇聚以太网的 RDMA(RoCE)[52]。InfiniBand 是一种高速、低延迟的网络技术,广泛用于 HPC(高性能计算)环境,如 Eagle 超级计算机 [125]。这项技术需要专用的网络基础设施,反映了它在提供卓越性能方面的设计重点。多年来,InfiniBand 在带宽能力方面有了显著的演进,从 EDR(增强数据速率)的 100 Gbps 发展到 HDR(高动态范围)的 200 Gbps,最近发展到每个链路的 NDR(下一个数据速率)400 Gbps [126]。RoCE 利用现有的以太网基础设施提供 RDMA 能力。这种方法提供了一种更具成本效益和更易于部署的解决方案,特别是在已经使用以太网的数据中心。RoCE 有两个版本:RoCE-v1 [52],作为以太网链路层协议运行,和 RoCE-v2 [53],通过 UDP 运行。行业领导者如字节跳动和 Meta 已经采用了这些技术来扩展 LLM 训练。另一种 RDMA 协议,互联网广域网 RDMA 协议(iWARP)[54],通过 TCP/IP 网络启用 RDMA。然而,由于其相对较低的性能,iWARP 通常不用于分布式 LLM 训练 [127]。
在 LLM 训练集群中,网络架构被结构化为前端和后端组件(见图 3)。前端网络处理各种流量,如作业管理、模型推理和存储活动,而后端网络专门用于训练过程中产生的大量流量。我们在优化 LLM 训练的主要关注点在于提高后端网络的性能和效率,以便将 AI 加速器扩展到数万个。
HPC 网络拓扑。 传统的 HPC 环境拓扑也可以在 AI 集群中用于分布式训练,如 Clos [55]、BCube [56]、DCell [57]、Jellyfish [58]、Torus [59]、Dragonfly [60] 和 Dragonfly+ [61]。Clos 网络架构,通常被称为 Fat-Tree 拓扑,在 LLM 训练集群中广泛使用。在基于 Clos 的集群中,每台服务器都配备了一个或多个 NIC,并被组织成连接到叶交换机的机架。这些叶交换机连接到脊柱交换机,提供机架间的连通性,并形成一个豆荚。豆荚进一步通过核心交换机相互连接,促进集群内服务器之间的任何到任何通信。例如,Meta 的最后一代 GPU 集群架构,支持多达 24,000 个 GPU,由八个豆荚组成,它们之间具有全脂带宽,并在核心层使用 7:1 的过度订阅比率 [9]。Meta 使用这个集群上的 24,000 个 GPU 来训练 Llama 3 405B。
训练优化拓扑。 许多网络拓扑是与分布式训练算法共同设计的。铁路优化拓扑 [62] 增强了 GPU 到叶交换机的连接。在每个铁路中,跨不同服务器具有相同索引的 GPU 通过同一个叶交换机相互连接。这种配置通过减少数据流之间的网络干扰来提高集体通信性能。SuperPod 架构使用了一个铁路优化网络,能够连接超过 16,000 个 GPU [128]。字节跳动在其 MegaScale 系统设计中采用了三层铁路优化网络,连接了超过 10,000 个 GPU [71]。然而,铁路优化网络设计可能效率较低,因为它们需要将 GPU 连接到远处的交换机,这需要昂贵且耗电的光收发器。这些光组件增加了功耗和热量,导致网络故障率更高,这对于分布式 LLM 训练来说是重要的。阿里巴巴进一步优化了铁路优化拓扑,采用了一个名为 HPN [63] 的 2 层双平面架构。这种架构采用了最新的 51.2Tbps 单芯片交换机,支持一个 tier1 网络中的 1,000 个 GPU,并且在一个豆荚内支持多达 15,000 个 GPU。对 GPT/OPT-175B 模型训练的网络流量分析显示,99% 的 GPU 对不承载任何流量,并且不到 0.25% 的 GPU 对处理流水线/张量并行和数据并行流量 [64]。基于这些发现,铁路唯一拓扑 [64] 在铁路优化网络中消除了不同铁路之间的连接。每个铁路通过专用但独立的 Clos 网络连接。不同铁路上的 GPU 之间的通信通过通过内部芯片到芯片互连转发数据来管理。这种方可以有效降低成本,同时保持性能。HammingMesh [65] 将 GPU 组织成具有 2D-环面拓扑的组,并通过稀疏连接的交换机连接这些 2D-环面组。这种设计旨在在不牺牲训练性能的情况下节省成本。鉴于 GPU 仅通过 PCIe 连接,BiGraph [129] 提出了一种新的网络架构,将节点内 GPU 通信导出到节点外,绕过 PCIe 带宽瓶颈。它采用两层网络,通过 Clos 架构相互连接,具有支持应用控制流量路由的唯一最短路径。
可重配置拓扑。 可重配置网络可以动态调整,以优化通信模式,提高训练性能。它们通常使用光交换和自定义配置来提高网络基础设施的带宽利用率、灵活性和可扩展性。在硅光子(SiP)接口的驱动下,SiP-ML [66] 推进了两个主要架构:SiP-OCS 和 SiP-Ring。SiP-OCS 采用完全连接的配置,通过使用商业上可获得的光电路交换机,通过 Tbps SiP 接口将 GPU 连接到所有交换机,从而最大化带宽。相反,SiP-Ring 利用无交换机的环形配置,通过在 SiP 接口中集成微环谐振器来减少重配置延迟。Wang 等人提出了 TopoOpt [67],用于共同优化网络拓扑和并行化策略,在分布式训练中。这种方法不仅优化了计算和通信需求,还解决了网络拓扑的物理层。TPUv4 [42] 具有光学电路交换机(OCS),允许动态重配置基于 3D-环面的互连拓扑,优化了 LLM 训练的多样化和密集通信模式的数据流。例如,使用 512 个芯片,TPUv4 提供了 4x4x32 或 8x8x8 等 3D-环面拓扑的灵活性。
负载均衡。 LLM 训练的网络流量特点是少量的大象流。具体来说,LLM 训练表现出由于梯度同步导致的网络流量的周期性突发。每个突发都需要大量的网络带宽。此外,参与 LLM 训练的每个计算节点生成的连接非常少 [63]。传统的负载均衡技术,ECMP(等成本多路径路由)[68],使用哈希算法将流量均匀分布在等价路径上,例如在 Clos 拓扑中从叶交换机到脊柱交换机的路径。然而,这种基于哈希的方案在处理 LLM 训练流量时是低效的,因为它们由少量的大象流组成。当多个大象流被路由到同一链路时,可能导致拥塞和高延迟。为了解决大规模 GPU 集群中的负载均衡挑战,已经开发了各种策略。在 Llama 3 405B 训练期间,集体库在两个 GPU 之间建立了 16 个网络流,而不是单个流,从而减少了每个流的流量并增强了负载均衡的机会 [9]。此外,增强的 ECMP(E-ECMP)协议通过在 RoCE 数据包头中对更多字段进行哈希,有效地在不同的网络路径上分布这 16 个流。数据包喷涂 [69] 将来自一个流的数据包分散到所有可用的并行链路上,这可能导致数据包乱序。NIC 需要处理乱序的 RDMA 数据包。基于 LLM 训练的流量模式,Ethereal [70] 展示了贪婪地为每个流分配路径可以均匀地将负载分布在所有网络路径上,并解决 ECMP 哈希冲突问题。在大规模 GPU 集群中,HPN [63] 通过识别精确的不相交等路径并在集体通信库内平衡负载来实现高效的负载均衡。MegaScale [71] 表明,铁路优化拓扑也可以缓解 ECMP 哈希冲突。
拥塞控制。 无损传输在 RDMA 集群中至关重要。基于优先级的流控制(PFC)[72] 是一种流控制机制,可以防止数据包丢失。当 PFC 启用的下游设备上的队列发生拥塞时,设备会指示上游设备在队列中停止流量,从而确保零数据包丢失。由于 PFC 是一种粗粒度机制,它可能导致头部阻塞 [130],这显著降低了网络吞吐量。为了解决这些挑战,开发了各种通用拥塞控制方案。这些技术包括 TIMELY [73]、数据中心量化拥塞通知(DCQCN)[75], [76]、Swift [74]、高精度拥塞控制(HPCC)[77]、边缘排队数据报服务(EQDS)[78] 和鲁棒拥塞控制(RoCC)[79]。这些方案监控网络拥塞,调整数据速率以缓解拥塞,并恢复速率以最小化吞吐量降低。当有并发训练作业时,许多拥塞控制方案利用突发性和周期性的流量模式有效地交错网络流量。MLTCP [80] 基于关键见解交错作业的通信阶段:训练流量应该根据每次训练迭代中发送的字节数调整其拥塞窗口大小。CASSINI [81] 通过考虑不同作业的通信模式来优化网络链接上的作业放置。MLT [82] 利用 LLM 训练的特点,其中早期层的梯度不如后期层的梯度重要,较大的梯度比较小的梯度更重要。因此,在通信拥塞的情况下,MLT 根据它们在开关级别中包含的梯度的重要性优先排队或丢弃数据包,以缓解通信拥塞问题。
存储系统在分布式 LLM 训练中起着至关重要的作用,需要满足几个关键要求。首先,它应该与 GPU 的计算能力相一致,以最大化它们的利用率,并避免由于存储瓶颈造成的资源浪费。其次,它应该支持大规模结构化和非结构化训练数据集的存储,并在分布式处理环境中具有可扩展性。此外,模型检查点的存储和检索在 LLM 训练中提出了挑战,要求系统满足由模型大小和训练时长决定的写入和读取带宽。最后,存储系统应该满足传统的企业级要求,如数据保护、高可用性和安全性。
在 LLM 训练中,模型检查点的尺寸非常庞大。随着参数数量的增加,需要写入的数据量也在增加,这要求存储系统提供更大的写入带宽。例如,对于一个具有 700 亿参数的 LLM,检查点大小为 980GB。为了管理大规模的 GPU 数据中心中的模型检查点,已经部署了多种存储系统。Meta 的分布式文件系统 Tectonic [83] 能够使数千个 GPU 同时保存和加载模型检查点,为广泛的训练操作提供高效和可扩展的存储解决方案 [131]。在字节跳动,HDFS [84] 被用于集中式模型检查点维护,确保规模上的一致性和可靠性 [71]。为了缓解检查点恢复期间的带宽瓶颈,一种常见的方法是指定一个单独的工作器从 HDFS 读取检查点分区,然后将其广播给共享相同数据的其他工作器。分布式对象存储,如 Ceph 对象存储 [85],由于它们没有层次目录树或名称空间,简化了一致性维护,从而提供了更容易的可扩展性。由于这些优势,对象存储已经广泛用于模型检查点存储。
LLM 训练的原始数据集是巨大的。LLaMA 3 在超过 15 万亿个 token 上进行了训练,这比 LLaMA 2 的数据集大七倍以上 [6]。每个 token 需要大约 2 个字节,相当于大约 30TB 的数据。准备训练数据集涉及广泛的预处理步骤,包括数据抓取和清洗,需要大量的实验。通常,在这些步骤中处理的数据量超过最终训练数据集大小的 100 倍 [132]。例如,WanJuan-CC 数据集 [132] 选择性地提取了大约 60 亿份文档,生成了大约 1 万亿个高质量 token,相当于 2TB 的数据大小,在丢弃了 99% 的原始数据后。因此,预计 LLM 训练的总数据量将超过数十 PB。并行文件系统,如 Lustre [86]、GPFS [87] 和 BeeGFS [88],经常部署在领先的高性能计算系统上,以确保高效的 I/O、持久存储和可扩展的性能。这些系统也广泛用于训练集群中的数据加载,为高效处理大规模训练数据提供了必要的基础设施。此外,对于文件系统来说,使工程师能够在使用数千个 GPU 的作业上执行交互式调试至关重要,因为代码更改需要立即对所有节点可访问 [131]。在大多数 LLMs 的训练中,每个 token 通常只遇到一次。然而,使用数据缓存仍然至关重要,以缓解数据加载期间的 I/O 瓶颈。这种策略涉及从较慢的后端存储预取训练数据到较快的缓存存储。Alluxio [89] 和 JuiceFS [90] 通过从底层存储系统(如 HDFS 或对象存储)有效缓存训练数据来增强 LLM 训练。Quiver [91] 支持跨多个作业和用户操作同一数据集时的缓存数据的透明重用。Fluid [92] 利用 Alluxio 进行数据缓存,并结合了一种机制,根据 I/O 条件启用缓存的即时自动扩展。
LLM 训练工作负载通常在大规模多租户基础设施(例如,GPU 集群、公共云)上运行,其中用户共享集群资源。有效的调度机制对于管理这些工作负载至关重要,确保资源的高效利用和任务执行 [133]。
与任务级调度(例如,流水线调度 [134]–[136])不同,它专注于单个作业执行的细粒度优化(§4.1.3),集群级调度旨在优化整个集群的资源分配和任务调度。我们根据它们主要优化的方面,将现有的集群级调度系统分为两种类型:工作负载调度和资源调度。
近年来,针对 DL 训练工作负载的调度器已经被积极探索 [93]–[95], [137]–[141]。为了提高资源利用率,通常实现了三个高级功能:(1) 异构感知调度器(例如,Gavel [96],Gandivafair [97])专注于在不同 GPU 代之间优化作业分配;(2) 作业打包调度器(例如,FGD [98],Lucid [99])启用细粒度的 GPU 共享,充分利用硬件能力;(3) 自适应缩放调度器(例如,Pollux [100],Sia [101])动态调整 GPU 数量以及训练超参数,以加速训练进度。然而,这些调度器是为通用 DL 工作负载设计的,可能不直接适用于 LLMs,因为 LLM 工作负载具有独特的特性 [23]。为了更好地管理 LLM 工作负载,一些最近的研究提出了针对 LLMs 的系统。Crius [102] 联合考虑了混合并行性(§4.1)和异构集群中的硬件亲和性。它调查了在集群调度级别集成自适应并行性配置的工作流程效率,为同时提高多个 LLMs 训练效率提供了重要机会。为了实现 LLMs 的高效超参数调整,Hydro [103] 将模型缩小到较小的代理模型进行超参数搜索,然后融合多个模型到一个单一实体中,以提高硬件利用率。此外,Hydro 通过将它们与流水线启用的 LLM 预训练任务交错,扩展了调整工作负载的资源,有效地利用了流水线泡沫。Acme [23] 进一步对 LLM 开发工作流程的工作负载混合进行了表征,并提出了一个系统,以有效调度与 LLM 训练相关的相关作业,包括解耦的评估调度,以便及时提供模型质量反馈,以及涉及 LLM 的故障诊断和自动恢复。
除了工作负载调度外,相关资源调度(例如,CPU、内存和网络)是集群级管理的另一个关键方面。对于网络,Cassini [81] 通过使用亲和图确定时间偏移值来调整不同作业的通信阶段,从而在不同作业的上下阶段期间交错带宽需求。HIRE [104] 引入了一个创新的数据中心交换机的网络计算调度系统,显著减少了网络绕路和尾部放置延迟。对于存储,SiloD [105] 将数据缓存和远程 I/O 视为一级资源进行联合分配,显示出显著的吞吐量改进。对于 CPU 和内存,Synergy [106] 通过优化 CPU 核心分配而不是依赖于 GPU 成比例分配,提高了训练效率。此外,一些工作专注于节能。EnvPipe [107] 利用了流水线中的时间扩展,通过降低 SM 频率来节省能源,同时拉伸了流水线单元的执行时间。Zeus [108] 自动配置批量大小和 GPU 功率限制,以提高训练期间的能效。Perseus [109] 引入了一个高效的基于图割的迭代算法,以获得大型模型训练作业的迭代时间 - 能量 Pareto 前沿。
LLMs 不断增长的规模需要大量的计算资源和内存容量。分布式训练,利用大规模 HPC 集群,已成为有效训练这些模型的关键方法。在本节中,我们调查了提出用于提高 HPC 集群 LLM 训练利用率的各种并行策略。我们将这些方法归类为三个主要组:混合并行、自动并行和异构并行。混合并行结合了多种手工制作的并行化策略,如数据并行、张量并行、流水线并行、序列并行和专家并行。自动并行根据模型和硬件特性自动确定最优的并行化策略。异构并行利用硬件或模型中的异构性进行有效训练。这包括利用不同类型的加速器或利用单个模型内的异构性(例如,RLHF 训练)来提高 HPC 集群上的整体训练效率。大多数当今最先进的并行化策略采用单一程序多数据(SPMD)编程模型,类似于 MPI 范式 [239],其中相同的程序在多个处理器上运行,每个处理器处理不同的数据 [225]。例如,数据、模型和序列并行使用 SPMD 编程模型。这种方法确保了操作的一致性和一致性,使其非常适合大规模、分布式训练环境。一些策略探索打破 SPMD 的限制,通过多程序多数据(MPMD)模型进一步提高资源利用率,其中不同的程序(或程序的不同部分)在不同的处理器上运行,处理数据或模型的不同部分 [225]。例如,流水线并行在不同的设备上运行 LLM 的不同部分。此外,自动并行和异构并行可以利用 SPMD 和 MPMD 模型来增加资源利用率。因此,我们根据并行发生的维度以及使用的计算设备是同构还是异构,而不是关注底层编程模型,来讨论这些方法。
混合并行通常结合多种手工制作的并行化策略,以划分 LLM 的不同可并行化维度。这些策略包括数据并行、张量并行、流水线并行和序列并行,如图 8 所示。数据并行、张量并行和流水线并行的组合也称为 3D 并行。
数据并行是分布式训练中最常用的并行化策略,由于其高可扩展性和易于实现。它遵循单一程序多数据(SPMD)编程模型。数据并行沿批量维度划分输入训练数据,每个 GPU 处理其分配的数据段,如图 8(a) 所示。在整个训练过程中,数据首先通过完整的模型权重逐层进行前向计算,然后按相反的顺序执行反向计算。每层生成的梯度将通过集体通信操作跨所有 GPU 聚合以进行优化器更新。数据并行包括各种分片策略,这些策略显著影响内存占用和通信开销。假设全局世界大小为 W(即设备数量),引入分片因子 F 来控制使用的分片策略 [146],定义为参数被划分的设备数量(1 ≤ F ≤ W)。我们有以下情况:
全复制(F = 1): 这种分片策略被简化为普通的数据并行。Pytorch-DDP [240] 和 Horovod [143] 在所有设备上完全复制模型,并使用 All-Reduce 进行梯度聚合。它们还将梯度划分为小桶,以重叠梯度通信和反向计算。全分片(F = W)。这种分片策略具有最低的内存消耗,但通信开销最大(比普通数据并行多 1.5 倍)。全分片策略完全分片模型,每个设备只持有 1/W 的模型参数。完整的权重和梯度在计算前按需通信和恢复,并在之后立即丢弃。ZeRO-3 [145] 采用每个参数分片来分片完整模型,并分别使用 All-Gather 和 Reduce-Scatter 进行无分片和分片通信。分片权重更新 [144] 也采用每个参数分片,但更侧重于在所有设备上分片冗余的参数更新计算。FSDP(完全分片数据并行)[146] 通过在模块单元粒度上分片模型参数来实现相同的功能,并提供更用户友好的 API。
混合分片(1 < F < W)。 在这种策略 [146] 中,所有设备被划分为 N×M 设备网格。模型参数沿网格的 N 维度分片,并沿 M 维度复制。MiCS [147] 调用 All-Gather 集体操作来收集分片参数,并使用 All-Reduce 聚合梯度。FSDP [146] 用 Reduce-Scatter 替换 All-Reduce,以减少内存和通信开销。与全复制和全分片相比,混合分片更灵活,可以通过调整 F 来提供内存消耗和通信开销之间的权衡,基于模型架构和硬件限制。
张量并行(图 8(b)),也称为层内模型并行,是一种提出的技术,用于在多个 GPU 上训练 LLMs。它沿多个维度划分每层的参数张量,有效地将模型参数分布在可用的 GPU 上。与数据并行不同,张量并行通常需要高带宽连接进行通信,因此更常用于单个 GPU 节点。张量并行可以按划分的维度分为 1-D [20]、2D [148]、2.5-D [149] 和 3-D [150] 并行。基于 Transformer 的 LLMs 的 MLP 和自注意力模块中都有两个参数矩阵。Megatron-LM [20] 首先采用 1D 张量并行将第一个参数矩阵沿其列划分,第二个参数矩阵沿其行划分。它复制每个划分模块的输入和输出张量,并引入两个 All-Reduce 集体通信,以适应多个 GPU 上的 LLM。受到可扩展通用矩阵乘法算法(SUMMA)[241] 和 Cannon 算法 [242] 2-D 并行矩阵乘法的启发,Optimus [148] 进一步将输入和参数张量沿 2 个维度划分,以提高 1-D 张量并行的通信和内存效率。Tesseract [149] 扩展了 2.5-D 矩阵乘法方法 [243],该方法被提出以提高 Cannon 算法的效率,用于 LLM 训练,并提出了 2.5-D 张量并行,以克服由于模型尺寸增加而导致的不必要通信过多。3-D 张量并行 [150] 采用并改进了 3-D 并行矩阵乘法算法 [244],用于线性操作,并在多个设备上实现 LLM 训练的完美负载平衡。
流水线并行(图 8©) [151],也称为层间模型并行,被提出以适应大型模型在多个 GPU 上,特别是在不同节点上的分布。流水线并行将模型的层划分为多个阶段,每个阶段由模型中的一组连续层组成,并映射到一组 GPU。与通常需要高带宽连接进行通信的张量并行不同,流水线并行只需要在指定的切割点交换中间张量,因此通信需求较少。因此,流水线并行适用于在连接带宽较小的多个 GPU 节点上扩展 LLM 训练。例如,Strati 等人 [245] 采用流水线并行充分利用地理分布式资源,以克服 GPU 短缺的问题。由于不同阶段的数据依赖性,流水线并行通常将输入数据拆分为多个微批次以实现流水线化,从而有效训练大型模型。然而,它带来了两个显著问题。首先,流水线泡沫问题由于等待前一阶段的输出而降低了 GPU 的利用率。其次,不同阶段之间的内存消耗不平衡,因为前几个阶段需要持有更多的活动微批次,以实现更好的流水线化和更高的利用率。我们下面详细说明每个问题。
流水线泡沫。 高效的微批次调度算法可以减少流水线泡沫。GPipe [151] 引入了一个填充 - 排空调度,一次性注入所有微批次进行前向传递执行,然后是反向传递。Gpipe 由于前向和反向传递的预热和冷却,引入了显著的流水线泡沫。PipeDream [134], [135] 引入了一个 1F1B(1 个前向 1 个反向)调度,一旦相应的前向传递完成,就执行微批次的反向传递,以减少异步场景中的流水线泡沫。DAPPLE [152] 采用早期反向调度,在每个阶段的开始首先注入固定数量的微批次,然后通过轮询交错前向和反向传递。交错 1F1B [153] 适应 1F1B 调度,但为每个 GPU 分配多个阶段(即循环流水线放置)。流水线泡沫以更高的通信和峰值内存消耗为代价减少了。Chimera [136] 引入了一个双向流水线来减少泡沫,通过权重复制。Hanayo [157] 进一步提出了一个波浪形流水线,将多个对称阶段分配给一个 GPU,以提高流水线利用率。零泡沫 [156] 将反向计算分为两部分:激活和参数梯度计算。它按 1F1B 安排前向和激活梯度计算,然后用参数梯度计算填充泡沫,这以更高的峰值内存消耗减少了泡沫。Breadth-First [159] 在循环流水线放置中一次性运行所有微批次,以减少与分片数据并行结合时的通信开销。TeraPipe [154] 沿序列维度拆分微批次,并利用更细粒度的标记并行来减少流水线泡沫。然而,TeraPipe 的内存开销很大,因为它基于 GPipe 调度。Seq1F1B [158] 将序列拆分成块,并使用 1F1B 调度来降低峰值内存消耗,同时实现低流水线泡沫率。DynaPipe [160] 使用动态微批次方法对具有可变长度输入的 LLMs 进行多任务训练。它引入了一种内存感知的自适应调度算法和提前通信计划,以进一步减少流水线泡沫率。Tessel [155] 是一个两阶段方法,包括重复模式构建和调度完成,以自动搜索指定分区策略的高效流水线调度。DISTMM [161] 启动双倍微批次以绕过由多模态训练的大批量需求引起的依赖性障碍,从而减少空闲周期。GraphPipe [162] 保持 DNN 图拓扑结构,并将其划分为可以同时执行的阶段,以提高流水线利用率并减少内存消耗。
内存不平衡。 流水线并行通常向开始阶段注入更多的微批次,以提高流水线利用率,导致这些阶段的激活内存消耗更高。为了解决这个问题,BPipe [163] 和 MPress [164] 采用 D2D(设备到设备)传输,在运行时将高负载 GPU 上的中间激活张量交换到负载较轻的 GPU。MPress 还结合了激活重计算以减少内存占用。Chimera [136] 引入了一个双向流水线,将两个不同方向的流水线结合起来,以实现更平衡的内存消耗。每个 GPU 持有两个对称阶段,导致权重复制。Hanayo [157] 将双向流水线变为两个数据并行流水线,通过对称地将多个阶段分配给一个 GPU 来消除权重复制,并实现平衡的内存消耗。V-Shape [166] 将模型划分为设备数量两倍的阶段,将两个阶段的一半以相反顺序放置。通过变化阶段之间的偏移量,V-Shape 在峰值内存消耗和泡沫利用率之间做出权衡。mCAP [165] 使用增量分析方法,根据峰值内存使用情况在 GPU 上均匀划分模型。峰值内存消耗限制了流水线并行中活动的微批次数量,从而限制了其效率。激活重计算可以有效地减少峰值内存消耗。Varuna [167] 结合了流水线并行和激活重计算来实现这一目标。它为给定的流水线设计了基于静态规则的调度,并列出了机会主义策略,以隐藏抖动并减少泡沫。静态调度是基于包括激活重计算时间、激活内存管理和反向计算优先级在内的约束生成的。为了解决内存不平衡问题并降低重计算开销,AdaPipe [168] 采用了自适应重计算,支持不同阶段的不同重计算策略,并根据 1F1B 调度自适应分区,以平衡每个阶段的计算。
当今 LLMs 的上下文窗口迅速增长,最强大的 LLM 可以支持数百万 token [7]。这种超长序列导致 LLM 训练的内存和计算需求显著增加:激活的内存占用线性增加,注意力机制的复杂度呈二次方增长。在反向传播中重新计算激活可以减少峰值内存消耗,但也引入了显著的开销(全部重新计算时为 30%)。大张量并行度引入了显著的通信开销。序列并行(图 8(d)) [169], [170] 被提出以适应长序列训练,并在内存容量内高效地在多个 GPU 上分布计算。它将输入数据沿序列维度划分为多个块,每个块被送入一个 GPU 进行计算。由于序列并行复制了模型参数,因此通常与张量和流水线并行结合使用,以扩展 LLM 训练。当与张量并行一起使用时,序列并行在多个 GPU 上分布了注意力的内存和计算,但引入了非张量并行区域内的冗余内存消耗和计算。Megatron-SP [170] 沿着序列维度拆分这些计算,减少了冗余激活计算和内存消耗,而没有增加通信。尽管序列并行在多个 GPU 上划分了内存、计算和通信,但二次因果注意力在训练效率方面仍然存在显著挑战,包括键 - 值张量通信开销、IO 感知注意力计算开销和由于因果注意力掩码导致的 GPU 之间负载不平衡。大多数序列并行方法的注意力是基于环形的 [169], [173], [175], [176], [178], [179]。Ring Self-Attention [169] 利用序列并行,并以环形通信方式计算自注意力,以扩展 LLM 训练的上下文窗口。它首先在 GPU 之间传输键张量以计算注意力分数,然后根据传输的注意力分数和值张量以类似方式计算自注意力输出。DistFlashAttn [176] 同时传输键 - 值张量块,利用 IO 感知 FlashAttention [115] 内核,并通过用早期 token 的空闲周期填充后期 token 的计算来平衡不同 GPU 的计算。Megatron Context Parallel [173] 也利用 FlashAttention 内核,并消除了由低三角因果掩码引起的不必要计算。它通过交换一半的块与对称 GPU,进一步平衡了 GPU 之间的计算。DistFlashAttn 和 Context Parallel 还在单独的 CUDA 流中重叠了键 - 值张量通信和注意力计算。Striped Attention [178] 通过为每个 GPU 分配一个在序列中均匀分布的 token 子集来解决不平衡问题,而不是连续的块。BurstAttention [179] 在每个 GPU 上使用 FlashAttention 计算注意力,并使用双缓冲区来重叠通信和计算。Blockwise Ring Attention [175] 将 Ring Self-Attention [169] 扩展到块状注意力,通过以小块计算注意力来减少内存占用。受到 N 体模拟的启发,WallFacer [180] 首先将 GPU 划分为子组,并通过异步 AllGather 在每个子组内复制查询和键 - 值张量。注意力计算利用多个环形 P2P 通信来提高效率。最后需要一个异步 ReduceScatter 来分发注意力输出。DeepSpeed-Ulysses [172] 与以前的基于环的方法不同,通过拆分头维度而不是序列维度,并利用 All-to-All 将分区维度从序列转移到头。DeepSpeed-Ulysses 可以与现有的注意力实现(例如,FlashAttention)无缝结合,并且 GPU 之间的工作负载自然平衡。然而,DeepSpeed-Ulysses 的并行度受到头数的限制,特别是对于使用 MQA [11] 和 GQA [12] 的 LLMs。LoongTrain [174] 和 USP [171] 是并发工作,集成了 DeepSpeed-Ulysses 和 Ring Attention 的优势。它们将 GPU 组织成二维网格,形成混合的 ulysses 和 ring 风格的进程组。在训练期间,它们首先在 ulysses 组之间执行 All-to-All 以将分区从序列维度切换到头维度,然后在 ring 组中使用 Ring-Attention 进行注意力计算。LoongTrain 进一步提出了 Double-Ring-Attention,以充分利用节点间通信的可用带宽,并重叠通信与计算。DSP [177] 根据多维变换器(如 DiT [246])的计算阶段动态切换并行维度。
Mixture-of-Experts (MoE) 目前是 LLMs 中最受欢迎的稀疏模型。虽然 MoE 显著增加了 LLMs 中的参数数量,但由于条件计算,它并没有大幅增加计算成本 [247]。MoE 的基本框架,如图 9 所示,由多个专家网络组成,这些网络处理不同的训练数据子集,以及一个门网络,该网络应用路由算法将输入 token 分配给不同的专家网络。MoE 能够训练具有超过万亿规模参数的大型模型,并且声称已在流行的 LLM 模型中使用,如 Mixtral 8x7B [248] 和 DeepSeek2 [13]。
稀疏激活。 随着模型大小的增加,并非所有专家都能适应并训练在单个设备上。因此,GShard [15] 将 MoE 的思想扩展到分布式设置中的 Transformers,专家分布在不同的工作器上,并通过 All-to-All 通信进行协作,如图 9 所示。随后的专家并行研究通常遵循相同的范式。例如,Switch Transformer [181] 在 T5 模型上结合了分布式 MoE 训练的设计。但与 GShard 使用的前 2 个路由算法不同,Switch Transformer 将每个 token 路由到仅有的前 1 个专家,以最大化计算效率。此外,DeepSpeed-MoE [183] 提出了一种新的分布式 MoE 架构,它在每个工作器中应用共享专家,并将更多的专家放置在更深层,以平衡通信成本与训练精度。专家并行可以有效地与常规的 3D 并行集成。例如,GShard、Switch Transformer 和 DeepSpeed-MoE 都把专家并行视为混合并行的正交维度。为了有效的混合训练,DeepSpeed-TED [249] 提出了一种混合并行算法,结合了数据、张量和专家并行,以启用 MoE 模型的训练。作者将 MoE 参数划分为预定义大小的'瓦片',以避免优化器内存峰值过高,并提出了像 Duplicate Token Dropping (DTD) 和激活检查点这样的通信优化,以消除 All-to-All 通信中的重复数据。然而,由于 MoE 的动态特性,选择最优的混合并行计划具有挑战性,并且在运行时在不同的并行策略之间切换也会带来相当的开销。因此,一些研究如 Tutel [182] 设计了一种自适应并行策略切换算法,它对所有可能的最优策略应用相同的分布模型布局,并且可以在每次迭代中动态切换并行策略,而不会产生任何额外开销。由于一般矩阵乘法(GeMMs)要求所有专家的输入大小一致,现有的 MoE 训练框架通常执行 token 丢弃和填充,以匹配相同的专家容量,这会浪费计算。Megablocks [185] 通过实现块稀疏矩阵乘法来优化分组 GeMMs,并支持单个内核中不同批量大小的专家计算,以避免 MoE 训练中的不必要 token 丢弃。另一个支持分组 GeMMs 的框架是 ScatterMoE [184],它实现了 ParallelLinear 内核,该内核融合了分组 GeMMs 和分散的读写操作,以减少 top-k(k ≥ 2)门控的内存占用。
通信优化。 专家并行中的 All-to-all 通信可能会严重影响 MoE 的训练效率,特别是在网络环境不佳的情况下。现有的分布式训练系统尝试通过使通信任务与计算任务重叠来优化 MoE 的性能,以便一些通信成本可以被隐藏。例如,Tutel [182] 沿专家容量维度将输入张量划分为组,并重叠不同组之间的计算和通信,以隐藏 All-to-All 开销。FasterMoE [191], [192] 使用与 Tutel 类似的策略,但沿专家维度拆分张量。此外,Tutel [182] 还通过在节点内部聚合小消息为单个大块来优化 All-to-All 内核实现,然后交换不同节点之间的数据。这种优化也用于 FasterMoE 和 ScheMoE [187]。基于 Tutel 中的重叠策略,PipeMoE [186] 基于工作负载对通信和计算任务的执行时间进行建模,并设计了自适应算法,以找到最小化训练时间的最佳分区数。ScheMoE [187] 考虑了 All-to-All 通信前的数据压缩方法,并将耗时操作模块化,包括数据压缩、集体通信和专家计算。ScheMoE 然后提出了一种自适应最优调度算法,以流水线方式调度通信和计算操作,以提高训练效率。专家并行通常与 MoE 训练中的其他并行策略相互作用。通过细粒度任务调度,可以减少通信开销。例如,Lina [188] 系统地分析了分布式训练和推理期间 MoE 的 All-to-All 开销,并发现当 All-to-All 与 AllReduce 操作重叠时,All-to-All 延迟会延长。Lina 提出优先考虑 All-to-All 而不是 AllReduce,以提高其带宽并减少分布式训练中的阻塞期。此外,Lina 结合了张量划分和流水线,以类似于 Tutel 的方式执行微操作调度。Lina 还根据专家的受欢迎程度在推理期间动态调度资源,以最小化开销。Janus [189] 设计了一种以数据为中心的范式,将数据保持在原地,并根据参数服务器上的 GPU 移动专家。以数据为中心的范式使用细粒度异步通信,并允许专家使用非阻塞通信原语(如 pull)在 GPU 之间移动。Janus 实现了一种拓扑感知策略,以有效地在节点之间拉动专家,并支持专家预取,将所有外部专家拉到本地 CPU 内存中。
有一些研究从模型 - 系统共同设计的角度优化 MoE 训练。例如,TA-MoE [190] 提出了一种大规模 MoE 训练的拓扑感知路由策略。TA-MoE 将调度问题抽象为一个优化目标,以在不同拓扑下获得目标调度模式,并根据调度模式设计了拓扑感知辅助损失。这种方法自适应地路由数据以适应底层拓扑,而不会牺牲模型精度。
负载平衡。 由于 MoE 的稀疏和条件计算特性,一个受欢迎的专家可能会收到比其他专家更多的 token(通常是由较差的路由算法引起的),导致严重的负载不平衡,并影响 MoE 的训练效率。FasterMoE [192] 提出了影子专家方法,根据前几次迭代的工作负载动态广播受欢迎专家的参数到所有其他 GPU。通过将受欢迎专家的工作负载分散到不同的设备上,影子专家方法减少了倾斜专家受欢迎程度的影响。SmartMoE [193] 采用两阶段方法搜索负载平衡的最优并行计划。首先,SmartMoE 设计了一个数据感知性能模型,将并行计划划分为池,在池内的并行模式切换成本相对较低。然后,SmartMoE 可以切换到适当的并行性(在 SmartMoE 中称为专家放置),以保持在线阶段的负载平衡。FlexMoE [194] 发现专家到设备的映射分布在短期内不会显著变化,因此它引入了细粒度的复制专家并行,将重专家复制到多个设备。FlexMoE 监控数据工作负载,并使用三种放置调整原语(即扩展、收缩、迁移),如果超出平衡比率,则生成最优放置解决方案。Prophet [195] 提出了一种系统化、细粒度和高效的大规模 MoE 模型的负载平衡训练方法。以 MoE 模型、设备池和 token 分布为输入,Prophet 的计划器迭代搜索和评估专家放置,最后输出一个平衡良好的专家放置。此外,Prophet 使用逐层调度策略隐藏了这些资源分配操作的开销。
给定任意 DNN 模型和 GPU 集群,存在着大量的并行选项,包括单个层的划分及其划分程度。设计可以最大化训练效率的手工混合并行方法是一个耗时且知识密集型的过程,需要专家对模型架构、硬件特性和并行化策略中涉及的复杂权衡有深入的理解。此外,有效实现最优并行化策略通常需要大量的人力努力。为了应对这些挑战,自动并行作为一种有希望的解决方案出现了,它寻求自动确定给定 DNN 模型在特定 GPU 集群上最有效的并行化策略。通过利用复杂的算法和启发式方法,自动并行系统可以分析模型架构、硬件规格和性能特征,以识别数据、张量和流水线并行等并行化技术的最佳组合。这种方法简化了在各种模型和基础设施上优化分布式训练的过程,提高了整体效率并减少了手动工作。此外,自动并行可以适应不断变化的硬件配置和模型架构,自动调整并行化策略以保持最佳性能。在下面,我们根据目标模型架构将现有的自动并行系统分为通用和特定于 Transformer 的框架。
通用自动并行框架专注于在特定计算集群上自动并行化通用 DNN。这些框架通常遵循三个步骤:(1) 定义并行化策略的搜索空间;(2) 开发性能模型以衡量不同策略的训练效率;(3) 设计算法以有效识别最优并行化策略。下面我们根据它们所涵盖的搜索空间来研究不同的方法。一些工作已经探索了 DNN 训练优化的混合数据和流水线并行策略的搜索空间。这些方法专注于自动划分 DNN 并设计流水线调度以提高流水线利用率。PipeDream [134] 通过最慢阶段的执行时间来衡量流水线分区的效率,并开发了一种动态规划算法,通过最小化最慢阶段来均匀划分 DNN。DAPPLE [152] 构建了一个分析模型来估计一个分区策略的执行时间,并使用动态规划来确定最优的流水线分区。AutoPipe [205] 构建了一个模拟器来模拟流水线执行,并提出了一种启发式算法来获得平衡的分区。AutoPipe 还自动拆分微批次以减少热身阶段的延迟。一些设备放置方法 [206]–[208] 使用强化学习来预测流水线并行的最优操作员放置。
研究人员还通过沿不同维度划分操作员来探索自动化数据和模型并行。OptCNN [198] 沿其输出张量的可划分维度划分操作员,并使用分析性能模型选择最优的并行化策略,包括可并行化的维度和并行度,这定义了如何在不同设备上并行化单个层。FlexFlow [199] 进一步将搜索空间扩展到 Sample-Operator-AttributeParameter (SOAP),其中几乎包括了输入和输出张量的所有可划分维度,并引入了一种新的执行模拟器进行准确的性能建模。FlexFlow 使用 MCMC 抽样有效地找到最优的并行化策略。Tofu [200] 和 HyPar [201] 开发了动态规划算法,通过最小化总通信成本而不是端到端性能,来确定每个操作员在混合数据和模型并行空间中的最优分区。TensorOpt [204] 在给定内存预算下使用前沿跟踪算法优化并行化策略。AutoMap [202] 采用蒙特卡洛树搜索(MCTS)选择由 PartIR [212] 定义的一系列分区规则,通过学习得到的评分器为一组选定的重要操作员传播整个并行化策略。最近的工作还设计了自动化数据、模型和流水线并行的方法。Piper [209] 设计了一个两级动态规划方法,以找到最优的混合数据、张量和流水线并行结合激活重计算。它首先将模型划分为流水线的小分区,然后在每个分区内拆分操作员。Alpa [21] 将并行化视为两个层次:操作员间和操作员内并行化,然后自动派生每个并行化级别的有效并行执行计划。Unity [210] 通过将它们表示为统一的并行计算图上的替换,联合优化并行化和代数变换。Aceso [211] 提出了一种迭代瓶颈缓解方法,显著减少了搜索时间。它在每一步识别性能瓶颈,并调整策略以缓解瓶颈,直到收敛。nnScaler [213] 引入了三个原语,以启用与任意分区和分区模型的空间 - 时间调度的搜索空间的组合。领域专家可以对原语应用约束,以构建有效且小的搜索空间,这些空间可以以低开销自动探索。AutoDDL [215] 通过迭代更新每层的 SBP [214] 分布,定制坐标下降算法,并快速发现具有近乎最优通信成本的最优策略。通用自动并行框架需要有效的系统支持各种并行化策略,以及用于发现最优并行化策略的快速优化算法。这是因为并行化通常涉及复杂的计算和通信操作员,特别是对于模型并行化,它划分了操作员。以前的工作已经开发了支持各种并行化策略的高效系统,要么通过构建在现代 DL 框架 [21], [213] 之上,要么从头开始实现 [199]。Mesh-TensorFlow [196] 观察到实现并行化策略的固有复杂性,并首先提出将设备集群抽象为多维网格,并将并行化抽象为划分迭代空间(即张量维度)。通过将张量和网格维度映射,可以轻松实现高性能的混合数据和模型并行化策略。例如,数据和模型并行化分别划分批次和隐藏维度。GSPMD [197] 进一步提供了一种统一的方式来实现各种通用并行方案,基于 JAX [250] 和 XLA [251] 的简单张量分片注释。OneFlow [214] 提出了 SBP(分割、广播、部分值)抽象进行分区,并允许用户指定张量的放置和 SBP 签名,以实现不同的并行化策略。PartIR [212] 将模型与其分区解耦,并为用户设计了一个编译器堆栈,通过计划逐步组合 SPMD 分片策略。与 TVM [252] 类似,Slapo [203] 定义了一组全面的调度原语用于并行化和子图优化,如操作员融合和激活检查点。这些调度与执行分离,并保留原始模型结构以进行逐步优化。
4.2.2 特定于 Transformer 的框架
由于 LLMs 基于 Transformer 架构,最近的工作为 Transformer 量身定制了自动化系统。DeepSpeed-Autotuning [216] 自动调整系统旋钮,以找出用户定义的调优空间中的性能相关配置,包括并行度。Galvatron [217] 设计了一个动态规划算法,以生成最有效的混合数据、张量和流水线并行策略。Merak [218] 引入了一个非侵入性的自动模型分区器和高性能的 3D 并行运行时引擎,以提高可用资源的利用率。Colossal-AI [219], [253] 为混合数据、张量、序列和流水线并行提供了统一的接口。Galvatron-BMW [220] 扩展了 Galvatron 的空间,包括分片数据并行和激活重计算,并在考虑内存消耗和计算的同时搜索最优策略,以最大化硬件利用率。
LLM 训练的不断增长的计算需求推动了异构硬件的发展,利用多样化的计算资源和全球分布的设备。这种异构性也反映在模型架构中,特别是通过人类反馈的强化学习(RLHF)。利用异构硬件和多样化的模型架构已成为高效训练 LLMs 的关键。
LLM 训练的大规模计算需求推动了加速器的发展,导致集群具有混合设备类型和不均匀的互连带宽。此外,现代数据和计算集群通常由于电力短缺等因素而分布在全球范围内。这些现象激发了采用异构并行,利用多样化的计算资源和地理分布的设备来加速 LLM 训练。
一些工作利用 CPU、GPU 和专用加速器等异构计算资源来提高 LLMs 的性能。这些设备的计算能力、内存容量和互连带宽的差异为有效 LLM 预训练带来了挑战。HetPipe [221] 将异构集群划分为多个虚拟工作区。每个虚拟工作区使用流水线并行处理小批量数据,并使用异步数据并行来提高吞吐量。AccPar [222] 提出了灵活的张量划分,以平衡不同加速器的计算,并使用动态规划自动决定异构设备上的张量划分。Whale [223] 提出了一个统一的抽象,以简化在异构集群上进行大型模型并行训练的努力。它通过自动图优化无缝适应异构 GPU,并通过硬件信息平衡工作负载。AMP [224] 使用异构感知性能模型来找到最优的混合数据、张量和流水线并行策略。HPH [226] 根据计算 - 通信比率的降序排列不同的 GPU,并将模型划分问题表述为整数规划问题,以最小化迭代时间。Pathways [225] 采用分片数据流模型和异步帮派调度,在异构集群上有效执行 ML 模型。SDPIPE [227] 引入了一种半去中心化方案,将通信模型同步去中心化,并将组调度过程集中化,以利用异构设备进行流水线并行。HAP [228] 使用 A*基于搜索算法生成最优的张量分片策略、异构设备上的分片比例和分布式训练的通信方法。PipePar [229] 提出了一种动态规划算法,考虑 GPU 的异构性和网络带宽,将模型划分为流水线阶段。
一些其他工作探索了以低网络带宽为特征的地理分布设备,以提高训练效率。Yuan 等人 [230] 将 LLM 划分为计算任务单元,并提出了一种新的调度算法,以有效利用由慢速异构网络连接的一组异构设备进行混合数据和流水线并行。SWARM 并行 [231] 将模型划分为等大小的阶段,并优先将输入路由到具有较低延迟的稳定对等体,以实现工作负载平衡。它还自适应地在阶段之间移动设备,以最大化训练吞吐量。FusionAI [232] 将训练计算图(DAG)拆分为子图(subDAG),并生成一个负载平衡的任务调度,以利用低带宽连接的异构消费级 GPU 进行流水线训练。还可以利用像 CocktailSGD [254] 这样的通信压缩方法,在低带宽集群中高效地训练 LLMs。
图 10: RLHF 的一个例子。推理过程:1 演员模型从给定的查询中生成响应。2 评论家模型、奖励模型和参考模型使用查询和响应对通过推理生成训练所需的值、分数和 KL 散度。训练过程:3 演员模型和评论家模型使用推理过程中收集的数据通过梯度下降更新它们的权重。
在 LLM 训练过程中,异构性不仅反映在硬件上,还反映在模型中。训练可能涉及几个不同模型的交互。一个具体的例子是通过人类反馈的强化学习(RLHF)。RLHF 是一种训练方法,旨在使 AI 系统更贴近人类的偏好 [255],利用人类在判断适当行为方面的优势,而不是展示。这种方法特别受到关注,尤其是用于微调大型语言模型。然而,由于近端策略优化(PPO)[256] 算法的特殊性,RLHF 训练引入了模型异构性,使得 RLHF 训练过程与预训练和监督微调非常不同。原则上,RLHF 包括三个不同的阶段:第一阶段是监督微调,第二阶段是奖励模型的训练,第三阶段是 PPO 训练。模型异构性在第三阶段呈现,如图 10 所示。PPO 训练阶段包括两个不同的过程,即生成数据的推理过程和更新演员模型和评论家模型权重的训练过程。PPO 训练通过这两个过程的协作来执行。此外,训练阶段引入了更高的内存成本,因为我们需要同时提供多个自回归生成模型和奖励模型的副本,以及更多的时间成本,因为我们必须等待经验生成完成后才能更新权重。已经提出了许多框架来支持 RLHF 训练。例如,DeepSpeed-Chat [233] 使用混合引擎在训练和推理之间无缝切换模型分区,如在推理期间使用张量并行来提高吞吐量,在训练期间使用 ZeRO [145] 或 LoRA [257] 来提高内存利用率,为 RLHF 训练提供了卓越的系统效率。HuggingFace TRL [234] 可以充分利用各种参数高效微调(PEFT)方法,如 LoRA 或 QLoRA [258],以节省内存成本,并使用由 unsloth [259] 设计的专用内核来提高 RLHF 的训练速度。ColossalAI-Chat [253] 是另一个端到端 RLHF 训练框架,也支持 LoRA,并支持使用 ZeRO [145] 来减少内存冗余。然而,上述工作采用了平面化策略进行模型放置,即将 RLHF 中的四个模型放置在同一设备上,然后使用 ZeRO 或 LoRA 等方法最小化内存成本。但是,仅使用 ZeRO 在训练较大模型时会导致内存瓶颈,而使用像 LoRA 这样的高效参数微调策略会损害模型性能。为了解决这个问题,OpenRLHF [235] 使用 Ray [260] 和 vLLM [261] 将奖励模型分布到不同的设备上,避免将所有四个模型都放在 PPO 的同一设备上。类似地,Adpative Placement and Parallelism (APP) 框架 [236] 提出了两种其他模型放置策略,即交错策略和分离策略。它捕捉到生成部分和训练部分在 PPO 训练期间可以独立运行的事实,并且一些串行化可以通过将它们放置在不同设备上来消除,这引入了额外的通信,但可以与计算很好地重叠。同时,还有一些工作将前两个阶段中的并行策略以细粒度调度的方式应用到 RLHF 的第三阶段。例如,ReaLHF [237] 通过重新分配参数,在第三阶段的不同子阶段之间切换最合适的并行模式,大大扩展了优化空间。PUZZLE [238] 根据不同阶段的亲和性重新安排任务执行顺序,以便具有更好亲和性的阶段可以有效覆盖执行并提高训练效率。
当今的 AI 加速器在 FLOPs 方面提供了前所未有的计算能力。然而,要充分利用这些 FLOPs,需要复杂的优化技术。本节介绍了 LLM 训练的系统和计算优化技术。我们首先详细讨论操作符优化,包括核心注意力操作符优化和通过编译器进行的自动优化。基于利用大量并行性和针对底层硬件特性的有效多级内存访问,基于操作符和计算图的显著性能获得了。其次,详细介绍了混合精度训练,其中计算通过降低精度而加速。16 位浮点混合训练已成为大多数训练系统的事实上的方法。低至 1 位的低位固定点也已被研究和用于高训练效率。
操作符优化可以分为手动优化和自动优化。手动优化主要关注注意力操作符,而自动优化应用更广泛。
注意力作为 Transformer 架构的核心,在 LLMs 的训练效率中起着至关重要的作用。给定一个查询 q 和列表 k1、k2、…、kn 以及值 v1、v2、…、vn,其中 q、ki、vi ∈ Rd,注意力计算如下:
自注意力相对于序列长度表现出二次时间和内存复杂度。由自注意力引起的大量内存消耗和对高带宽内存(HBM)的频繁访问限制了变换器模型的性能和上下文长度。为了优化自注意力,提出了大量工作。我们关注精确的注意力优化,而像线性注意力这样的有损优化不在讨论范围内。内存高效注意力主要是为了缓解大内存成本。Rabe 等人 [287] 证明,自注意力需要 O(logn) 的内存复杂度而不是 O(n^2)。通过使用延迟 softmax,softmax 中的除法∑j exp(sj) 可以推迟到注意力操作的最后来进行。因此,求和可以递增地处理,只需要一个标量(即 O(1))来维护中间结果,而不会改变输出。自注意力需要额外的 O(logn) 内存复杂度来保持查询列表中的附加索引,以顺序计算所有查询的结果。FlashAttention 系列进一步展示了快速和内存高效的精确注意力与 IO 感知性、高并行性和平衡的工作负载在 GPU 上。在 FlashAttention [115] 中,提出了一种 IO 感知平铺算法,以减少基于在线 softmax 的慢 HBM 和快片上 SRAM 之间的内存读写次数。更具体地说,softmax 可以通过跟踪包括最大分数和指数分数总和在内的归一化统计数据,逐块计算。因此,平铺算法将自注意力中的所有计算操作链,包括矩阵乘法、softmax、矩阵乘法等,融合在一个 cuda 内核中,以减少 HBM 访问。FlashAttention2 [116] 通过在序列长度维度上增加并行性和改进线程块内数据共享的 warp 级调度,进一步提高了 FlashAttention 中的低占用率和不必要的共享内存读写。此外,流行的训练系统 [174] 通常采用 FlashAttention-2 进行高性能。FlashAttention-3 [262] 通过挖掘新推出的硬件功能来加速 H100 GPU 上的注意力,因为以前的 FlashAttention 实现基于 A100 GPU。基于 FlashAttention-2,重新设计了一种交错块状 GEMM 和 softmax 算法,以隐藏 softmax 中的非 GEMM 操作与异步 WGMMA 指令的 GEMM。此外,通过利用 Tensor Cores 和张量内存加速器(TMA)的异步性,通过 warp 专门的软件流水线方案,使整体计算与数据移动重叠。Blockwise Parallel Transformer (BPT) [263] 通过扩展 FlashAttention 中的平铺算法来进一步减少大量内存需求,将前馈网络融合起来。还针对各种架构优化了注意力机制,利用硬件特定特性。例如,SWattention [264] 为 Sunway 新架构设计,建立在 FlashAttention 之上。同样,Bikshand 等人 [265] 在 H100 GPU 上实现 FlashAttention-2,使用 Cutlass 库。他们利用 TMA 和 WarpGroup Matrix-Multiply-Accumulate (WGMMA) 指令分别优化数据复制和 GEMM 操作。此外,基于 Cutlass 库,精心设计了基于张量布局变换和两个 GEMM 之间的数据复制与计算之间的软件流水线。注意力机制还针对分布式 LLM 训练中常见的变长序列进行了优化。这些变长序列如果填充到最大长度,可能会产生显著的内存和计算成本。FlashAttention-2 通过不可分割地并行化序列长度维度来高效处理变长输入。ByteTransformer [266] 专注于无填充变换器,针对变长输入,维护一个位置数组,在计算期间记录原始张量和中间打包张量之间有效 token 的映射关系。这种长序列的融合多头注意力算法采用针对未填充张量的优化分组 GEMM。这种优化减少了填充引起的内存和计算开销,从而提高了性能。
DNN 编译器在优化 LLM 训练中的关键计算中发挥着重要作用。自动生成高效的操作员内核,大大减轻了不同硬件供应商对库基础内核优化的负担。在训练过程中,通过自动分析计算图执行操作员融合。
高效的操作员内核生成。 Halide [267] 和 TVM [252] 自动生成高性能的操作员实现,依赖于多个有效的调度原语,利用各种后端上的并行性和数据局部性。此外,Roller [268] 优化了在内核实现的大搜索空间中寻找最优替代方案的成本。它主要生成一个由 Load、Store 和 Compute 接口组成的 tile 内核,然后通过扩展然后扩展的方法构建完整的操作员内核。Triton [269] 提供了一种基于 C 的语言和编译器,便于表达和优化 tile 张量程序,以获得竞争性的性能。特别是,通过机器依赖的编译通道支持有效的优化,如层次化平铺和共享内存分配。ALCOP [270] 对操作员在 GPU 上执行自动负载 - 计算流水线,以重叠高延迟内存访问与计算。通过流水线缓冲区检测以及复杂的索引分析和替换,在复杂的循环结构中使用多阶段流水线。
图级别操作员融合优化。 随着计算核心速度和内存带宽差异的扩大,现代 DNN 受到内存访问的限制。通过编译器执行的操作员融合,挖掘了操作员间的数据重用。许多编译器工作 [288]–[291] 通过设置专家规则执行操作员融合。特别是,Chimera [271] 致力于优化计算密集型操作员链。首先将操作员链分解为一系列计算块,然后根据分析模型选择最优的块执行顺序,以最大化数据重用。此外,设计了可替换的微内核以利用硬件特定的块内优化。Welder [272] 将计算图降低到平铺级别的数据流图,其中节点是操作员平铺,边用连接节点的张量数据重用的内存级别标记。在平铺级别搜索最大化不同内存层次结构中数据重用的操作员融合组合。
Pytorch2 [273] 展示了两个扩展,即 Python 级 JIT 编译器 TorchDynamo 和相应的编译器后端 TorchInductor,以在各种后端上实现更强大的图编译,显著提高性能,同时不牺牲 Python 的灵活性。Slapo [203] 提出了一种调度语言,将模型执行与定义解耦。声明一组调度原语,用户可以转换模型以获得高性能内核。JIT-Q [274] 提出了一种针对权重的即时量化方法,该方法在训练期间只存储高精度副本的权重,并根据商业 PIM(内存内处理)解决方案的内存内 ALU 增强创建基于低精度权重副本。
低精度训练是降低大规模模型训练中的计算、存储和通信成本的有效方法。如今,LLM 训练通常利用 FP16 和 BF16 数据类型。特别是,BF16 可以表示与 FP32 相同的值范围。由于在 FP16 中损失标量变得太低时损失会缓慢发散 [293],BF16 训练在 BLOOM [292] 等模型中得到应用。然而,快速的 bfloat16 支持仅在 TPU 或在 NVIDIA Ampere 系列或之后开发的 GPU 上可用。此外,混合精度训练和技术,如损失缩放,被利用以确保由于表示的动态范围有限而减少精度时的数值稳定性。8 位甚至更低位训练也成为量化研究的焦点。
流行的训练系统通常采用 FP16/BF16 混合精度策略,在训练期间降低精度,如 Megatron-LM [20] 和 ColossalAI [253] 的工作所强调。FP16 混合精度训练方案 [275] 利用 IEEE 半精度格式存储权重、激活和前向和后向算术操作的梯度。为了在降低精度的同时保持模型精度,每个优化器步骤都会保留权重的单精度副本进行累积。还应用了损失缩放,以保留小幅度梯度的值。Campo [276] 通过自动图重写优化了 FP32 和 FP16 之间转换引起的铸造成本。这很关键,因为铸造成本有时可能会抵消低精度的性能优势。Campo 还使用离线训练的线性回归模型预测 FP32 和 FP16 操作的铸造成本和执行时间。BF16 [277] 也在各个领域 [294], [295] 的混合精度训练中得到广泛应用。它具有与 FP32 相同的表示范围,并且不需要为收敛进行超参数调整。此外,THC [278] 通过消除压缩和解压缩的需求,解决了参数服务器架构中的计算开销,从而提高了效率。
随着以 FP8 等低精度数据类型为特征的新芯片的推出,混合精度训练旨在使用更低精度进行训练。新设计的数据格式结合了确保数值稳定性的技术,主要利用 FP8 训练深度学习神经网络。Wang 等人 [279] 使用新的 FP8 浮点格式进行数据的数值表示和计算。块计算和随机舍入分别在浮点累加和权重更新过程中使用,以保持模型精度。Sun 等人 [280] 提出了一种适用于整个深度学习模型的混合 8 位浮点训练,而不会降低精度。新颖的混合 FP8 格式为前向和后向传播使用不同的指数位和尾数位,因为前向和后向传递在范围和精度之间的最佳平衡不同。此外,使用如损失缩放等技术以避免精度降低。随着更多 FP8 数据类型的加速器的成熟,基于 NVIDIA H100 GPU [296] 的 FP8 自动混合精度框架(FP8-LM)[281] 被提出,逐渐融合了 8 位梯度、优化器状态和分布式并行训练,并且指定了包括张量、流水线和序列并行在内的 FP8 低位并行性。此外,设计了精度解耦和自动缩放,以解决由于动态范围较窄和精度降低导致数据下溢或上溢的问题。FlashAttention-3 还采用了块 GEMM 量化和非相干处理,利用硬件对 H100 GPU 上 FP8 低精度的支持。此外,Rouhani 等人 [282] 利用与细粒度子块相关联的缩放因子的微缩数据格式,在权重、激活和梯度的 1 位精度下训练 LLMs,以最小的精度损失。
LLM 训练中的低位固定点训练也被研究。Jetfire [283] 保持 INT8 数据流,其中输入和输出以 INT8 数据格式加载和存储,以加速计算约束的线性操作员和内存约束的非线性操作员。此外,利用平铺算法挖掘共享内存数据访问,采用每个块的量化方法,其中执行更高精度的计算,即在线性操作员的 WMMA 张量核心操作中使用 INT32 和非线性操作中的 FP32,以保持预训练变换器的精度。Xi 等人 [284] 提出了一种针对变换器模型的新型 INT4 训练算法。在前向传播中,首先将激活矩阵转换为块对角哈达玛矩阵,以减轻激活中的异常值引起的精度降低,然后对转换后的矩阵进行量化。在反向传播中,利用位分割和利用分数采样根据激活梯度的结构稀疏性选择信息性梯度进行量化。最近,LLMs 的低精度训练已经发展到使用 1 位精度。BitNet [285] 在变换器块内使用 1 位权重和 8 位激活的新型低位精度矩阵乘法,将模型权重集中在零点周围,以在有限的数值范围内最大化容量,然后使用符号函数将其二值化为 +1 或 -1。为了确保训练稳定性和精度,梯度、优化器状态和高精度潜在权重副本被维护以进行参数更新。在 BitNet 的基础上,BitNet b1.58 [286] 通过将模型权重降低到三元值 {-1, 0, 1} 来增强建模能力。权重矩阵通过其平均绝对值进行缩放,每个值四舍五入到 -1、0 和 +1 中最近的整数。
LLM 训练期间的内存消耗可以分为四个主要部分:模型状态、激活、临时缓冲区和内存碎片化。
模型状态 :模型状态包括优化器状态、梯度和模型参数所占用的内存。在混合精度训练 [275] 中,模型参数和激活以 16 位精度存储。当训练具有Φ参数的模型时,需要 4Φ字节来存储参数和梯度。32 位参数、动量和方差的副本各需要 4Φ字节,总计 12Φ字节。因此,存储模型状态所需的总内存为 16Φ字节。
激活 :激活指的是在前向传播期间生成的张量。这些张量在后向传播阶段用于梯度计算。
临时缓冲区 :临时缓冲区用于存储中间结果。例如,像梯度 AllReduce 这样的操作通常将梯度在单个扁平缓冲区中融合,以提高吞吐量。
内存碎片化 :内存碎片化可能导致尽管有大量可用内存,但内存请求失败,因为可用内存变得碎片化,没有足够的连续内存来满足内存请求 [145]。
为了解决 LLM 训练的内存限制,提出了各种内存高效技术。这些包括激活重计算策略,通过增加计算换取减少内存使用;冗余减少方法,通过减少训练过程中的数据复制来最小化内存冗余;碎片化整理技术,通过优化内存分配和释放来减少碎片化,提高内存利用率;以及交换和卸载方法,利用 CPU 内存和 NVMe SSDs 来补充 GPU 内存。图 12 概述了这些优化在内存高效 LLM 训练中的分类。
在模型训练的后向阶段,激活对于计算梯度至关重要。随着模型大小的增加,训练期间存储这些激活所需的内存可能超过 GPU 内存容量,从而限制了可以训练的模型规模。激活重计算 [325] 通过策略性地在前向传播期间丢弃某些激活,并在后向传播期间按需重新计算它们,提供了一种解决方案。这种方法已成为减少 LLM 训练中内存消耗的标准方法。有效激活重计算的关键是平衡内存节省与额外计算开销。我们将这些方法分为两种主要方法:静态逐出和动态逐出。
静态逐出涉及建立一个固定的计划,在前向传播期间丢弃激活,并在后向传播期间重新计算它们。Checkmate [300] 将这个激活重计算问题表述为一个混合整数线性规划问题,以确定静态深度学习模型的最佳重新材料化计划。然而,Checkmate 在扩展到像 LLMs 这样的大型模型时遇到困难,因为搜索空间巨大。最近,一些工作提出了针对 LLM 训练量身定制的激活重计算策略。Selective-checkpointing [170] 选择性地丢弃内存密集型注意力模块的激活。FlashAttention [115] 将注意力模块融合到一个单独的内核中,并且也采用选择性检查点来减少内存消耗。DistFlashAttn [176] 解决了长序列中由于重新计算注意力模块而导致的高计算开销,采用了重新材料化感知的梯度检查点策略。具体来说,DistFlashAttn 在 FlashAttention 内核的输出处而不是 Transformer 层边界处放置检查点,从而在后向传播期间消除了注意力模块的重新计算,只要求存储其输出。LoongTrain [174] 引入了 selective-checkpoint++,在长序列训练中通过将注意力模块添加到白名单中来进一步优化检查点过程。这种方法在前向传播期间保存白名单模块的输出,在后向传播期间检索这些存储的输出,而不是重新计算它们,继续计算图,从而减少了重新计算注意力的需要。与最近主要关注为 LLM 训练的注意力模块手工检查点策略的工作不同,Yuan 等人 [301] 仔细测量了在模型训练期间重建每个激活张量所需的最小计算成本。他们通过枚举所有可能的检查点方法,推导出内存和计算成本的 Pareto 前沿。从这个 Pareto 前沿中,他们选择了一个最优平衡计算和内存成本的解决方案。
动态逐出根据训练过程的当前状态实时决定哪些激活要丢弃和重新计算。DTR [297] 提出了一种贪婪在线算法,用于在运行时对静态和动态模型的张量进行启发式逐出和重新材料化。MegTaiChi [298] 引入了一种动态张量逐出,利用在运行时跟踪的张量访问模式。Coop [299] 提出了一种方法来缓解由于激活重计算方法在丢弃张量时不考虑它们的连续性而引起的内存碎片化问题。Coop 采用高效的滑动窗口算法,确保只有连续的内存块被逐出,从而最小化内存碎片化。
传统的数据并行方法在所有 GPU 上复制整个模型状态,导致大量的冗余内存使用。冗余减少技术被提出来通过消除或减少每个设备上的内存冗余来优化内存使用。这些技术通常寻求在内存效率和引起的通信开销之间取得平衡,从而以可接受的成本促进更大规模或批量大小的训练。
零冗余优化器(ZeRO)[145] 通过在所有 GPU 上完全分片模型状态来优化内存冗余,通过三个阶段:ZeRO-1、ZeRO-2 和 ZeRO-3。ZeRO-1 在所有 GPU 上全局分布优化器状态。在训练过程中,每个 GPU 进行独立的前向和后向传播来计算梯度,然后使用 ReduceScatter 操作在数据并行组内的所有 GPU 之间同步梯度。每个 GPU 负责更新特定部分的模型参数。随后,更新后的模型参数片段从其他 GPU 上收集,使用 AllGather 操作,确保所有 GPU 都有最新的模型参数。ZeRO-1 将优化器状态的内存消耗从 12Φ降低到 12Φ/N,其中 N 是数据并行的大小。在 ZeRO-1 的基础上,ZeRO-2 进一步在所有 GPU 上分片梯度,每个 GPU 只更新其参数片段,将用于存储梯度的内存需求从 2Φ降低到 2Φ/N。ZeRO-3 除了优化器状态和梯度外,还分片参数。每个 GPU 只持有部分参数。当需要远程 GPU 上的参数进行即将到来的计算时,通过 AllGather 操作收集它们,并在之后丢弃。在 ZeRO-3 中,每个 GPU 只持有与其特定参数分区对应的权重、梯度和优化器状态,将整体内存消耗从 16Φ降低到 16Φ/N。ZeRO 被许多框架广泛采用,如 DeepSpeed [183]、PyTorch-FSDP [146] 和 ColossalAI [253]。
ZeRO 面临着通信挑战,因为集体通信操作的延迟随着通信规模的增加而增加。在分布式 LLM 训练中,内存利用率和通信成本之间存在权衡。通过在较小的 GPU 组中分片模型状态,可以优化通信开销,这些小组是大型 GPU 集群中的较小 GPU 集合。这种方法减少了节点间通信和通信规模,尽管由于模型状态的冗余增加,可能会导致更高的内存使用。关键是要在通信规模和内存利用之间找到平衡 [305]。已经提出了几种在 ZeRO 框架上构建的方法,以解决通信效率低下的问题,同时提高内存利用率。ZeRO++ [302] 按照 ZeRO-3 在全球所有设备上分片所有模型状态,并进一步在 GPU 的子组内引入参数的二次分片。在前向阶段,它利用主要分片在所有 GPU 上收集参数,并在子组内维护参数的二次分片,通常在同一节点内。在后向阶段,它从这个二次分片中收集参数,减少通信规模和节点间通信。此外,ZeRO++ 使用量化来压缩参数和梯度,有效地减少通信量,但以精度为代价。MiCS [147] 和 FSDP [146] 在子组内分片所有模型状态组件,并在子组之间复制它们,从而减少通信规模和通信延迟,提高训练性能。AMSP [305] 和 PaRO [303] 结合了三种灵活的分片策略,包括全副本、全分片和部分分片,允许模型状态的每个组件独立选择分片策略。AMSP 制定了一个优化问题,以在内存约束下找到最小化通信成本的最佳分片策略。此外,AMSP 提出了一种定制的通信和计算重叠策略,结合这些灵活的分片策略,实现优化的训练效率。RTP(Rotated Tensor Parallelism)[304] 通过策略性地分片激活和旋转权重/梯度来最小化内存冗余。
GPU 内存碎片化指的是在相邻张量之间出现的分散、无法使用的内存块。由于不同张量的生命周期不同,以及一般深度学习框架(如 PyTorch [240] 和 TensorFlow [326])的内存分配和释放方案效率低下,这个问题在 LLM 训练中尤为突出。此外,像重计算和卸载这样的内存优化技术通过引入更频繁和不规则的内存分配和释放请求,加剧了这个问题 [299], [306], [308]。碎片化问题可能导致高峰值内存和内存不足(OOM)错误,限制了批量大小和整体训练效率。碎片整理工作提出了通过内存管理技术来缓解这些问题。
深度学习框架通常使用带有内存池的缓存分配器,以实现快速内存分配和释放,而不需要设备同步。已经提出了几种基于缓存分配器中张量分配和释放方案来减少内存碎片化的方法。ROAM [306] 通过考虑张量的生命周期和大小,共同优化操作员的执行顺序和张量分配。它引入了一个高效的基于树的算法,搜索最大化张量重用和减少数据碎片化的执行计划。ROAM 已在单 GPU 场景中进行了评估,特别是使用最大的模型是 1.5B GPT-2 XL [5],但尚未在具有更大模型的分布式训练场景中进行测试,其中计算图可能变得显著更大。Imanishi 等人 [307] 提出了一种离线优化方法,将张量分配建模为二维装箱问题。在这个模型中,每个张量分配表示为一个垂直可移动的矩形,反映了模型训练期间的周期性分配模式。他们提出了一种使用模拟退火的启发式算法,优化分配的拓扑排序,旨在最小化碎片化。虽然有效,但这种方法在应用于 LLMs 时可能会遇到可扩展性问题,因为涉及的分配数量众多,模式复杂。MegTaiChi [298] 和 Coop [299] 在减少内存消耗时考虑内存碎片化问题。
GMLake [308] 和 PyTorch 可扩展段 [309] 提出利用低级 CUDA 驱动应用程序编程接口的虚拟内存管理(VMM)功能来缓解碎片化。这个低级 API 为开发人员提供了直接控制 GPU 虚拟内存操作的功能,如保留、映射和管理虚拟内存地址。基于此,GMLake [308] 引入了一种虚拟内存拼接机制,通过虚拟内存地址映射将非连续的内存块合并为更大的块,最小化数据移动和复制。同样,PyTorch 的可扩展段 [309] 使分配的内存段能够扩展到更大的尺寸以供重用。这两种方法对不同模型和内存高效训练技术都是透明的,并且可以无缝集成到现有的深度学习框架中。此外,GMLake 证明了在多 GPU 上具有出色的可扩展性,开销最小,不需要修改用户代码。PyTorch-v2.1 也已经集成了可扩展段。
为了在较少的 GPU 上有效训练 LLMs,提出了各种利用交换和卸载方法的工作。这些技术将部分计算和数据从 GPU 内存转移到外部资源,这些资源虽然价格低廉且速度较慢,但拥有巨大的容量。
许多研究提出了有效利用 CPU 内存以增强分布式 LLM 训练的方法。这些技术大致可以分为两种主要方法:静态卸载和动态卸载。
静态卸载。 静态卸载方法涉及预先确定的模型组件在 GPU 和 CPU 内存之间的分配。L2L [310] 逐层管理和移动张量。L2L 同步获取即将到来的计算层所需的张量到 GPU 内存,同时保持剩余层的张量存储在 CPU 内存中。L2L 允许将模型扩展到任意深度,但无法在多 GPU 上扩展。相比之下,ZeRO-Offload [311] 专注于多 GPU 训练。它将模型参数保持在 GPU 上,并将优化器状态和梯度存储在 CPU 内存中。此外,它将优化器更新计算卸载到 CPU。这种方法使得使用 16 个 V100s 训练高达 700B 的模型成为可能。然而,ZeRO-Offload 可能会留下一些未使用的 GPU 内存,并遭受缓慢的 CPU 优化器更新 [312]。
为了解决这个问题,Elixir [312] 使用搜索引擎通过利用预运行时模型分析来找到内存分区和卸载的最佳组合。与 ZeRO-Offload 不同,Elixir 通过在 GPU 和 CPU 之间分割模型状态和优化器块,有效利用了所有可用的 GPU 内存。Mobius [315] 通过引入流水线并行方案来解决在有限的 GPU 间带宽和高通信争用下的商品服务器上的多 GPU 训练问题。这个方案为每个 GPU 分配多个阶段,并在 GPU 和 CPU 内存之间动态交换它们。此外,Mobius 通过预取和交叉映射优化通信,以减少开销和争用。Yuan 等人 [301] 提出通过在流水线阶段的粒度下卸载和重新加载激活来减轻激活瓶颈,同时最大化激活传输与计算之间的重叠,从而避免放慢训练过程。与其他卸载工作相比,这项工作更侧重于改善计算和内存利用率之间的平衡,而不是在极其紧张的内存预算下进行训练。
动态卸载方法根据实时优化内存利用率和数据传输,在 GPU 和 CPU 内存之间适应性地分配模型或张量的分区。STRONGHOLD [318] 提出在 GPU 和 CPU 内存之间动态卸载模型状态,并保持适当的工作窗口大小,以最小化卸载期间的 GPU 停顿。Harmony [316] 采用基于启发式的调度器将计算和模型状态映射到物理设备。Harmony 通过减少交换和快速的点对点交换来减少卸载开销。TMOF [317] 引入了不相交交换和双向重叠协调机制,以防止交换和卸载期间的 PCIe 通道争用。对于 MoE 模型,MPipeMoE [327] 设计了一种自适应和内存高效的流水线并行算法。具体来说,MPipeMoE 通过消除内存冗余和自适应选择组件来决定是卸载还是重新计算所需张量,以减少内存需求。
为了促进更好的内存管理,一些研究提出了将张量分解为更细粒度单元的系统。TSPLIT [313] 和 PatrickStar [314] 是两个动态内存管理系统,它们优化了 GPU 的峰值内存使用。TSPLIT 将张量分解为微张量,并在微张量级别执行操作,实现精确和动态的内存操作。PatrickStar 将模型数据组织成内存块,这些块在 CPU 和 GPU 内存之间动态分配,并优化 CPU-GPU 数据传输以及带宽利用。此外,TSPLIT 使用模型引导的计划算法为每个张量找到最优的内存配置,而 PatrickStar 采用运行时内存跟踪、块逐出策略和设备感知操作员放置,进一步最小化 CPU 和 GPU 之间的数据移动。
为了促进万亿规模 LLMs 的训练,仅依靠 CPU 卸载的方法是不足够的,因此提出了在训练期间将数据卸载到 CPU 内存和 NVMe SSDs 的几种方法。ZeROInfinity [319] 将所有分片模型状态卸载到 CPU 或 NVMe 内存,并将激活卸载到 CPU 内存。这种方法支持在 32 个节点(总共 512 个 V100s)上训练高达 32T 参数的模型。然而,CPU 卸载对激活仍然需要大量的 CPU 内存。例如,存储 10T 模型的训练需要大约 0.76TB 的 CPU 内存,而对于 100T 模型则需要大约 4TB。Fuyou [322] 专注于在 CPU 内存容量有限且仅有单个 GPU 的商品服务器上训练 LLMs。与 ZeRO-Infinity 相比,Fuyou 进一步将激活卸载到 SSDs,并结合 SSD-CPU 通信作为额外的优化维度。它还引入了同步的外核 CPU 优化器,与后向传播阶段重叠,并引入了自动激活交换机制,从而最大化 GPU 利用率。Smart-Infinity [321] 提出使用近存储处理设备减少二级存储带宽需求,用于参数更新。MoESys [323], [324] 结合了各种存储设备(GPU、CPU 内存和 SSDs),以保存稀疏参数状态和密集参数状态,并提出了一种 2D 预取调度策略,以便 MoE 训练的参数计算可以与调度重叠。
不同的并行机制引入了不同的网络通信流量模式。例如,张量并行需要在张量并行等级之间进行 AllReduce 操作。另一方面,数据并行在每个迭代结束时需要在数据并行等级之间进行 AllReduce 操作以同步梯度。流水线并行涉及在每个阶段结束时将激活值传递到下一个阶段。通常,训练框架将张量或序列并行通信组(这些组需要高带宽)放置在高带宽域内(例如,同一节点内),而将数据并行或流水线并行通信组(这些组的带宽需求较低)放置在高带宽域之间。图 13 显示了 LLM 训练实践中的通信热图,并很好地反映了不同并行策略带来的数据流量。从这个热图中可以观察到,LLM 训练通信流量呈现出明显的模式和层次结构,大部分通信发生在较小的范围,只有一小部分流量穿过整个集群。这一洞见激发了像铁路优化拓扑 [62] 这样的方法,它减少了不必要的核心交换机以降低成本。本节介绍分布式 LLM 训练的集体通信性能优化的系统和技术。如图 14 所示,我们首先讨论集体通信库,它们利用预定义和合成算法。接下来,我们探索旨在重新组织通信操作以与计算重叠,从而减少延迟并加速训练过程的通信调度技术。最后,我们深入研究了网络内聚合(INA),它利用网络设备的计算能力来执行聚合操作,例如求和深度学习模型的梯度。
有效压缩模型参数和梯度可以显著减少分布式 LLM 训练期间的通信开销。各种研究探索了稀疏通信和量化方法。例如,ZeRO++ [302] 在通信之前将权重量化为 FP16 到 INT8 数据类型,以缩小每个模型参数。然而,这些工作通常涉及有损稀疏化或量化技术。我们没有调查有损数据压缩技术,因为这超出了这项工作的范围。
消息传递接口(MPI)是大规模科学应用在并行计算架构上的一种广泛采用的编程模型。MPI 有几种实现,包括 OpenMPI [328]、MPICH2 [329] 和 MVAPICH [330]。这些库提供了多种 CUDA 感知原语,如 AllReduce、AllGather 和 ReduceScatter,这些对于分布式 LLM 训练至关重要。在实践中,当前的训练框架更倾向于使用针对特定 AI 加速器的集体通信,这些通信具有预定义或合成算法。
NVIDIA 的 NCCL [331] 和 AMD 的 RCCL [332] 是高度优化的库,通常在它们各自的 AI 加速器上比基于 MPI 的集体通信库表现更好。这些库通常根据网络拓扑和输入张量大小等条件选择预定义算法来执行集体操作。
环形算法。 环形算法用于像 AllReduce 这样的集体通信,以在所有 GPU 之间移动数据。在这个算法中,输入张量被分成多个块,并且在操作过程中逐个传输。这个流水线减少了每个设备等待数据的空闲时间。百度使用了带宽最优的环形 AllReduce 算法 [333] 进行分布式深度学习模型训练。Horovod [143] 用 NCCL 替换了百度的环形 AllReduce 实现,并为分布式训练设计了用户友好的接口。
树算法。 环形算法的延迟随着 GPU 设备数量的增加而增加 [346]。双二叉树算法 [334] 被提出来解决这个问题。双二叉树依赖于这样一个事实:在二叉树中,一半或更少的等级是节点,一半或更多的等级是叶子。因此,可以使用叶子作为节点,反之亦然,为每个二叉树构建第二个树。这个算法在基于 MPI 的库、NCCL 和 RCCL 中实现。
混合算法。 几种方法提出了使用混合算法来处理具有异构节点内和节点间通信带宽的训练集群上的集体通信任务。两级 AllReduce [335] 将单个 AllReduce 操作分为三个步骤:节点内 Reduce 使用 PCIe/NVLINK,节点间 AllReduce 使用网络,节点内 Broadcast。2D-Torus AllReduce [336] 和 ACCL [337] 将单个 AllReduce 操作分解为三个阶段:节点内基于环形的 ReduceScatter,节点间基于树的 AllReduce,节点内基于环形的 AllGather。BlueConnect [338] 将单个 AllReduce 操作分解为多个可并行化的操作,每个操作可以映射到不同的网络结构,利用每种特定结构的最佳预定义实现。Plink [339] 可以探测网络拓扑,并有效地生成两级混合通信计划,利用数据中心网络中的局部性。
一些方法已经出现,它们合成了针对硬件拓扑的集体通信算法和内核,旨在在许多情况下胜过通用预定义算法。
GC3 [340] 引入了一个面向数据的特定领域语言(DSL),用于设计自定义的集体通信算法。它包括一个优化编译器,将这些算法转换为针对特定硬件配置优化的可执行形式。SCCL [341] 将集体通信合成问题编码为可满足性模理论(SMT)公式。这种方法旨在推导出精确的时间表,用于 Pareto 最优算法,优化延迟和带宽利用率。TACCL [342] 将寻找最优通信算法的问题表述为混合整数线性规划(MILP)问题。它利用通信草图抽象来有效收集必要信息并减小搜索空间,目标是最小化总体执行时间。Blink [343] 通过在运行时探测每个作业的可用链接集来动态构建具有适当链路容量的拓扑。使用这个拓扑,它通过创建数据包生成树和生成 CUDA 代码来优化通信速率。P2 [344] 使用并行矩阵在系统级别划分并行轴,从而生成拓扑感知的并行放置和简化策略。通过模拟和预测通信成本,这种方法减少了实际评估的数量。
分布式训练中的通信调度重新组织通信操作,使其与计算重叠,从而减少延迟并加速训练过程。通信调度的关键概念涉及根据并行训练的数据依赖性重新排序通信操作。混合并行 LLM 训练需要多维通信调度方案来管理由数据、流水线、张量和序列并行生成的通信,以及它们的组合。
在后向阶段,而不是等待所有梯度计算完成后再启动通信,通信可以在每个梯度准备好后立即开始。这种无等待反向传播方法利用依赖于有向无环图来有效管理任务。Poseidon [345] 使用先进先出(FIFO)队列来调度 AllReduce 操作员,确保每层在生成梯度后立即开始其通信。受到大型张量上集体通信效率的启发,GradientFlow [346] 和 Pytorch DDP [142] 将多个顺序 AllReduce 通信操作合并为单个操作。这种方法通过等待短暂的时间,然后将多个梯度合并为一个 AllReduce 操作,避免了在网络上传输大量小张量。
FIFO 方案通常是次优的,因为在后向阶段生成的通信序列与前向阶段的计算序列不同。这种不匹配可能导致通信阻塞计算,即使启用了重叠。因此,许多方法采用优先级队列来有效调度通信操作员。P3 [347] 以更细粒度调度 AllReduce 操作员,将当前层的梯度通信与下一层的前向计算重叠。与基于 FIFO 队列的调度不同,这种方法将层划分为固定大小的切片,并根据它们在前向传播中处理的顺序优先同步切片。因此,第一层获得最高优先级,优先级随着层的结束而递减。当使用参数服务器架构进行分布式模型训练时,TicTac [348] 优先考虑加速底层计算图中的关键路径的传输。ByteScheduler [349] 和 PACE [350] 被提出以在训练框架中推广基于优先级的通信调度。具体来说,ByteScheduler [349] 引入了一个统一的抽象,以在不破坏框架引擎内原始依赖关系的情况下促进通信调度。ByteScheduler 通过使用贝叶斯优化自动调整两个关键参数:分区大小和信用大小,实现了良好的性能。PACE [350] 通过将原始 AllReduce 操作分割成更小的部分来实现抢占式通信。被抢占的 AllReduce 操作员可以在以后恢复。这种抢占策略防止了大通信张量的队首阻塞。此外,PACE 使用动态规划方法将小通信张量融合,以减少处理大量小张量的开销,从而实现更有效的带宽利用。为了提高 MoE 系统中的带宽效率,Lina [188] 优先考虑 All-to-All 操作而不是 AllReduce。通常,专家并行(All-to-All)和数据并行(AllReduce)过程使用单独的 CUDA 流,可能导致没有协调的潜在重叠和带宽共享。Lina 将张量分解为更小的块,确保 All-to-All 操作获得全部带宽,同时允许 AllReduce 微操作在空闲时间运行。此外,微操作使 All-to-All 操作与专家计算重叠。
一些进步集中在将通信和计算操作分解为细粒度任务,以更大的灵活性重新排序这些操作,以最大化重叠并优化执行效率。
流水线阶段分解。 当使用传统的流水线并行时,每个 GPU 存储一层的连续段。Breadth-First [159] 进一步将这些连续阶段分解为分布在不同 GPU 上的更细粒度的阶段,通过将第一个和最后一个 GPU 连接起来形成一个循环,使每个 GPU 被分配多个阶段。这允许给定的微批次更早地到达流水线的末端,减少流水线泡沫。Breadth-First 使用广度优先调度策略实现更大的计算 - 通信重叠。Fold3D [351] 采用全进全出的调度策略,使流水线的梯度同步过程与计算重叠。这涉及进一步折叠流水线内的模型片段,其中每个设备包含两个模型片段,允许一个片段的梯度同步与另一个片段的前向或后向计算重叠。异步流水线并行放宽了梯度和参数更新之间的数据依赖性。利用这一特性,TriRace [352] 推迟参数更新以最大化计算与梯度通信的重叠。此外,TriRace 将流水线阶段之间的双向 P2P 通信分解为两个单独的单向操作,并根据关键路径分析对它们进行优先级排序。
通信分解。 通信原语可以被分解为具有高调度灵活性的细粒度操作。Wang 等人 [353] 将通信操作(例如 AllGather 和 ReduceScatter)分解为一系列细粒度的点对点集合。此外,计算操作(例如爱因斯坦求和)被划分为细粒度任务,每个任务执行部分计算。这种分解创造了更多的机会,使通信与计算重叠。SYNDICATE [354] 将通信操作分割成更小的子操作,称为 Motifs,并使用中央优化器使用马尔可夫链蒙特卡洛搜索来实现最优的重叠执行计划。Centauri [355] 采用不同的方法,使用原语分区、组分区和工作负载分区将通信操作分解为细粒度的原子操作。然后使用工作负载感知调度、反向调度和弹性调度来调度这些操作,以最大化重叠效率。DeAR [356] 还分解通信原语,特别是将 AllReduce 分解为 AllGather 和 ReduceScatter。这种分解允许后续操作与模型的前向传播过程重叠,从而消除了等待两个通信步骤完成的需要。
计算分解。 当使用张量并行时,在前向阶段需要一个 AllReduce 通信来同步矩阵乘法输出。CoCoNet [357] 通过将输出划分为更小的块,并在计算每个结果块后立即启动 AllReduce 内核,来促进矩阵乘法和 AllReduce 的重叠。为了最小化等待 AllReduce 内核的时间,数据块以精心安排的顺序输入到矩阵乘法内核。T3 [358] 应用了一种硬件 - 软件共同设计方法,它透明地重叠矩阵乘法与通信,同时最小化资源争用。在硬件层面上,T3 引入了一个跟踪和触发机制来协调生产者的计算和通信活动。此外,它采用了计算增强内存来处理通信过程中所需的相关计算操作。反向传播生成两种类型的梯度:输出梯度,用于计算前一层的梯度,以及权重梯度,用于更新层的权重参数。这些权重梯度需要使用 AllReduce 与其他等级同步。传统的框架同时执行权重和输出的梯度计算。Out-of-order backpropagation (ooo-backprop) [361] 将权重和输出的梯度计算解耦,灵活地安排权重梯度计算的顺序。这允许更关键的计算被优先考虑和相应地调度。因此,ooo-backprop 通过基于这种乱序计算策略的通信调度来优化整体性能。Zero Bubble [156] 也使用这种方案来降低流水线并行的泡沫率。使用激活检查点时,训练框架需要在后向阶段重新计算激活。这种重新计算也涉及在张量并行时的 AllReduce 通信。Oases [359] 通过始终将 AllReduce 通信作为重新计算单元的最后一个前向通信操作来减少冗余通信,并进一步将批量划分为更小的子批量,允许两个批量的通信和计算重叠。Lynx [360] 也利用了重新计算和通信重叠的潜力,使用两种重新计算调度算法 OPT 和 HEU 来寻找最优或近似最优的重新计算调度策略,实现最佳的重叠和训练性能。
网络内聚合(INA)使用网络设备的计算能力来执行聚合操作,如求和深度学习模型的梯度。这种技术之前已被提出以加速大数据处理。值得注意的是,像 NetAgg [369]、SwitchAgg [370] 和 CamDoop [371] 这样的框架通过在直接连接拓扑内的高性能中间盒或服务器上执行数据聚合,展示了显著的性能优势。许多方法已被提出用于将网络内聚合应用于深度学习模型训练,旨在减少在后向阶段梯度的 AllReduce 操作期间节点间交换的数据。
许多基于以太网的网络内聚合系统依赖于可编程交换机,并且可以用于分布式 LLM 训练。SwitchML [362] 支持在分布式训练的后向阶段将集体通信操作卸载到可编程网络交换机。由于完整的模型更新可能超过交换机的存储容量,SwitchML 通过交换机流式传输聚合,每次处理有限数量的向量元素的聚合函数。SwitchML 有两个限制。首先,处理浮点运算时,SwitchML 不能直接对浮点张量执行集体通信(如 AllReduce)。相反,它使用类似块浮点的方法将浮点值转换为 32 位整数。其次,SwitchML 主要在 DPDK 上实现,虽然有 RDMA 功能实现,但很难与训练框架集成。为了更好地促进分布式模型训练,FPISA [363] 在可编程交换机上直接运行 P4 [373] 程序来实现浮点计算。因此,训练框架可以将 FP16 张量的集体通信操作卸载到交换机,而无需将它们转换为 32 位整数。NetReduce [364] 支持与 RoCE 兼容的网络内聚合,充分利用 RoCE 的拥塞控制和可靠性设计,无需在交换机中进行昂贵的网络协议处理堆栈。NetReduce 使用附加到以太网交换机的 FPGA 板进行了原型设计。AllReduce-Switch [365] 与 NetReduce 密切相关,并且与其网络协议兼容。它引入了一种针对网络内聚合任务量身定制的新型交换机架构,并使用 FPGA 硬件实现了原型。PANAMA [366] 和 ATP [367] 也通过它们为共享环境设计的网络内聚合框架为该领域做出了贡献。PANAMA 专注于通过管理多个并发训练作业之间的带宽分配来优化网络负载。它解决了传统拥塞控制可能无法充分支持同时训练操作的挑战。ATP 则支持多个并发租户同时运行多个作业,强调对共享环境中不同工作负载的支持。某些工作针对特定的训练工作负载进行了定制,使它们不适合分布式 LLM 训练。例如,Libra [374] 是为使用参数服务器架构的稀疏模型训练而设计的。它将频繁更新的参数聚合卸载到可编程交换机,而将不经常更新的参数留给服务器处理。这种方法有效地减少了服务器负载。另一方面,iSwitch [375] 是为强化学习训练任务中的参数聚合而设计的。尽管其基于 FPGA 的实现支持本地浮点运算,但它在显著较低的带宽下运行。此外,iSwitch 在聚合期间存储了整个梯度向量,这对于强化学习工作负载是可行的,但对于大规模模型,尤其是 LLMs 来说,并不具有很好的扩展性。
NVIDIA Mellanox 的可扩展层次聚合协议(SHARP)[368] 是一种专有的网络内聚合方案,可在某些 InfiniBand 交换机和 NVIDIA GPU 上使用。SHARP 建立在 InfiniBand 之上,利用链路层流量控制和无损失保证,并采用专用的片上浮点运算单元进行集体卸载。SHARPv1 在 InfiniBand EDR 交换机上推出,SHARPv2 在 InfiniBand HDR 交换机上增强,具有支持集体通信(例如 Barrier、Reduce、AllReduce 和 Broadcast)、整数和浮点运算(16/32/64 位)以及 GPUDirect RDMA 等功能。SHARPv2 还使用流式聚合在线路速率下进行大型向量缩减,与 NCCL 集成,并可轻松用于现有的训练框架。在最新的 InfiniBand NDR 交换机上启用的 SHARP 已经为分布式 LLM 训练做好了生产准备,并已部署在许多训练集群中。除了 InfiniBand,NVIDIA 的 NVSwitch-v3 [46] 也集成了 SHARP,以加速基于 GPU 的集群中的集体操作。
LLM 训练通常涉及从几周到几个月的延长训练周期,使用数万个 GPU 的集群。涉及的组件数量庞大,从底层基础设施到训练系统优化,需要强大的容错机制来确保训练过程的可靠性。这是因为系统中任何一点的错误都可能导致训练过程因同步训练而暂停。在本节中,我们首先介绍 LLM 训练的故障分析,然后研究快速故障检测和恢复的方法。
来自不同来源的实证证据强调了 LLM 训练中故障的频繁性。例如,在 384 个 GPU 的集群上,Bloom 的平均每周经历 1-2 次 GPU 故障 [292]。Meta 的 175B OPT 模型在 992 个 A100 GPU 上进行了全面的培训记录 [398],在两周内记录了 40 多次中断,归因于硬件、基础设施和其他外部因素。更多近期的研究进一步强调了这一问题。Acme [23] 在其使用超过 1,000 个 A100 GPU 的训练过程中平均报告每 1-2 天发生一次故障。字节跳动的 MegaScale 项目 [71],使用 12,288 个 Ampere GPU,在几周内经历了 100 多次故障。Meta 的 LLaMA3 在 16,384 个 H100 GPU 的集群上进行了 54 天的预训练,期间经历了 466 次作业中断 [9]。频繁的故障主要是由于这些系统的复杂性和规模,以及延长的训练周期。整个训练系统包括我们之前调查过的各个组件的庞大数组。此外,同步训练进一步加剧了这个问题,因为任何单个节点的错误都可能导致整个作业失败,使系统特别容易受到即使是孤立的硬件或软件故障的影响。即使在单个节点上观察到的低至 1.5% 的日故障率,如阿里巴巴的集群 [383] 所观察到的,当扩展到具有 1,000 个 GPU 的系统时,也会转化为惊人的 84.8% 的日故障率。然而,训练系统规模的扩大趋势仍在继续,强调了容错机制在维护系统可靠性方面的挑战。这些故障背后的原因多种多样,来自 LLM 训练系统的各个组件。根据 Acme [23],最严重的影响来自硬件故障,如 GPU(例如,CUDA-Error、ECC-Error)、NVLink 和网络系统(例如,NCCL-Timeout-Error、Connection-Error)。阿里巴巴 C4 [377] 也有类似的观察结果。C4 进一步观察到,大多数错误(约 82.5%)局限于特定节点甚至单个设备,尽管用户观察到的大多数错误是 NCCL 错误。LLaMA3 预训练 [9] 还报告说,78% 的故障是硬件问题。此外,最新一代 GPU(A100 和 H100)倾向于表现出较高的错误率,可能由于快速发展、匆忙交付和增加的功耗 [377], [399]。除了硬件,分布式训练框架、数据预处理管道或库依赖中的软件相关问题可能导致崩溃或意外行为 [23], [378], [399]。模型本身的复杂性质可以引入不稳定性,如损失尖峰、数值溢出或下溢、梯度爆炸或优化困难 [398], [400]。数据中心的外部因素,如停电或冷却系统故障,进一步导致系统不稳定。例如,集群服务器房间的高温也可能导致 GPU 过热,从而导致 NVLink-Error 或 ECC-Error [23] 或训练速度不稳定 [9]。这些频繁和多样化的 LLM 故障导致 GPU 的显著浪费。这种效率低下主要表现为两种方式:故障恢复和性能下降。首先,LLM 训练作业定期保存检查点,以维护进度。发生故障时,系统维护人员必须首先定位和诊断问题,然后通过回滚到先前的检查点来重新启动训练。然而,一些硬件故障可能很难主动检测,并且通常需要大量的时间来诊断和恢复,导致 LLM 训练的长时间停滞。其次,由于网络链路故障 [377] 或异常计算减速 [71] 导致的集群中的落后者,可以显著降低 MFU,进一步加剧整体训练效率的下降。Meta 的 175B OPT 模型的训练就是一个例子 [398]。虽然基于 MFU 的理想训练时间估计为大约 25 天,但实际训练持续了大约 57 天。这意味着惊人的 56% 的总时间被浪费在处理各种故障上,突显了系统不稳定对资源利用率和 LLM 训练效率的严重影响。
快速检测和诊断 LLM 故障对于维护训练稳定性和效率至关重要。这个过程,称为异常检测,主要采用两种方法:统计监控和主动验证。
统计监控是一种系统性的方法,用于观察和分析 LLM 训练过程中的各种指标和指示器。该方法涉及收集、处理和解释数据,以识别异常或与预期行为的偏差。在典型的设置中,每个 GPU 都有一个专用的监控进程,负责收集基本信息和运行时统计数据 [71], [378], [379]。然后将这些统计数据作为心跳消息传输到中央监控节点进行进一步分析。未能发送心跳消息的节点被视为已失败。这个监控系统的主要目标是及时检测异常,允许快速恢复,以最小化训练中断并保持整体效率。LLM 训练中监控的大多数运行时统计数据与硬件相关,包括 GPU 和网络指标。最近的工作 [71], [378], [379] 使用 NVIDIA DCGM [401] 收集与 GPU 相关的统计数据,包括 SM 块利用率、SM 占用率、SM 管道利用率、PCIe 流量率、NVLink 流量率等。一个经常出现的问题是 GPU 内存行重映射,它在硬件中无缝地用稀疏的单元替换已知的降级内存单元。Vela [378] 通过利用 DCGM 中的 DCGM_FI_DEV_ROW_REMAP_PENDING 统计数据来检测这个问题。Megascale [71] 和 Transom [380] 还通过分析训练日志中发生的错误来检测错误。除了 GPU 指标外,网络统计数据对于监控分布式训练性能至关重要。MegaScale [71] 跟踪 RDMA 流量指标以检测潜在的异常。它还开发了可视化系统,手动识别效率低下的 GPU。Unicorn [379] 检测诸如 NCCL 超时、TCP 超时和任务挂起等错误,并在训练期间延迟通知。C4 [377] 收集连接细节,如 RDMA IP 和 QP 号码,以及传输层的消息统计数据,包括传输计数、大小和持续时间,以检测训练减速和挂起。集体通信活动也可以通过 PyTorch 内置的 NCCL 飞行记录器 [273] 进行监控,该记录器将集体元数据和堆栈跟踪捕获到环形缓冲区,以便稍后诊断。Meta 进一步与 PyTorch 共同设计了 NCCLX [9],允许 PyTorch 访问其内部状态,以实现快速准确的故障检测。NCCLX 跟踪每个 NCCLX 通信的内核和网络活动,这可以帮助诊断通信问题。Vela [378] 实现了一个增强的 MultiNIC 健康检查器,收集每对节点上每个端口的节点网络带宽数据。这些信息可以用来检测具有降级 RoCE/GDR 性能的节点。利用 LLM 训练的关键特征作为先验知识,Transom [380] 开发了机器学习算法进行异常检测。统计监控还使 Google 的 TPUv4 超级计算机 [376] 具有弹性。每个 TPUv4 机器都配备了一个 healthd 守护进程,实时监控 IC 互联(芯片间互联)、PCIe 链接和 TPU ASIC。检测到的严重症状将通知集群调度程序采取适当措施,如驱逐受影响的作业或重新安排它们。
主动验证提供了一种替代基于在线统计监控的反应性故障排除的方法,旨在在故障发生之前验证训练系统。然而,验证测试时间和准确性之间存在权衡,因为全面的验证可能会显著影响有效的训练时间。MegaScale [71] 引入了一系列轻量级测试,包括网络主机和 NCCL 测试,以诊断广泛的潜在故障。Vela [378] 采用了两级策略,定期在每个节点上运行轻量级测试,并且仅在节点空闲时执行更具侵入性的测试。Google 的 TPUv4 超级计算机在用户作业之前实施了预飞行检查 [376],包括端到端检查和针对硬件健康的意向驱动检查器。SuperBench [381] 是一个全面的基准测试套件,用于评估单个硬件组件,结合了一个选择器来平衡验证时间和潜在问题相关处罚。
定期保存模型状态,即检查点,以及在发生故障后从最新检查点恢复计算是容错 LLM 训练的常见做法。然而,这带来了一个困境:频繁的检查点会带来高 I/O 开销,而不频繁的检查点在发生故障时会导致大量的进度损失。为了解决这个困境,设计了快速持久和内存中的检查点方法。
持久检查点涉及将模型状态保存到非易失性存储中,例如 SSD 和远程云存储,确保系统故障期间数据的持久性。该过程通常包括两个阶段:首先,快照阶段将模型状态从 GPU 复制到 CPU 内存,其次,持久阶段将快照写入持久性存储设备。尽管由于存储设备的低带宽而产生相当大的 I/O 开销,持久检查点仍然是由于其易用性和可靠性而广泛使用的一种容错方法。已经提出了先进的持久检查点方法,以减少训练停滞,从而实现更频繁的检查点,而不会显著影响性能。
同步检查点。 为了保持模型参数的一致性,DeepSpeed 的默认同步检查点 [22] 和 Varuna [167] 定期停止训练过程,以在数据并行等级 0 上同步执行持久存储的检查点。这种方法在快照和持久阶段导致 GPU 空闲时间,导致资源利用不足。认识到大多数故障可以归因于单个 GPU 或网络设备,JIT-Checkpointing [382] 提出了一种替代策略。它在发生故障后立即进行即时检查点,允许从这些即时检查点恢复训练。这种方法显著降低了浪费 GPU 时间的成本,将其限制在最多一个微批次迭代的工作上。DLRover Flash-Checkpoint [383] 利用分布式缓存服务加速迁移效率。Universal Checkpointing [384] 引入了一种通用检查点表示,将分布式检查点存储与并行化技术解耦。Universal Checkpointing 可以根据不同的需求轻松地将检查点从一个并行化策略转换为另一个。
快照 - 停滞检查点。 为了在保持检查点一致性的同时减少 LLM 训练停滞,Check-N-Run [385] 将快照和持久阶段解耦。它通过仅在快照阶段停滞训练并使用专用的后台 CPU 进程异步持久化快照,实现原子检查点。TorchSnapshot [386] 通过张量分块和多线程磁盘写入进一步优化了这个过程。通过创建分块的快照,TorchSnapshot 允许持久阶段通过并行写入更早地开始,从而减少整体训练停滞时间。MegaScale [71] 和 InternEvo [18] 也采用快照 - 停滞方法进行快速检查点和恢复。快照阶段停滞训练几秒钟以捕获模型状态,而持久阶段异步地将检查点从 CPU 内存传输到分布式文件系统。MegaScale 通过指定数据并行组中的单个工作器从分布式文件系统中读取来优化恢复过程,从而减轻了低带宽瓶颈。然后,这个工作器将检查点数据广播到其他 GPU,实现整个系统的更快、更有效的恢复。为了节省存储空间,InternEvo 还异步地将检查点从昂贵的热存储移动到更便宜的冷存储。
异步检查点。 异步检查点旨在通过将快照和持久阶段与训练同时执行来最小化训练停滞。DeepFreeze [387] 将轻量级(快照)和重量级(持久)持久策略应用于后台,将检查点跨数据并行 GPU 分片,以分布 I/O 工作负载。CheckFreq [388] 仔细地将快照和持久阶段与后续迭代的前向和后向传递进行流水线化,确保在下一次参数更新之前完成快照。它还动态调整检查点频率,以平衡恢复成本和运行时开销。LightCheck [389] 利用迭代间数据依赖性,引入分层检查点流水线以减少停滞。DataStates-LLM [390] 通过预分配固定主机内存进行快照,并利用高效的计算、快照和持久分层流水线。FastPersist [391] 识别了完全异步持久阶段的风险,并将其与下一次迭代的参数更新同步。它通过使用双缓冲固定内存提高 SSD 带宽利用率,并减少硬件争用,通过使用数据并行等级的一个子集进行检查点写入。
远程持久存储的低带宽严重限制了检查点的频率,内存中的检查点通过将检查点存储在其他计算节点的内存或专用的内存存储系统中来解决这个限制,显著减少了 I/O 开销,并实现了更高的检查点频率。Gemini [392] 提出将检查点保存到 CPU 内存以实现更快的故障恢复,以及检查点放置策略以最小化检查点丢失和流量调度算法以减少与训练的干扰。REFT [393] 异步地将模型状态缓存到主机内存和像 Redis 这样的内存存储中,绕过了检查点 I/O,并实现了高检查点频率。它还利用纠删码实现了 RAIM5(受 RAID5 启发,将'磁盘'替换为'内存'),以保护数据免受节点故障的影响。虽然这些方法通过实现更频繁的检查点而不带来性能惩罚,显著推进了 LLM 训练的容错能力,但它们可能无法提供与传统基于存储的方法相同的长期数据持久性。因此,结合内存和持久检查点的混合方法对于全面容错策略是必要的。
无检查点恢复方法旨在通过消除在发生故障时重新启动和回滚到先前检查点的需求来最小化训练停滞。这些技术依赖于自动故障检测机制及时识别问题。当检测到故障时,无检查点方法自动解决问题,并允许训练过程在不中断的情况下继续。通过避免从检查点加载和重复计算的耗时过程,这些方法可以显著减少停机时间并提高整体训练效率。无检查点恢复策略大致可以分为两种主要方法:实时迁移和模块冗余。
实时迁移利用分布式 LLM 训练设置中固有的冗余,特别是不同数据并行管道中的模型副本,以在发生故障时恢复模型状态。当检测到故障时,实时迁移方法使用剩余的健康实例或通过将新实例纳入训练集群来动态重新配置并行化策略。然后,将当前模型状态传输到这些重新配置的节点,允许训练过程以最小的中断继续。Parcae [394] 提出了三种不同的迁移机制,每种机制具有不同的通信开销,以有效地在不同的并行化策略之间传输模型状态。Oobleck [395] 采用了基于流水线模板的方法进行实时迁移。它维护一组预定义的流水线模板,并在检测到故障时,迅速根据这些模板实例化新的异构流水线。
模块冗余与实时迁移类似,也利用模型状态的冗余。然而,它不是在不同 GPU 上恢复最新模型状态,而是通过将计算路由到冗余模块来继续训练。Bamboo [396] 在持有同一流水线中相邻流水线阶段的 GPU 中放置一个冗余流水线阶段。这个冗余阶段在训练期间执行冗余计算,利用流水线泡沫,并在发生故障时作为正常阶段激活。SlipStream [397] 利用跨模型副本流水线的冗余,将故障节点的计算路由到不同数据并行流水线中的节点。SWARM [231] 提出了一个类似的解决方案,但更侧重于连接不良、异构和不可靠的设备。除了冗余计算外,SWARM 还结合了实例迁移来重新平衡流水线,结合了冗余计算和实时迁移方法的方面。
LLMs 的兴起已经改变了人工智能,使得像个人助手、代码生成和科学研究等应用成为可能。像 GPT、LLaMA 和 Gemini 这样的模型树立了新的标准,但训练这些庞大的模型,如在 16384 个 GPU 上进行的 LLaMA-3 的 54 天过程,提出了可扩展性、效率和可靠性方面的挑战。管理庞大的 GPU 集群需要创新的硬件和网络解决方案。高效的训练需要优化计算、通信和内存使用。可靠性涉及强大的机制,以便在长时间的训练期间检测和从故障中恢复。本综述回顾了 LLM 训练系统和基础设施的最新进展,突出了提高可扩展性、效率和可靠性的方法。传统的基于数字电路的计算系统,受摩尔定律和丹纳德缩放的指导,正面临满足训练和部署 LLMs 的计算需求的重大物理和经济限制。因此,AI 行业需要创新的解决方案。一种有前景的方法是大规模光电集成技术,它用集成硅光子学取代传统的数字电路,增强计算和通信能力 [402]。这种光电混合数据中心技术结合了光计算和光网络,提高了单节点计算能力,以及大规模分布式计算的效率。一些工作已经提出利用光网络进行 LLM 训练。例如,TopoOpt [67] 优化了分布式训练中的光网络拓扑和并行化策略,提高了计算和通信效率。TPUv4 [42] 使用光电路开关动态重构其 3D-Torus 互连拓扑,改善了 LLM 训练中密集通信模式的数据流。此外,Taichi [403] 探索了一种分布式衍射 - 干涉混合光子计算架构,有效地将光学神经网络扩展到百万神经元级别,能效为每秒 160 万亿次操作每瓦(TOPS/W)。未来可能需要在 LLM 训练和推理方面向硅光子学的范式转变。然而,这一转变将需要在系统设计和实施方面的广泛创新。