突破“延迟墙”:大规模模型训练中的数据移动瓶颈

科技   2024-11-20 20:35   北京  

大会预告

12月5-6日,2024中国生成式AI大会(上海站)将举办,30+位嘉宾已确认参会演讲。其中,上海交大副教授、无问芯穹联合创始人兼首席科学家戴国浩,GMI Cloud 亚太区总裁King.Cui,英飞流创始人张颖峰,Jina AI联合创始人兼首席技术官王楠,中昊芯英芯片软件栈负责人朱国梁将在AI Infra峰会带来分享,欢迎报名~


01 核心观点


02 作者洞察


随着模型规模和训练数据集的不断增大,训练所需的计算量快速增长,远超摩尔定律的速度。为了维持这种增长速度,需要进行三方面的扩展:延长训练时间、增加参与训练的GPU数量以及使用性能更高的GPU。

然而,随着训练计算规模的扩大,数据移动成为一个瓶颈,限制了效率的提升。在单节点内部,GPU内存带宽和PCIe带宽限制了数据传输速度;在节点之间,网络带宽和延迟成为主要瓶颈。这些瓶颈导致GPU利用率降低,甚至使得在固定时间窗口内完成训练变得不可能。

延迟墙

“延迟墙”(latency wall),是指在固定时间窗口内,无论使用多少GPU,都无法完成训练任务的上限。

原因:

  1. 更大规模的模型通常需要更多梯度步数进行训练。如果这些梯度步数必须在固定时间窗口内完成,这意味着每个梯度步数所需的时间必须缩短。
  2. 当每个梯度步数的时间变得非常短时,延迟限制将不允许训练该规模的模型。

03 作者论断


  1. 当前数据移动瓶颈限制了LLM的扩展,使其无法超过2e28 FLOP,并在2e31 FLOP处存在“延迟墙”。预计可能会在约3年内遇到这些限制。
  2. 改进硬件互连最多可以将训练运行规模扩大两个数量级,前提是采用类似于当前范式的技术。
  3. 超越硬件互联的改进,关键的创新必须来自机器学习算法。关键挑战是将两个串行依赖关系(批次之间和层之间)转化为并行机会,方法是增大批大小(可能通过稀疏性实现)并使模型更宽更浅。
  4. 激进的批大小缩放可能会克服数据移动瓶颈的限制。

04 常见问题问答


什么是数据移动瓶颈,为什么它对大规模模型训练很重要?

数据移动瓶颈是指在训练大型机器学习模型时,数据在GPU内或GPU之间移动的速度限制了整体训练速度。当模型大小和训练数据集大小增加时,这个问题变得尤为突出。如果数据移动速度跟不上计算速度,GPU的计算能力将无法得到充分利用,从而降低训练效率。

影响大规模模型训练的主要数据移动类型有哪些?

  • GPU内数据移动:主要指数据在GPU的DRAM内存和计算单元之间移动,例如将模型权重从DRAM加载到片上SRAM或寄存器。
  • GPU间数据移动:主要指数据在不同GPU之间传输,例如在使用数据并行或张量并行时进行梯度同步。

有哪些常用的并行策略来加速大规模模型训练?

  • 数据并行:将训练数据分割到多个GPU上,每个GPU训练一部分数据,然后汇总梯度更新模型。
  • 张量并行:将模型参数分割到多个GPU上,每个GPU计算一部分模型,然后进行结果同步。
  • 流水线并行:将模型的不同层分配到不同的GPU上,数据像流水线一样依次经过各个GPU进行处理。

专家并行:将模型分割成多个专家模型,每个专家模型处理一部分数据,然后汇总结果。

如何缓解GPU内数据移动瓶颈?

  • 加模型宽度和批大小:可以增大矩阵乘法运算的规模,提高算术强度,从而更好地利用GPU的计算能力。
  • 使用更高效的内存访问模式:例如使用矩阵分块技术,将数据加载到片上SRAM后进行多次计算,减少DRAM访问次数。
  • 使用更高带宽的内存:例如HBM2e等高带宽内存可以提供更高的数据传输速度。

如何缓解GPU间数据移动瓶颈?

  • 使用更高带宽的网络互联技术:例如NVIDIA NVLink和InfiniBand等高速互联技术可以提供更高的数据传输速度。
  • 使用更高效的通信算法:例如使用环形或树形拓扑结构进行All-reduce操作,可以减少通信次数和数据传输量。
  • 优化并行策略:例如根据模型结构和硬件特性选择合适的并行策略组合,可以最大限度地减少通信开销。

为什么即使网络带宽无限,延迟也会限制大规模模型训练的效率?

即使网络带宽无限,每次数据传输的延迟仍然存在。当模型规模很大,需要频繁进行GPU间通信时,这些延迟会累积起来,形成一个显著的瓶颈,限制整体训练速度。

未来有哪些技术发展方向可以帮助突破大规模模型训练的数据移动瓶颈?

  • 更高带宽、更低延迟的网络互联技术:例如光互联技术和新型网络拓扑结构可以提供更快的通信速度。
  • 更接近计算单元的内存:例如将内存集成到芯片上,可以减少数据移动距离和延迟。
  • 更高效的并行策略和算法:例如异步通信和模型压缩技术可以减少通信开销。
  • 专用硬件加速器:例如专门用于机器学习模型训练的芯片可以提供更高的计算效率和更低的数据移动开销。

基于当前硬件和算法,大规模模型训练的规模极限在哪里?

根据Epoch AI的研究,基于当前硬件和算法,使用NVIDIA H100 GPU训练密集模型,在保证高利用率的情况下,最大训练规模约为1034 FLOP以上。

05 引言


近五年来,大语言模型(LLMs)性能显著提升,这主要归功于计算资源在更大规模模型和数据集训练上的快速扩张。据我们估算,从2010年至2024年,前沿AI模型的训练算力需求年均增长约4至5倍。这一增长速度远超摩尔定律,要维持如此高速增长,需要从三个维度同步扩展:延长单次训练时长、扩充训练所需GPU数量,以及提升单个GPU的性能。

延长GPU集群的训练时长相对简单。但实践中,训练周期很少超过半年。这是因为硬件和软件的迭代周期较短,研究机构都不希望其模型在发布时就已显得过时。这实际上为训练时长设定了上限。

扩大训练集群规模是另一个主要方向。例如,从Vaswani等人2017年提出的首个Transformer模型(8个GPU训练12小时),到Llama 3.1 405B(约1.6万个GPU训练2个月),我们可以看到这种显著的扩展。

然而,扩大集群规模远比延长训练时长复杂:数据传输带来的根本限制决定了在特定时间内能有效协同工作的GPU数量。参与同一训练任务的GPU越多,每秒需要传输的数据量就越大,这既包括GPU内部的数据传输,也包括GPU之间的数据交换。

根据Epoch AI研究员Ege Erdil和David Schneider-Joseph的最新研究,在当前技术条件下,数据传输瓶颈很可能会限制训练规模的扩展。研究估算表明,假设最大训练时长为3个月,当训练计算量超过2e28 FLOP时,继续扩展将失去效率。这是因为在超过该规模后,数据传输时间将主导整个计算过程,导致硬件资源无法被充分利用。

图1:顶级机器学习(ML)模型的训练计算量正在逼近当前GPU架构的理论极限,数据传输瓶颈日益凸显。目前的训练规模与这一极限值仍有约100倍的差距,一旦超过这个阈值,扩展效率将显著下降。值得注意的是,尽管NVIDIA H100在算术吞吐量(arithmetic throughput)上有所提升,但由于其互连(interconnect)和内存带宽(memory bandwidth)的增长速度未能匹配算术吞吐量的提升,导致数据传输瓶颈更为突出,其最大有效训练规模实际低于A100。

即便不考虑效率损失而强行扩大集群规模,最终也会遇到网络延迟瓶颈,这限制了设备间的通信速度上限。如果没有互连技术的重大突破,或者没有找到比随机梯度下降(Stochastic Gradient Descent,SGD)更高效的优化算法,这些限制将难以突破。

本文将深入探讨分布式训练中的数据传输模式,分析大语言模型规模扩展的制约因素。我们将研究GPU内部和GPU之间的数据传输特点,讨论各类并行策略,并探讨训练规模扩展可能遇到的瓶颈。通过本文,读者将了解未来规模扩展面临的挑战,包括利用效率下降、延迟瓶颈等问题及其可能的解决方案。

06 分布式训练中的数据传输模式


要理解数据传输如何制约分布式训练,首先需要了解其中涉及的主要数据传输类型。总体可分为两类:GPU内部数据传输(intra-GPU data movement)和GPU间数据传输(inter-GPU data movement)。

GPU内部数据传输

在当前分布式训练体系中,高带宽内存(HBM,也称为VRAM)承担了几乎所有GPU内部的数据传输开销。通常,执行两个大矩阵的乘法运算时,需要从HBM读取这两个矩阵,并将计算结果写回HBM。由于内存带宽有限,这个过程需要一定时间。

当矩阵规模较大时(如总元素数超过1000万),GPU核心的计算时间远大于内存访问时间,此时GPU内部数据传输开销可以忽略。通过优化内核编程,还可以实现计算与内存访问并行,从而“掩盖”内存访问延迟。但随着矩阵规模减小,计算时间降低速度快于内存访问时间,当降至某个临界点时,内存访问将成为瓶颈。此时将无法掩盖内存访问延迟,计算资源利用率开始下降。

这种情况对分布式训练构成了约束,原因如下:根据Hoffmann等人(2022)提出的Chinchilla规模定律(scaling law),计算最优的训练要求模型的训练计算量与参数数量成二次方关系增长。然而,单个批次中每个Transformer块所需的计算量通常与参数数量呈线性关系,具体取决于批次大小能否随模型规模同步扩展。这种扩展不匹配意味着在大规模训练中,矩阵乘法操作需要被分解为更小的片段;当这些片段在矩阵某一维度上变得足够小时,内存带宽约束可能会显著影响性能。

GPU间数据传输

GPU间数据传输的目的是实现计算结果的共享,确保所有GPU能在同一训练任务上协同工作。训练任务的并行化主要分为四种模式,每种模式对应不同的GPU通信模式:

数据并行(Data Parallelism):各GPU处理不同的输入序列或批次中的不同token。每个GPU可独立完成前向和反向传播,但在更新权重前需要聚合梯度信息,这要求在每次梯度更新前进行一次通信。

流水线并行(Pipeline Parallelism):各GPU负责处理Transformer中不同的连续层或块。GPU间需要按序传递输出,直到输入数据流经模型的所有层块。每个流水线阶段的交界处都需要通信。值得注意的是,如果在单个批次上直接实施流水线并行,会出现除当前计算GPU外其他GPU空闲的“气泡”(bubble)现象。为避免这种情况,通常将批次划分为多个微批次(micro-batches),并按序输入流水线。

张量并行(Tensor Parallelism):各GPU负责处理模型权重矩阵的不同分片进行矩阵乘法。与流水线并行的水平分层不同,张量并行采用垂直分片策略。其特点是每完成一到两次矩阵乘法(具体取决于实现方式)就需要同步结果。在所有并行方式中,张量并行要求最频繁的GPU间通信,因此通常只在高带宽、低延迟的互连网络(如NVIDIA NVLink)上部署。

专家并行(Expert Parallelism):这种并行方式专门用于混合专家(Mixture of Experts, MoE)架构的模型。每个块包含多个“专家”,分布在不同GPU上。系统通过轻量级路由机制决定每个输入在各块中应该由哪些专家处理。专家计算完成后,需要合并结果再传递给下一层。因此,专家并行在每个块都需要进行两次通信:专家路由和结果合并。

在特定规模的集群中,要最小化GPU间数据传输量,需要组合使用多种并行策略。这是因为每种并行方法的数据传输开销都随其并行度呈线性增长(张量并行呈次线性增长),而集群规模则是所有并行度的乘积。因此,同时扩展多种并行方法能够最小化特定训练任务的总体数据传输量。例如,Falcon-180B的训练采用了64路数据并行、8路流水线并行和8路张量并行的组合,总计使用64×8×8=4096个GPU。

07 扩展的极限


经过详细建模分析各项影响因素后可以得出:在现有技术和硬件条件下,若要保持良好的硬件利用率,无法在3个月内完成超过2e28 FLOP计算量的训练任务(对于计算最优的密集模型来说,大约对应15万亿参数)。更进一步,当训练计算量超过2e31 FLOP时,受限于延迟约束,除非牺牲随机梯度下降(Stochastic Gradient Descent, SGD)的逻辑一致性,否则训练任务将无法实现。

图2:模型预测结果展示了在特定训练计算量下,密集训练过程能够达到的最优模型FLOP利用率(Model FLOP Utilization, MFU)。

在制约分布式训练扩展性的多个因素中,“延迟墙”(latency wall)是最易量化分析的限制因素。它为固定时间窗口内可完成的训练任务设定了理论上限,这一限制与投入的GPU数量无关。“延迟墙”产生的根本原因是:随着模型规模增长,所需的梯度更新步数也相应增加;在固定时间窗口约束下,要完成更大规模模型的训练,每步梯度计算的时间必须不断缩短。这种趋势最终会导致单步梯度计算时间小于系统延迟下限,使得该规模模型的训练在物理上变得不可行。

08 “延迟墙”的定量分析


从定量角度看,假设模型参数量为N,根据Chinchilla规模定律(scaling law),需要使用D=20N个token进行训练。若批次大小为B个token,且仅训练一个epoch,则需要执行n_steps=D/B=20N/B次梯度更新。在总训练时间t_train的约束下,每次梯度更新的最大允许时间为(t_train×B)/(20×N)。

另一方面,在高度张量并行的训练场景中,每次矩阵乘法后都需要跨GPU通信,每次通信至少消耗网络延迟时间t_latency。对于包含n_layer层的Transformer结构,考虑到每层有两次矩阵乘法,以及每次梯度更新包含前向和反向传播,在延迟约束下的可行性条件为:

代入实际参数:网络延迟t_latency=9微秒,层数n_layer=120,批次大小B=400万token(参考GPT-3和Falcon-180B的实践经验),训练时间t_train=3个月,可得出参数量上限N=400万亿。按Chinchilla最优训练计算,这相当于约2e31 FLOP。

实际应用中,层数(n_layer)和批次大小(B)会随模型规模增长而变化。如果这两个变量的增长效应大致相互抵消(这种情况是可能的),上述计算结果仍然有效。更复杂的情况(如两种效应不完全抵消)在我们的论文中有详细讨论,此处不再展开。

在上述讨论的各项限制因素中,我们预计“延迟墙”(latency wall)将成为未来最难突破的瓶颈,尽管在当前的训练任务中延迟因素的影响还可以忽略不计。这主要基于以下两个原因:

与内存带宽和网络带宽限制不同,即使牺牲训练效率也无法突破“延迟墙”对固定训练时间内模型计算量的限制。

自2017年NVIDIA V100发布以来,在所有相关硬件指标中(包括HBM带宽、L2缓存带宽、网络带宽等),设备内部延迟(如内核启动和终止延迟)和设备间延迟(如NVLink和InfiniBand延迟)的改进最为缓慢。

09 突破方案探讨


图3:展示了延长高效(或线性)扩展阶段并克服“延迟墙”的潜在路径。蓝线代表了一种理想情况:假设带宽无限(包括内存带宽和网络带宽),同时所有延迟降至原有水平的十分之一。

尽管上述模型似乎为训练规模设定了严格上限,但仍存在几种可能的突破路径。最直接的方案是通过硬件升级降低延迟,但从近几代GPU的发展趋势看,这一方向进展缓慢。

另一种选择是保持硬件不变,转而改变训练方法。让我们回顾一下由“延迟墙”(latency wall)导致的模型参数量上限公式:

由于这个上限是基于Chinchilla规模定律(scaling law)推导得出,突破延迟墙的一个直接解决方案是对模型进行欠训练(undertrain):具体来说,就是将具有N个参数的模型训练的token数量控制在20×N以下。这种方法理论上可行,因为在固定的训练时间窗口内需要执行的前向传播和反向传播次数减少了,从而缓解了带宽和延迟的约束。然而,根据Chinchilla规模定律(scaling law),这种策略会降低训练算力到模型能力的转化效率。因此,我们认为这不是一个突破延迟墙的可行方案。

如果要在保持训练逻辑一致性且避免欠训练(undertrain)的前提下提升上限,我们可以考虑调整公式中的变量。由于训练延迟(t_latency)受硬件限制,训练时间(t_train)被视为常量,可调整的空间只剩下增加批次大小(batch size)和减少模型深度(model depth)两个方向。

扩大批次规模

如果能找到方法使批次大小(batch size)随模型规模同步大幅扩展,我们就能有效延缓“延迟墙”的影响。这种扩展同时也为数据并行(data parallelism)和流水线并行(pipeline parallelism)提供更多可能性,从而推迟内存或网络带宽成为主要瓶颈的时间点。然而实践中,批次大小的扩展会遇到临界批次大小(critical batch size)的瓶颈,超过这个阈值后训练效果会显著下降。

McCandlish等人(2018)研究表明,模型性能提升会带来临界批次大小的增长。由于规模更大的模型在训练后期表现更优,这意味着更大的语言模型应能支持更大的批次大小。根据他们的模型,临界批次大小应与模型可降低损失(reducible loss)的倒数成正比。基于这一结论并结合Hoffmann等人(2022)的规模定律(scaling law),我们模型的基准结果已经允许批次大小按7e31 FLOP,则需要更激进的扩展策略。

在我们研究发布前几天,Zhang等人(2024)发现了一个新的扩展关系:B = 17.75 × D^0.47(单位为tokens)。如果在我们的模型中严格采用这种更激进的扩展关系,利用率下降的时间点将推迟两个数量级,从约2e28 FLOP延后到3e30 FLOP。当然,Zhang等人(2024)的研究并未探索其他可能的激进扩展方法,如使用替代优化器等。

总体而言,如果要通过算法创新突破“延迟墙”,最可能的方向是找到一种使训练批次大小具有比当前假设更强扩展性的方法。

降低模型深度

另一个可选方案是降低模型深度的扩展速度。目前学术界针对模型深度和宽度独立进行的扩展法则研究还较为匮乏。我们只能通过分析一些在不同规模下大量训练模型的研究(如Hoffmann等人,2022)来推测当前领域的最佳实践。

根据Hoffmann等人(2022)的研究,模型层数的规模定律(scaling law)估计为 n_layer ~ N^0.27,这与我们模型中临界批量大小(critical batch size)的缩放指数N^(1/3)接近,因此这两个规模效应在数量级上基本相互抵消。然而,需要指出的是,Hoffmann等人的研究存在两个局限性:首先,未进行实验验证该缩放策略的最优性(optimality);其次,没有建立用于量化偏离最优策略的性能损失(performance penalty)的缩放律。因此,这种扩展关系最多只能视为未来模型深度扩展的初步参考。

如果保持模型深度相对稳定,同时按预期速度扩大批次大小,训练计算的“延迟墙”可能会推迟多个数量级。相比之下,单纯采用更激进的批次大小扩展策略的可能性较低。

10 结论


我们的模型表明,常规的模型规模和集群规模线性扩展方式将在约2e28 FLOP规模处遇到瓶颈。考虑到已知最大模型的估计计算量约为5e25 FLOP,且历史计算扩展速率约为每年4-5倍,我们预计将在约3年内达到2e28 FLOP的临界值。

通过硬件改进将这一临界值显著推移(即使是一个数量级)都极具挑战性,因为这需要内存带宽(memory bandwidth)、网络带宽(network bandwidth)和网络延迟(network latency)的同步显著提升。相比之下,更有希望的突破方向是在训练过程中采用更激进的批次大小扩展策略。这种策略的可行性取决于临界批次大小(critical batch size),而在大语言模型(LLM)中,临界批次大小与模型规模的关系目前仍不明确;尽管我们推测,领先研究机构可能已掌握尚未公开的临界批次大小扩展规律。

临界批次大小扩展规律、模型深度和纵横比(aspect ratio)扩展规律以及稀疏度因子扩展规律(如MoE模型采用的策略)对理解分布式训练的未来至关重要;但目前,这些领域的可靠研究都很稀少。如果能获得更多这些方面的可靠研究成果,我们对分布式训练未来发展的理解将大大提升。

----------

参考资料:Erdil, E., & Schneider-Joseph, D. (2024). Data movement limits to frontier model training. arXiv. https://arxiv.org/abs/2411.01137

—END—


点击下方名片


即刻关注我们


算力猩
隶属于智猩猩,关注计算芯片创新,解读中国算力突破。
 最新文章