23年1月论文“Systems for Parallel and Distributed Large-Model Deep Learning Training“, 来自UCSD。

深度学习(DL)已经改变了各种领域的应用,包括计算机视觉、自然语言处理和表格数据分析。对提高DL模型精度的探索促使探索越来越大的神经架构,最近的一些Transformer模型跨越了数千亿个可学习参数。这些设计为DL空间带来了规模驱动系统挑战,例如内存瓶颈、运行时效率低和模型开发成本高。解决这些问题的努力已经探索了一些技术,如神经架构的并行化、在内存层次结构中溢出数据以及高效内存的数据表示。这项调查将探索大型模型训练系统的前景,强调关键挑战和用于解决这些挑战的各种技术。

DL实践的最新发展为DL研究引入了系统模型规模的新挑战。实践者已经开始探索将非常大的神经结构图用于DL模型,其中一些包含数十亿甚至数万亿的可训练参数!关键示例包括NLP Transformer模型BERT Large[16]、GPT-3[13]和Meta的深度学习推荐模型(DLRM)[41]。这些模型的庞大规模在三个关键领域带来了严峻挑战。

  • (1) 内存可扩展性。标准DL训练通常将模型的参数保存在加速器(例如GPU)的存储器上,并使用采样数据来计算每个参数的梯度更新。对于非常大的模型,保存参数、中间计算和梯度更新所需的空间通常会超过加速器相对有限的内存资源。高端消费级GPU,如特斯拉V100[2],具有16-32GB的设备内存,但大型DLRM可能需要数百GB的空间。
  • (2) 性能。参数计数的增加通常与较高的执行时间有关。此外,复杂的模型架构往往需要大型数据集来提高模型的学习能力——例如,GPT-3在300Btoken上进行训练[13],开源BLOOM语言模型在366B[12]上进行训练。训练这样的模型可能需要数周甚至数月的时间[12,41]。因此,可以提高执行性能的优化对大规模DL模型的开发人员非常有益。
  • (3) 训练费用。前两个挑战的标准解决方案通常涉及跨多个设备并行存储或执行。然而,这种方法可能会显著提高计算成本。BLOOM使用416个A100 GPU进行训练,Megatron LM使用512个[5]。这对大多数从业者来说是不现实的,尤其是当这些GPU需要保留几周甚至几个月的训练时间时。在AWS上复制BLOOM的训练程序将花费550万美元。请注意,这甚至没有考虑到模型选择的额外成本,包括训练多个模型以评估最佳超参数设置和配置[29]。

随着不断突破模型规模的界限,解决这些挑战以实现大模型DL空间的进一步发展变得越来越必要。因此,已经开发了各种系统和技术来解决这些问题。一些方向包括重具体化(rematerialization)[15]、数据溢出/CPU卸载[23,36,37,45–47]、流水线/模型并行[21,25,27,33,40]和混合并行[26,31,36,37,60]。这些主题属于“大模型训练技术”的保护伞,已成为工业界和学术界研究人员的重点,但该领域的工作量之大使该主题难以确定方向。本文回顾大模型DL训练系统空间的现状,并评估该领域未来的增长和发展方向。

最近发布了一些关于该领域的高级别、简短的综述[20,54],但并不全面,也没有涉及模型选择、混合并行性和技术交叉等关键主题。本文将讨论这些领域,并进一步深入研究最先进的技术。

模型并行是指将神经架构图划分或分片为子图,并将每个子图或模型分片分配给不同设备的技术。在前馈网络中,这些碎片可能指的是堆叠层的组。

模型并行性的加速潜力在很大程度上取决于体系结构和分片策略。前馈网络上的序列模型分片(如图A所示)将不提供并行执行的范围,而是在加速器之间引入相关图(dependency graph)。然而,这种分片策略在序列体系结构(如Transformers)中仍然很流行,因为它可以在多个设备之间分配内存需求,而且设置起来相当简单。在其他情况下,神经计算图为算子间并行提供了自然的机会(如图B所示)。

如图所示:A) 说明了如何将不适合单个GPU的大型前馈网络在三个GPU上进行模型并行化以实现执行;(注:执行速度并没有加快——没有并行执行,只有分区内存需求)B) 三个GPU上的性能模型并行分片策略的示例,并行执行层用共享颜色表示;这种策略利用了算子图中现有的并行执行机会——如果图本质上更序贯处理,这可能并不一定。

大型深度学习训练の并行分布式系统_数据

另一种方法是实际划分网络中的各个算子。一些算子(如嵌入表)可以以最小的开销按宽度进行分片(sharded)。其他,如矩阵乘法,仍然可以进行分割(例如使用并行GEMM策略[55]),但需要更多的通信步骤。如图展示了一个模型并行嵌入表的示例。这些宽度分片策略,通常被称为张量并行,因为它们需要输入张量分区,可以实现比算子间模型并行更高性能的算子内并行,但需要更多的精力和思考来实现。此外,大多数张量并行算子至少需要一个全聚集(all-gather)通信步骤来重新聚集分区输出,这一事实削弱了性能优势。Mesh TensorFlow[51]提供了一个通用的框架和语言来指定张量并行计算,尽管它不再被支持或维护。

大型深度学习训练の并行分布式系统_并行性_02

任何类型的模型并行都引入了GPU-GPU通信。最新的英伟达GPU支持“NVLink”互连,即提供高达900GB/s带宽的高速GPU-GPU通信路由,这有助于最大限度地减少开销。然而,NVLink并不总是现成的,尤其是当用户无法轻松定制的云端机器。当不支持NVLink时,GPU-GPU通信通过PCIe互连运行,速度要慢得多。特斯拉V100通常被认为是DL应用程序的标准高性能GPU,支持16GB/s的16通道PCIe 3.0互连。

为了避免通过缓慢的互连传输过多数据,模型并行性用户通常会选择一种分区策略,该策略将最小化需要在片之间传输的激活的数量,或者平衡计算以隐藏通信成本。为此存在各种分片算法[14,25,26,37,46,61]。

数据并行是一种常见的深度学习执行策略,可以并行消费多个小批量数据。数据并行执行技术可以分为两大类——异步数据并行和同步数据并行。

最著名的数据异步并行技术是Parameter Server,其中一个核心主服务器持有一组基线参数,而分布式worker持有在不同的小批量上训练的模型副本。分布式工作程序偶尔会向基线服务器发送更新,而基线服务器又会向分布式工作程序发送替换参数,以保持它们的更新。工作程序可能会彼此不同步,因为它们只需要与基线服务器通信/同步。异步技术带来了许多挑战,例如与单个worker训练相比,准确性下降,以及由于worker返回时间的差异而导致的不可复制的结果。由于这些原因,在现代DL训练环境中,异步技术通常被逐步淘汰,取而代之的是同步技术。

最流行的同步数据并行执行技术是分布式数据并行(Distributed Data Parallelism,DDP)。DDP复制模型并将副本分配给 不同的加速器。首先接受一个初始的“全局小批量”,然后在副本之间均匀分解,为每个副本生成本地梯度更新。然后,这些梯度在副本之间聚合,产生全局更新,通常使用 all-reduce通信模式。然后将此全局更新并行应用于每个复制副本。该技术在数学上等效于使用原始全局小批量的单个GPU训练。虽然这种技术引入了一个 all-reduce通信步骤,但这些开销通常可以在模型执行时间下重叠和隐藏。

All Gather和All Reduce通信模式通常用于数据并行和更广泛的分布式DL。这些模式的目的是在不同的处理器上获取单独的数据,然后将它们聚合并分发回处理器,这样每个处理器现在都拥有相同数据的副本。all-together模式都使用一种算法,其中每个处理器将其数据传送给其他每个处理器。如果每个处理器都有一个需要全局广播的数据分区,则通常会使用此方法。每个带宽使用率高的处理器通常需要 个通信步骤——每个处理器必须与所有其他处理器进行通信。All-reduce模式是All-together之上的一层,它将与一些reduction函数(例如总和、平均值)相结合。在同步过程中,运行函数与运行all-together然后运行局部函数应用程序相比,省去了一个步骤。例如,Horovod[50]实现了一个为数据并行梯度聚合的带宽最优reduce模式,其中每个处理器需要2×( − 1) 个通信步骤完成完全的gather和educe。

混合并行是指将不同的并行策略结合起来以实现更高的整体性能的策略。例如,将数据并行性叠加在模型并行性之上,可以使用户实现跨多个设备的内存可扩展性,同时加快数据并行性的执行速度。这些策略需要在设计中进行权衡。一个简单的覆盖混合并行,模型并行性的多设备需求要乘以数据并行性的复制要求。任务并行性的进一步叠加(例如,在多模型训练中)可以将另一个乘法因子添加到等式中。更复杂的混合可能会将数据并行应用于模型并行体系结构的某些阶段,让其他阶段串行执行,如图所示。注意这个设计打开了新的“搜索空间”——数据并行复制应该选择哪些阶段?模型并行分片决策如何影响数据并行性能?有限的资源应该如何分配到各个阶段,设备互连和拓扑结构如何影响性能?

大型深度学习训练の并行分布式系统_并行性_03

模型架构的规模大致分为两类——深度规模化和宽度规模化。深度规模化是像Transformers这样的长序列链架构最常见的需求。宽度规模化通常用于非常宽、易于并行化的算子(例如表查找)。

模型并行等技术对于大型Transformer训练和一般的深度模型训练来说是必不可少的。然而,为非常深的模型启用并行执行可能具有挑战性。对于一个深的层序列,最自然的分片策略是将序列划分为子序列。但这种方法迫使用户添加GPU,而实际上并没有从任何性能优势中获益——将序列划分为子序列并不能提供任何并行执行加速的机会。考虑一个万亿参数模型,它甚至需要使用1024个GPU才能适应内存。所有这些GPU都只是用于“启用”执行,并没有提供任何性能优势。事实上,由于GPU间的通信成本,该策略可能比同等的内存内训练作业慢。

存在一些宽度分片(width-wise sharding)策略,例如跨多个GPU并行处理注意块中的操作。然而,这些方法需要更多的定制,增加通信开销,并且需要模型设计者付出大量的努力来实现。因此,大多数用于深度模型训练的系统更喜欢应用可以针对所有深度模型类进行优化的广义深度分片策略,而不是一次针对单个架构。

尽管存在顺序依赖性的挑战,但深度分片也可以带来许多机会。流水线并行和溢出(spilling)等技术只适用于深度分片模型(depth-wise sharded models)。

在推荐模型中,嵌入表通常是宽度分片最受欢迎的候选者。大多数公司都使用基于嵌入的电子商务模型,这些公司收集特定实体的数据(如Meta、Netflix、TikTok)来创建定制的体验。一种标准的方法是创建一个表,将用户ID映射到可训练向量,然后将这些向量馈送到顶部的其他DNN。然而,要想在Facebook这样的数十亿用户平台上运行,相应的表格必须非常宽。一个30亿的索引表,大小为1024个可训练向量,填充单精度(32位)浮点,需要12TB的内存。真实世界的推荐模块可能包括用于查找的多个表(例如,用户表、业务表、视频目录表),这进一步增加了内存成本。

对嵌入表进行分区是一项简单的任务,因为表的查找是并行的——一个索引的查找不依赖于其他索引。因此,将表分配给不同GPU成为子表,是分配内存成本的常见策略。跨片并行执行,可以简单地将小批量中的索引查找请求路由到适当的GPU。然而,为了在并行表查找之后重新聚合小批量馈送到顶部DNN,需要一个潜在的昂贵的all-together通信步骤。

将宽度分片应用于其他算子(如矩阵乘法)并不常见,但也并非闻所未闻。但总的来说,嵌入表是内存最密集的单个操作[9]。考虑到宽度分片的主要用例是嵌入表,针对这种情况进行优化可能显得过于特殊。然而,嵌入表和推荐模型在DL工作负载中占了很大比例——Meta报告称,他们50%的DL训练周期都花在了基于嵌入表的推荐模型上[9]。因此,优化非常广泛的模型情况是非常值得的,即使适用性比序贯深度模型可扩展性的优化更有限。

如图对不同训练系统和技术进行比较:

大型深度学习训练の并行分布式系统_人工智能_04

一些基本技术,如再具体化,通常被用作更先进大型模型训练系统的常见构建块。一般来说,这些技术对组织和结构的影响很小,可以与其他方法集成。

再具体化,也称为梯度检查点,试图最大限度地减少反向传播的内存需求[15,19]。反向传播需要保存中间算子输出,以便正确应用梯度计算的链式规则。然而,中间输出张量可能需要大量内存!一些分析[54]表明,激活占ResNet[22]内存消耗的95%,占某些Transformer内存使用的80%。最初先丢弃除少数检查点之外的大多数激活,用检查点重新计算反向传播过程中丢弃的激活,再具体化以计算换取内存。通过这种方式,在任何给定点,只有检查点之间的中间点需要存储在内存中。这种方法确实会产生计算开销——前向传播有效地进行了两次。然而,前向传播中的算子通常比反向传播中使用的自动微分程序更快,因此开销比看起来更小。一些梯度检查点系统声称只有30%的开销,可以节省6-7X内存[15]。

累积,是针对反向传播中分批梯度的内存需求而言[25]。随机梯度下降将样本分批放入模型馈送的小批量中,反过来可以将参数更新生成的梯度视为每个样本更新的聚合。累积延迟了这些聚合梯度的应用,而是计算新的小批量梯度更新,并将它们累积到聚合梯度向量上。新梯度现在是2个小批量更新的总和,而不是1个。通过这种方式,可以扩大有效的小批量大小和梯度影响,而无需实际训练更大的批量。将较小的单个批次称为微批次(micro batch),并将有效的合计批次称为小批次(mini-batch)。累积对于流水线并行性至关重要,并且经常与其他技术结合使用。

大多数训练框架(例如TensorFlow、PyTorch)[8,42]使用梯度和参数的单精度浮点(32位)表示。双精度表示(64位)相对不常见。减少训练模型的内存需求的一种方法是使用数据的半精度(16位)表示,即低精度表征。自然地,当数值被近似时,这会导致精度损失[38]。然而,这种方法可以提供加速和内存节省。为了尝试和平衡这一点,自动混合精度(AMP)[3]将自动尝试并确定何时可以安全地将数据压缩到16位,而不会造成精度损失。AMP在训练大型模型时,精度损失很少甚至没有损失,同时实现了高达5.5X的加速[3]。由于AMP直接在非常低的级别修改数值,因此该技术通常与用于大模型训练的实际系统方法正交。

在某些情况下,DL训练中使用的向量非常稀疏。例如,嵌入表查找通常只涉及表的几个索引。应用于表的梯度向量将仅在使用的索引处具有非零值,而其余部分置零。实际上,在内存中保留所有这些零是不必要的,并且会浪费内存。稀疏表示试图将这些向量压缩直至非零值,同时避免任何信息丢失。默认情况下,通常用于嵌入表的最简单方法是将梯度表示为将索引映射到梯度值的K-V对。当将稀疏表示与假设标准向量表示的操作相结合时,会出现一些复杂情况,例如all-reduce通信模式。一些工作[35]展示了如何通过替代通信模式或将数据转换回标准表示来解决这一问题。稀疏向量表示解决了一个非常具体的问题,但对于一些算子(如宽嵌入表)的有效训练至关重要。

流水线并行性针对“序列深度模型”设置[25]。它是模型并行训练范式的直接扩展。模型并行性创建一个分阶段的片序列,创建一个自然的“流水线”结构。流水线只是通过尝试用执行操作填充阶段来利用这种流水线结构,从而减少纯序列模型并行性中存在的空闲。每片都可以被视为流水线的一个阶段,因此一个在三个GPU上三次分区的模型现在是一个三阶段流水线。

在CPU流水线中,用发送到CPU的各种指令填充流水线[52]。对于DL流水线,用微批次填充流水线,就像在梯度累积中一样[25,27]。从本质上讲,流水线并行是梯度累积和模型并行的结合。独立的微批次在分片流水线中穿梭,然后为每个流水线阶段积累每个微批次的梯度。一旦整个小批次(所有微批次的组合)的梯度全部聚合,就可以将其应用于模型。这种设计几乎就像一个模型-并行-数据-并行的混合,其中数据片是并行处理的,但在不同的模型-并行片上。如图对此进行了说明:输入的小批次被划分为微批次,然后在流水线阶段中运送。, 指的是带有小批量 的片 前阶段, 虽然 , 指的是小批量 的片 后阶段, 通过这种方式,实现了一种“流水线式”数据并行,即在不同的模型并行级之间并行处理数据。请注意,在反向传播之前,必须清除正向流水线。

大型深度学习训练の并行分布式系统_人工智能_05

反向传播对流水线并行训练提出了挑战。中间输出必须可用于反向传播。然而,当与累积相结合时,这将要求为每个微批次存储不同的中间输出集,从而剥夺了累积提供的任何可扩展性优势。GPipe[25]是最早的流水线并行训练系统之一,提出了将累积与检查点相结合来解决这个问题。激活将仅存储在片/流水线阶段的绑定中,在反向传播过程中,随着梯度在流水线中向后移动,将进行重新计算。检查点方法现在是大多数(如果不是全部的话)流水线并行训练系统的标准方法[21]。另一个挑战是流水线的结构。片流水线必须是双向的。输入和激活在预测期间向前流动,梯度在反向传播期间向后流动。这导致了一个问题——流水线中的数据在两个方向上流动时会在阶段“碰撞”。因此,在预测和反向传播之间会发生流水线冲洗(flush)。如果管理不当,冲洗可能会严重影响性能。上图展示了一个流水线并行化模型。请注意,很大一部分时间都花在了“气泡”期,即必须完全冲洗流水线的时间。

主要的流水线并行技术如下。

GPipe[25]建议在保持加速器计数不变的同时增加微批次的数量,这样流水线可以更长时间保持满状态。这不会消除冲洗,但会提高整体效率。然而,这种方法将需要更多的内存来存储很多具备检查点的微批次激活。DAP-PLE[17]提出了一种可替代的流水线调度,该调度可以保持GPipe的收敛行为,但空闲时间较少。不幸的是,它还同时保持更多的微批次“活跃”而大幅增加了内存成本,这使得调度对于已经突破内存边界的应用程序来说是不可行的。

异步流水线并行的形式还有另一种解决方案,以保持严格的收敛行为为代价,重新排列流水线阶段和反向传播以消除冲洗。这种序列的“解耦”将问题放松为更有效的——以影响数据消费和消费顺序为代价[21,32,39,59]。例如,PipeDream[21]提出的1F1B模式,为每个后阶段(在不同的微批次上)运行一个前阶段,保持完美的比例和利用率。但它的设计需要更仔细的分区和打包,而缓解陈旧权重更新导致的准确性下降,需要存储多个权重副本[21],从而增加了内存成本。虽然像1F1B这样的异步流水线可以很好地执行,但它并不是一个通用的解决方案——精度损失是特定情况下的,通常可能是巨大的。精度至关重要且收敛行为必须可复制(例如模型选择)的应用程序不适合异步流水线并行。

虽然模型并行性着眼于在多个GPU上执行分布内存需求,但一些系统试图利用主系统内存(DRAM),而不是在更多GPU上横向扩展。这种方法的主要动机是,虽然GPU内存有限且昂贵,但DRAM实际上更便宜且可访问。

最初的工作[23,30,34,49,56]将卸载(offloading)视为一个“交换”问题——决定何时将张量从GPU内存交换到DRAM上。大多数使用图分析算法来确定在哪里“注入”交换操作,这取决于激活、梯度或参数下一次可能在执行图(execution graph)中使用的时间。SwapAdvisor是这些交换系统中最先进的,它使用并行遗传搜索算法来分析交换算子应该放在哪里以获得最佳性能。它也是最早支持卸载参数和激活的系统之一,这对于训练十亿参数模型架构至关重要。

这些复杂的交换过程可能很难设置——SwapAdvisor的搜索算法大约需要一个小时才能完成。此外,它们很难扩展到多GPU训练,因为没有明确的方法来扩展交换注入图的技术来覆盖多GPU并行性。

ZeRO-R[46]提出了另一种方法,这是一种向DRAM动态发送激活和参数的卸载系统。这种方法“在需要时卸载”,而不是预先计划卸载。设计的不规则性可能会引入内存碎片等问题,但与基于图的设计相比,它增加了很大的灵活性。在后来的版本中,ZeRO Infinity[47]将其扩展到卸载到非易失性快速存储(NVMe)/磁盘存储,实现进一步的可扩展性。

Hydra[37]选择了“独立块”策略,将模型体系结构划分为子模型(如模型并行)然后可以在DRAM和GPU存储器之间自由地溢出。可以将其与RDBMS中的溢出进行类比,在RDBMS中,独立的数据块可以向下发送到较低级别的内存。与其他溢出系统不同,Hydra的执行模式与模型并行性相同,并完全分离每个模型片的执行。它仍然试图重叠通信和计算,但忽略了其他CPU卸载技术所探索的细粒度张量卸载的复杂性。这种泛化使其不太适合单GPU执行,但使其更适合与多GPU并行化技术混合。

如图所示:Hydra的溢出策略只是简单地提升和降级进出GPU内存的模型并行分片。其他溢流设计,如ZeRO Offload使用的,虽然结构不太严格,但也类似。

大型深度学习训练の并行分布式系统_并行性_06

L2L[45]使用了类似于Hydra的设计,但在其分片方法上受到了更多限制。它专门针对Transformer架构,并将自注意块(标准Transformer运算器)与专门为其目标模型类选择的启发式方法进行交换。这使它能够在Transformer架构上表现出色,但无法实现Hydra的灵活性或ZeRO-R的动态通用性。

请注意,这些技术通常用于深度大模型分布其内存需求,因为它们在执行中都利用了某种次序。一个非常宽的算子(例如嵌入表)如果没有性能的大幅降低就无法串行化,也不容易在DRAM和GPU内存中溢出。宽算子在混合设备执行的唯一选项是串行化并行算子(在表的情况下即索引查找),并将一系列操作重写为一个深度,而不是宽阔的模型,或者在CPU上实际执行宽阔算子。

有些系统更甚,实际上是在CPU上执行操作。通常,最好完全使用GPU或TPU计算来运行模型,因为大多数DL操作符在支持高度并行的加速器上运行得更快。然而,通过卸载,数据无论如何都会在CPU上——因此,GPU操作并行地执行CPU操作不应增加开销。

ZeRO[48]提出在GPU执行期间在CPU上运行参数更新,特别是针对流行的Adam优化器[28]。Adam优化器保存一些状态参数(通常是32位),需要在32位参数上运行以避免精度下降。不幸的是,这阻止了用户为了减少内存需求而部署16位表示的工作。Adam优化器的ZeRO版本在DRAM上保持32位版本的参数,在GPU上保持低精度的16位版本,消耗更少的内存。在执行过程中,系统将梯度和优化器状态溢出到DRAM上,然后使用CPU处理对32位参数运行参数更新。在与CPU-GPU通信和GPU计算重叠的第二步骤中,更新被传播到16位参数。

混合CPU-GPU计算在非常大的推荐模型中也很常见。嵌入表是非常广泛的内存密集型算子,通常会输入一些较小的DNN进行进一步处理。如果没有任何优化,嵌入表的庞大规模将迫使只执行CPU[9]。或者,用户可以将嵌入表放置在CPU上,而DNN位于GPU内存中,并享受GPU加速的好处。一些工作,如Hotline[10]尝试和流水线数据通过模型,从基于CPU的嵌入表到GPU加速的DNN。他们证明,这种混合计算方法甚至可以比宽度方向的多GPU模型并行更快,因为它消除了all-to-all通信步骤的需求。

并行化技术可以以不同的方式进行组合。各种系统试图将各种“基本”并行方法(如数据并行、模型并行)的优点结合起来,为用户提供更高的性能和可扩展性。混合并行技术可以分为两大类——“真正的”混合,从底层开始集成并行技术,以及自上而下的混合,在不同的执行阶段选择不同的策略。

接地式混合

传统上,从一开始就将模型并行性与其他技术相结合是一项具有挑战性的任务。模型并行性提高了GPU对标准执行的要求,这可能会使与基于复制或多实例的并行技术(如数据并行性、任务并行性)的组合变得不切实际,因为它们进一步扩大了模型并行性的设备要求。

为了解决这个问题,Hydra[37]建议使用溢出技术来减少可扩展模型并行训练所需的GPU数量,然后在顶部应用任务并行性一层来支持高效的多模型训练。然后,Hydra系统利用模型并行性的分段特性,实现混合的“细粒度并行”日程,可以优于标准的任务并行性和模型并行。如图对此进行了说明。目前,Hydra是唯一一个明确针对大模型设置多模型的系统,但随着从业者努力解决模型选择和多用户集群管理的成本,这一领域的重要性可能会增加。

大型深度学习训练の并行分布式系统_并行性_07

最初由ZeRO[46]引入的完全分片数据并行性(FSDP,Fully Sharded Data Parallelism)提供了模型并行性和数据并行性的混合。与Hydra不同的是,Hydra仍然以模型并行分片的方式执行,FSDP只使用模型并行性将模型分布在数据并行的实例上,每个数据并行克隆都持有一个层组的部分参数集。当执行一个层组时,FSDP运行一个all-together步骤,在每个数据并行实例上生成完整的、未分片的层组。然后层组以纯数据并行方式执行。在执行该层之后,可以立即对其进行重新丢弃,重新分配内存占用空间。反向传播也采用了类似的方法。

在FSDP中,每个加速器的内存需求,减少到单层的最小footprint加上其余部分的分区内存需求。将单层需求分解为一个常数因子,可以将其表示为 ( / )减少,其中 是原始模型内存占用, 是数据并行实例的数量。这使得用户能够同时受益于数据并行性的性能和模型并行性的可扩展性。请注意,这确实增加了大量的通信开销——对每一层都运行all-gather——而且与基于溢出的技术不同,这仍然需要横向扩展以实现可扩展性。

ZeRO Offload[48]提出将FSDP与每个加速器溢出相结合,卸载在近期不会使用的分片层参数。这提供了更好的可扩展性,但通过CPU-GPU通信引入了更多的通信开销。ZeRO的工作是将通信与计算重叠,但一些速度减慢通常是不可避免的。分析表明,FSDP比标准数据并行性慢(尽管更具可扩展性,并且能够运行更大的模型)。FSDP的支持者声称,用户可以利用其更高的可扩展性来增加批次大小(从而使执行时间与分布数据并行DDP性能保持一致),但批次大小会影响准确性收敛行为。为了获得更好的FSDP性能而规模化批次大小,可能会导致与异步流水线相同的问题(尽管不那么极端)。3D并行性将FSDP与流水线并行性和张量并行性相结合,利用可扩展的数据并行性以及并行的深度和宽度分片执行操作。通常采取的形式,是在模型的某些部分应用FSDP,在另一个部分应用流水线,在更适合宽度分片的另一个分段中应用张量并行。3D并行通常需要基于模型架构进行大量定制——它不能像Hydra或FSDP那样开箱即用。也就是说,它已经使用Megatron[53]等系统成功地应用于许多非常大规模的模型的训练,如Megatron-LM[5]和BLOOM[12]。未来,将3D并行混合与Hydra的分片任务并行相结合,一种新的“4D并行”成为可能。

策略发现

策略发现系统试图自动化在模型中组合并行化技术的过程。最近的几个例子是FlexFlow[26]和Alpa[60]。   

FlexFlow是在开发高级DL并行技术(如流水线并行、FSDP和分片任务并行)之前构建的,它只探索数据、张量和模型并行,主要针对卷积神经网络。FlexFlow构建了一个设备拓扑图,将加速器建模为节点,将互连(例如NVLink、PCIe、Infiniband网络)建模为边缘。这允许它产生混合并行执行策略,该策略考虑了给定设备配置中边缘之间的数据移动成本。它使用模拟器来评估不同的划分策略,使用试点通道(pilot passes)来建模运营商运行时间,并基于边缘带宽进行理论计算来建模通信开销。使用模拟器作为启示(oracle),它评估了划分算子的不同方法。请注意,这种基于“分割”的并行表示不能支持在不同任务上利用独立执行的并行化技术(例如任务并行、流水线并行),尽管它可能支持FSDP。此外,它没有明确说明内存的可扩展性或在特定配置中设备内存耗尽的可能性[11]。

Alpa更明确地考虑了内存可扩展性,并考虑了算子间的并行性(例如,模型并行性、流水线并行性),而不仅仅是像FlexFlow那样的算子内分割。它使用指令级并行(ILP)公式来确定如何设置并行化策略,然后该阶段将超过设备内存限制时修改执行规划[60]。占据更广阔的策略搜索空间,这种方法可以实现比FlexFlow更好的性能。

这些混合并行化策略非常适合静态的、非数据依赖的执行任务(例如非递归神经网络)。然而,它们不能很好地扩展到更动态的任务,如多模型训练——它们是用于训练的编译器,而不是调度器。未来的工作可以考虑弥合这一差距,构建一个动态混合并行执行器。

推荐模型的模型数据并行性

DLRM给从业者带来了独特的挑战,因为它们结合了两种不同的扩展挑战。嵌入表是非常明智的,并且保证了模型并行执行的宽度分割。顶级DNN是计算密集型的,但规模较小,并且将从数据并行性中获益最多。因此,将张量并行性应用于模型的表格,并将数据并行性应用到DNN,这种混合策略将在推荐模型上表现良好。这种方法已成为完全GPU加速DLRM训练的标准[41],尽管异构CPU-GPU执行也适用于访问GPU资源较少的用户。

混合并行DLRM训练在多个GPU上划分嵌入表,并在每个GPU上放置顶部DNN的局部副本。分片的表处理在样本维度上分片的输入,然后运行分区的all-gather来重新聚集表输出,并在批次维度上为每个数据并行副本进行分区。如图对此进行了说明。

大型深度学习训练の并行分布式系统_反向传播_08

这种方法使从业者能够从神经架构中的数据和模型并行性中受益。通信步骤是密集的,通常会带来沉重的开销[35],但并行执行的好处通常会超过这一点。

总的来说,混合并行性在适当的时候结合不同并行化策略的优点,为用户提供了高效训练模型的能力。混合并行技术,如分片任务并行和FSDP,从一开始就结合了可扩展性和效率,而策略发现和DLRM混合并行可以帮助训练模型架构,其在这个图的不同阶段具有混合并行需求。