MPress 通过存储保存算子间并行性在多GPU服务器上实现十亿规模级模型训练的民主化
MPress: Democratizing Billion-Scale Model Training on Multi-GPU Servers via Memory-Saving Inter-Operator Parallelism
这篇论文的题目是《MPress: Democratizing Billion-Scale Model Training on Multi-GPU Servers via Memory-Saving Inter-Operator Parallelism》（MPress：通过节省内存的操作间并行化民主化十亿规模模型训练）。论文来自中国科学技术大学、安徵省高性能计算实验室和美国休斯顿大学。
问题背景：
深度神经网络（DNN）模型的规模不断扩大，参数数量已经从百万级增长到数十亿级。这导致了GPU内存的极大需求，单一GPU无法满足如此大规模模型的训练需求。现有的一些内存节省技术如GPU-CPU交换、重计算等虽然可以减少部分内存消耗，但会带来额外的计算或通信开销。
论文创新点：
MPress提出了一种新的单服务器多GPU系统，旨在突破十亿规模模型训练中的GPU内存壁垒，同时最小化额外的成本。它结合了多种内存优化技术，包括：
- 操作间并行性，减少GPU间的通信开销。
- 新的D2D交换技术，利用多条高速NVLink链路将张量交换至轻载GPU，以充分利用闲置的GPU内存。
- 综合了重新计算、GPU-CPU交换和D2D交换的优势，通过智能调度优化性能和内存使用。
实验方法：
MPress被集成到现有的两种操作间并行训练系统中（PipeDream和DAPPLE）。实验使用了BERT和GPT两种模型，分别在现代GPU服务器（如DGX-1和DGX-2）上进行训练。实验结果表明，MPress在相同内存优化下训练速度大幅提升，并且能够训练比基线更大的模型。
结论：
MPress显著提升了多GPU服务器上十亿规模DNN模型的训练效率，能够在不牺牲性能的前提下，突破GPU内存瓶颈。
II. 背景与动机
由于深度神经网络（DNN）计算资源的高需求，将模型训练任务并行化到多个GPU设备上已经成为常态。主要有三种并行训练方法，每种方法对应不同的分区策略，分别是通过输入样本分区（数据并行）、通过网络结构分区（模型并行）和通过层级分区（流水线并行）。为了便于理解模型并行和流水线并行的区别，本文参考了最近的Alpa工作，将现有解决方案分类为两种正交方向：即操作内并行和操作间并行。
操作内并行依赖于一个操作符在具有多个维度的张量上工作，将张量沿某些维度进行分区，并将结果分配给多个设备。数据并行作为最简单的操作内并行，将输入张量进行分片，并将数据片段分配给GPU设备来训练共享的复制模型。与此不同，操作间并行训练将目标DNN模型划分为不相交的阶段，每个阶段对应一组连续的模型层，并映射到单独的GPU上进行计算。微批次训练数据通过这些阶段以流水线的方式进行处理。
不幸的是，所有这些并行策略在支持十亿规模的单服务器训练时都会遇到GPU内存瓶颈。我们选择操作间并行作为研究的出发点，原因如下：
首先，与其他两种方法相比，数据并行训练带来了最重的内存负担和跨GPU通信开销，因为每个GPU都会复制相同数量的模型数据，并定期交换与模型参数等大小的梯度。因此，单靠数据并行难以满足快速增长的模型规模所带来的巨大内存需求。操作内并行则通过将操作符拆分成更小的部分来工作，但它需要大量通信来收集和汇总部分结果，以触发后续的计算，而这些通信位于训练的关键路径上。
相比之下，操作间并行训练引入的通信开销最少，因为对于大规模的自然语言处理模型，只有激活值在阶段之间传输，而且这些激活值通常很小。例如，对于Bert-0.64B模型，两个主机连续阶段的GPU之间仅交换每个微批次1.5MB的数据。此外，操作间并行已受到工业界和学术界的广泛关注，许多训练系统，如PipeDream、DAPPLE、GPipe、DeepSpeed和Megatron-LM，已经集成了该技术。因此，我们专注于通过操作间并行实现十亿规模模型的高效训练。
操作间并行
图1展示了操作间并行DNN训练的工作流程。每个GPU设备（即“工人”）负责训练一个不相交的连续模型层集。每个训练数据的微批次被进一步分为多个子批次，进入整个执行流水线。工人1从第一微批次开始前向传播，并将计算传递给工人2，以启动第二阶段。同时，工人1处理第二个微批次。工人3则处理第三阶段的子批次。当工人3完成第一个微批次的前向传播计算后，立即开始相应的反向传播，并从工人3回流至工人1。
在相邻的微批次之间有两种调度执行方式，即异步和同步模式。PipeDream中使用的异步模式允许第二个微批次的前向传播与第一个微批次的反向传播并行进行。例如，第二个微批次中的第七个子批次在第四个子批次的反向传播完成后立即由工人1执行。与此相反，GPipe和DAPPLE中使用的同步模式则要求不同微批次的计算是串行的。
GPU内存消耗的问题
为了探索操作间并行训练中GPU内存的利用情况，我们使用PipeDream和DAPPLE两个代表性系统训练了两个流行的DNN模型Bert和GPT。在一台AWS EC2 p3dn.24xlarge GPU服务器上进行实验，该服务器配备了8个V100 GPU（每个GPU具有32GB内存），并设置了最大可持续的模型规模。
对于PipeDream，设置微批次大小为12时，它能够支持训练参数最多为6亿的Bert模型，再往上则会出现GPU内存不足的错误。当微批次大小缩小到2时，PipeDream可以支持2亿参数的Bert模型。在相同的硬件配置下，DAPPLE能够训练最多具有53亿参数的GPT模型，微批次大小同样为2。PipeDream和DAPPLE之间可支持模型规模的差异在于，PipeDream使用异步调度，需要保存多版本的模型数据。
内存优化及其局限性
内存平衡的分区策略能够解决GPU内存不平衡的问题。然而我们发现，采用这种分区策略付出的代价是计算时间不均衡，导致训练性能下降34%。重计算技术可以通过丢弃前向传播生成的激活值并在需要时重新计算，节省内存，并已被许多主流系统采用。然而，单靠重计算有两个主要缺点：一是它无法减少优化器状态、参数和梯度的内存消耗，而这些数据占用了大量GPU内存；二是重新执行前向传播会与反向传播竞争GPU资源，并引入额外的延迟，可能导致训练时间延长33%。
GPU-CPU交换利用了CPU内存的大容量来扩展GPU内存空间，然而，由于PCIe带宽有限，GPU-CPU交换会大幅降低训练吞吐量。例如，在PipeDream中对39%的BERT数据应用GPU-CPU交换后，训练吞吐量降低了67%。ZeRO系列通过在数据并行训练中消除数据冗余来减少GPU内存消耗，但也引入了不可忽视的跨GPU通信开销。
硬件趋势与机遇
随着硬件的快速发展，现代GPU服务器已经集成了超高速的GPU间互连。2016年，NVIDIA的P100 GPU引入了第一代NVLink，提供高达160GB/s的双向数据传输带宽，是PCIe Gen3x16带宽的近5倍。
我们提出了一种新的GPU D2D交换技术，通过NVLink将张量从高内存压力的GPU卸载到有空闲内存的GPU，从而缓解GPU内存的限制。
III. MPress 内部原理
A. 设计原则
我们提出了MPress，这是一个高效的操作间并行深度神经网络（DNN）训练系统，使用异构内存优化技术来应对著名的GPU内存墙问题。MPress采用了一种新的设备间（D2D）交换技术，利用多个高带宽的NVLink互连将模型数据从内存压力较大的GPU卸载到轻负载的GPU。由于GPU资源有限，我们只对少量生命周期较短的模型数据应用D2D交换，这些数据的生命周期过短，无法承受重新计算或GPU-CPU交换的延迟。
D2D交换的快速速度对于避免重新计算和GPU-CPU交换所带来的高额性能损失至关重要，并且允许D2D交换更好地与DNN计算重叠。通过精心映射流水线阶段到GPU设备，我们可以积极满足每个GPU的内存交换带宽需求。为了克服GPU总体交换空间有限的局限，MPress还采用了重新计算和GPU-CPU交换技术来进一步减少GPU内存消耗。
B. MPress的总体概述及工作流程
图5展示了MPress系统架构的高级视图，MPress的逻辑横跨静态和运行时部分。静态部分的任务是生成内存节省计划，该计划确定在内存压力下应用内存优化的张量候选、应用哪种优化以及何时执行相应的优化或恢复被节省的张量。
在MPress的静态部分，首先通过训练目标DNN模型，获取基础的张量大小、正向和反向计算的延迟等基本统计数据。表III展示了收集的统计数据示例。然后，规划器根据成本模型探索可能的配置，该模型比较不同优化方法的时间成本，以选择对性能影响最小的优化方法。
随后，生成的初步计划将被传递给重写器，重写器进一步修改输入的数据流图，将这些优化策略嵌入到合适的位置，以符合操作的依赖关系。我们使用模拟器根据修改后的数据流图执行一次训练迭代，并收集GPU内存节省量及引入的开销。最终，反馈信息将从模拟器传递给规划器，以确定当前的配置是否接近最佳配置，并与之前的运行结果进行比较。
需要注意的是，MPress静态部分是离线运行的，因此不会带来运行时开销。与实际训练可能需要运行数百万次迭代以实现模型收敛不同，模拟器只需要执行有限的几次训练步骤。
MPress运行时部分与操作间并行训练框架的运行时部分共同工作，主要包括执行器、内存管理器和压缩库三个关键系统组件。首先，压缩库提供三种内存优化技术的高效实现，其中包括我们提出的D2D交换。同时，运行时流程如下：执行器接收来自MPress静态部分的仪器化数据流图，并触发启用了内存压缩的操作间并行训练。普通操作将直接通过底层训练框架运行时处理，而内存优化操作则由执行器执行，以释放已使用的GPU内存，并根据需要恢复状态。
C. D2D交换
在设计D2D交换技术时，我们必须综合考虑以下两个关键因素：1）不同GPU之间的带宽和连接链路数量可能存在硬件异构性；2）内存交换需求与目标操作间并行训练任务中空闲GPU内存资源的多样性。
为此，我们引入了以下两项关键技术来优化D2D交换。
数据分条：每个GPU设备可以将张量交换到多个NVLink可达的轻负载GPU，因此我们引入了数据分条技术，将目标张量划分为若干子块，并通过不同链路并行传输。在对称NVLink拓扑结构中，GPU间完全通过同质的链路连接，我们将子块划分为等大小，并让其总数量与目标接收GPU的数量一致。
在不对称的拓扑结构中（如DGX-1架构），GPU对之间的带宽可能会有所不同。为了反映这种差异，我们进一步演进了数据分条技术，结合NVLink链路带宽对子块大小进行加权，大小与链路带宽成正比。
此外，我们还管理了一个元数据表，用于跟踪经过D2D交换的张量状态。对于每个张量，我们在执行交换操作之前记录子块数量、每个子块的大小和目标GPU设备的索引。
设备映射：我们设计了一种设备-阶段映射算法，能够将轻负载的GPU分配到高内存压力的GPU附近，以最大化这些GPU的交换带宽。这个算法首先枚举所有可能的设备映射方案，并为每个映射方案确定合适的内存分配策略，最终选择时间成本最小的映射方案。
D. 内存压缩规划
探索将D2D交换、GPU-CPU交换和重计算结合起来的最佳配置，以最大化GPU内存节省的同时最小化训练过程中引入的额外延迟是一个非常具有挑战性的问题。为了解决这一问题，我们提出了一个近似搜索算法，并基于以下几个关键观察来进行简化。
与重计算相比，两种交换方法的优势在于它们不消耗GPU的计算资源，并且能够与GPU上的前向和后向计算并行运行。当被适当地应用时，也就是当目标张量的生命周期（即从生成到再次使用之间的时间）长于这两种交换方法的时间成本时，就不会引入额外的时间延迟。
我们应该优先考虑对模型的后期层使用重计算来缓解内存限制，原因有两个：首先，反向传播首先从流水线执行的后半部分的后期层开始；其次，不可避免的额外重计算延迟会扩大早期层的生命周期，从而为应用GPU-CPU交换腾出更多空间。
选择对连续层的张量应用重计算是一个不错的选择，这可以进一步减少除了第一层操作符的输入之外的内存消耗，因为第一层操作符的输出是前一层的输入。
由于GPU的空闲内存资源有限，且D2D交换比GPU-CPU交换要快得多，为了发挥其全部潜力，我们应该只将D2D交换应用于性能至关重要的情况，以最大限度地减少GPU-CPU交换和重计算带来的额外延迟。
基于以上观察和权衡，我们引入了一个近似的搜索算法，首先积极地为适当的张量分配GPU-CPU交换和重计算优化，然后通过逐步将部分GPU-CPU交换和重计算替换为D2D交换来优化配置。具体而言，我们首先执行生命周期变量分析来计算每个张量的生命周期。接着我们构建了初步配置，将GPU-CPU交换分配给生命周期特别长的张量，同时将重计算分配给激活张量，前提是重计算引入的额外延迟小于GPU-CPU交换的成本。最后，我们将GPU-CPU交换分配给剩余的张量，以满足目标内存节省需求。
我们的算法会经历一些迭代步骤来逐步更新内存减少优化配置。在每个步骤中，我们使用模拟器运行最新的配置（仅需要执行一次训练迭代）来过滤出一组减少的张量，这些张量的优化引入了最大的额外开销。对于这些张量，我们尽量使用D2D交换来减少它们的内存占用，前提是存在空闲的GPU内存。如果新的配置比之前的配置性能更好，我们就接受它。该算法会在后续配置相较于前一个配置的性能提升不明显时终止。
E. 实现细节
我们将上述设计原则集成到开源训练系统MPress中，该系统由2000行C++和Python代码组成。为了展示其在操作间并行训练中优化GPU内存使用的通用性，我们将MPress集成到了PipeDream和DAPPLE这两个最近的操作间并行训练系统中。其后端引擎是PyTorch。我们通过将PipeDream的原始PyTorch版本从1.1升级到1.2，并启用其NCCL库以使用NVLink在阶段之间传输数据，改进了PipeDream系统。需要注意的是，MPress是通用的，可以应用于其他操作间并行训练系统，如GPipe。我们将继续使MPress适用于这些系统。
内存管理：我们的内存管理器负责为张量分配和释放GPU/CPU内存空间，并监控每个设备的内存使用情况。首先，针对GPU内存分配，管理器直接使用PyTorch中的原生GPU内存分配器。其次，在主机内存请求方面，考虑到固定内存（pinned memory）与GPU内存之间的数据传输速度快于普通可分页内存（pageable memory），我们决定使用固定内存作为交换空间。为了避免频繁分配和释放固定内存的高成本，我们进一步构建了一个独立于PyTorch运行时的主机固定内存池。
内存交换：对于D2D交换，执行器为执行交换操作的swap-in和swap-out任务分别管理两个额外的线程，这些线程使用系统启动时通过调用cudaStreamCreate创建的不同CUDA流（streams）。这种设计允许执行器在不阻塞主线程的情况下，启动张量传输任务并与DNN计算同步检查其状态。因此，GPU之间的数据移动可以与DNN计算异步进行。
IV. 评估
我们的评估旨在回答以下几个问题：
- MPress是否能有效缓解GPU内存限制，支持操作间并行训练中的大规模模型训练，同时比基线方法提供更好的训练性能？
- 独立的D2D交换、设备映射和内存压缩策略的性能影响是什么？
- 这三种内存优化技术对GPU内存节省的贡献各自是多少？
A. 实验设置
机器配置：我们在DGX-1和DGX-2 GPU服务器上进行实验以评估MPress和基线方法。DGX-1服务器是AWS EC2 p3dn.24xlarge实例，配备96个虚拟CPU、8个NVIDIA Tesla V100 GPU（每个GPU的内存为32GB，连接方式为不对称的NVLink），以及768GB的CPU内存。我们还使用了另一台提供商的DGX-2服务器，因为AWS EC2上这种高端GPU服务器的配额非常有限，我们的配额申请多次失败。这台DGX-2服务器有164个虚拟CPU，8个NVIDIA Tesla A100 GPU（每个GPU的内存为40GB，连接方式为对称的NVLink），以及948GB的CPU内存和6TB的NVMe SSD。两台机器都运行了Ubuntu 18.04、CUDA 11.7、NCCL 2.8.4、PyTorch 1.2.0等软件。
模型和数据集：我们选择了两个广泛使用的DNN模型BERT和GPT，它们都来自自然语言处理领域。我们使用SQuAD v1.1数据集训练BERT，使用维基百科数据集训练GPT。
我们在PipeDream上运行MPress，训练了不同规模的BERT模型。根据文献，我们通过调整编码器层的数量和隐藏层的大小，使BERT模型变得更深、更宽。表II展示了我们使用的BERT变体的参数规模。BERT-0.35B表示最小的BERT模型，其总GPU内存需求为108.8GB，每个阶段最大内存需求为24.77GB。显然，这一GPU内存需求可以通过我们测试的GPU服务器在没有内存优化的情况下满足。BERT-0.64B是一个中等规模的模型，其最大和最小阶段内存需求分别高于和低于每个GPU的内存容量。此外，BERT模型的参数数量分别为1.67B和4B，表示内存需求超过每个GPU容量的大模型。最后，BERT-6.2B是一个超大模型，其总内存需求是服务器GPU内存供应量的5倍。我们将微批次大小设置为12，这是文献中推荐的大小。
同样，我们使用调整后的参数对GPT及其变体进行训练，并在DAPPLE上运行MPress。如表II所示，GPT-5.3B是最小的模型，原生的DAPPLE可以训练该模型。然而，其他四个GPT配置的每阶段最大GPU内存需求为56.4-140.1GB，已超过单个GPU的容量。我们将微批次大小设置为2，这是DAPPLE文献中推荐的最小批次大小。
基线和系统配置：对于基于PipeDream的操作间并行训练，我们使用原始的PipeDream作为没有内存优化的基线操作间并行训练系统。我们还部署了两种系统作为基线，分别是启用GPU-CPU交换的系统和启用重计算的系统，通过启用PipeDream中的内存交换和重计算优化实现内存压缩。此外，我们运行了两种MPress变体，其中一种仅使用D2D交换，另一种结合了三种内存优化技术。
对于基于DAPPLE的MPress，我们将原始DAPPLE作为自然基线，并启用了高性能重计算的DAPPLE+Recomp。此外，我们还运行了两个最先进的训练系统，ZeRO-Offload和ZeRO-Infinity，它们可以通过数据并行进行大规模模型训练。
我们使用BERT和GPT模型推荐的计算平衡阶段分区策略。此外，重计算基线系统按照文献的建议丢弃特定的张量。最后，MPress变体采用我们之前算法生成的设备映射和内存节省计划。
评估指标：我们测量了每秒处理的样本总数、每秒浮点运算次数（FLOPS）、D2D交换、GPU-CPU交换和重计算的时间成本，以及三种方法各自的内存减少情况。类似于其他现有的FLOPS计算工具或方法，我们通过测量模型前向传播的FLOPS，并估算相应的反向传播的FLOPS为前向传播的两倍。
B. MPress在PipeDream上的性能
图7对比了在五种不同系统配置下，使用不同规模的BERT变体进行操作间并行训练时的性能（用TFLOPS表示）。我们通过以下几个方面分析了结果：
小规模模型：我们首先分析了参数规模为0.35B的最小BERT模型。五个系统都成功训练了该模型，且报告的性能数据相同。这是因为普通的操作间并行训练已经能够满足该模型的GPU内存需求，因此不需要触发任何内存压缩优化。
中等规模模型：当模型从BERT-0.35B扩大到BERT-0.64B时，PipeDream出现了GPU内存不足的错误，而其他四个系统成功执行了训练任务。GPU-CPU交换的性能最差，因为PCIe带宽的限制导致交换操作非常慢，从而延迟了相应的DNN计算。重计算比GPU-CPU交换快了143.4%，因为重新计算丢弃激活值的前向传播引入的延迟通常远低于GPU-CPU交换。在这种情况下，两种MPress变体表现最好且性能相同。原因是D2D交换足够缓解内存限制，并且比重计算更快，因此MPress没有使用其他两种优化。
大规模模型：当模型扩大到1.67B时，单独的D2D交换无法满足需求，因为在高内存压力下，空闲的GPU内存不足以容纳从负载较大的GPU卸载的张量。与GPU-CPU交换相比，重计算提高了125.4%的性能，但比MPress慢19.5%。此时，MPress结合了三种优化策略，GPU-CPU交换和重计算为D2D交换腾出了更多空间，提升了整体性能。
有趣的是，重计算在BERT-4B及以上模型中失败了。这是预期中的结果，因为重计算只能减少前向传播生成的激活值的内存消耗，无法处理剩余的模型数据，如参数和梯度，这些数据占用了更多的内存资源。相比之下，GPU-CPU交换依然可行，因为它可以应用于所有类型的模型数据，只要有足够的主机内存空间。然而，通过优先使用更快的D2D交换和尽可能使用重计算，MPress相比GPU-CPU交换实现了1.8倍的训练性能提升。
超大规模模型：最后，我们测试了参数规模为6.2B的BERT模型。在这种情况下，只有GPU-CPU交换和MPress能够成功执行训练任务，而其他系统因内存不足而失败。尽管GPU-CPU交换可以支持相同的超大模型，但MPress通过设备映射优化和三种优化策略的组合，训练性能提升了3.1倍。此外，与最先进的重计算系统相比，MPress支持的模型规模提高了2.7倍。
C. MPress在DAPPLE上的性能
我们使用不同参数规模的GPT模型在DGX-1和DGX-2 GPU服务器上对MPress进行了测试。这里，我们将MPress与三个强大的基线系统进行了对比，这些系统应用了最先进的内存节省技术：
- DAPPLE+Recomp：启用了高性能重计算。
- ZeRO-Offload：将优化器状态从GPU卸载到CPU。
- ZeRO-Infinity：通过GPU-CPU交换并进一步扩展到NVMe设备。
值得注意的是，ZeRO系列是DeepSpeed框架的一部分。我们在一台高端GPU服务器上运行测试，该服务器具有与前面实验相同的GPU配置，但拥有更大的CPU内存和额外的NVMe SSD存储空间。我们无法在上述Amazon EC2实例上进行这组实验，因为ZeRO-Infinity需要大量的CPU内存来进行初始化，并且需要额外的具有高I/O带宽的存储空间进行张量交换。
图8a总结了在DGX-1 GPU服务器上的性能对比情况。DAPPLE无法支持超过5.3B参数规模的模型，因为其最大每GPU内存需求超过了32GB。与此不同，启用重计算的DAPPLE能够成功训练最多10.3B参数的模型，但其性能比MPress低19.2%。相比之下，ZeRO系列和MPress能够支持从5.3B到20.4B不等规模的所有训练任务。ZeRO-Infinity比ZeRO-Offload的GPU计算效率高出20.6%-23.8%，这是因为ZeRO-Offload将优化器状态卸载到CPU会在每个微批次中导致频繁的数据移动，而ZeRO-Infinity通过GPU-CPU交换取而代之。然而，MPress在模型规模上始终能够保持稳定的训练性能，这得益于D2D交换技术的使用，以及多种内存压缩优化技术的结合。此外，MPress比ZeRO-Infinity的性能提升了37.0%-40.8%。这意味着，仅靠GPU-CPU交换的方案可能可以支持非常大规模的模型训练，但需要牺牲训练速度。MPress通过进一步利用空闲的GPU内存资源，补充了它们的局限性。
图8b展示了在DGX-2 GPU服务器上的类似趋势，性能提升幅度比DGX-1服务器翻了一倍多，原因是DGX-2上A100 GPU的计算密度高于DGX-1上的V100。此外，重计算能够支持模型规模达到15.4B，比在DGX-1上支持的规模更大，因为A100 GPU的内存为40GB，而V100的内存为32GB。与MPress相同，ZeRO系列可以支持参数规模最多达到25.5B的模型，但其训练性能分别比MPress降低了30.4%-44.8%和23.2%-70.0%。有趣的是，在大模型上，ZeRO-Infinity的表现不如ZeRO-Offload。这是因为我们租用的DGX-2服务器的SSD存储带宽显著低于DGX-1。然而，找到可公开访问的、具备可扩展GPU计算能力和存储容量的高端GPU服务器几乎不可能。需要注意的是，即使具有足够的SSD带宽，ZeRO-Infinity也不应显著优于ZeRO-Offload，这一点已被ZeRO-Infinity的原始论文和图8a中的结果验证。
PipeDream和DAPPLE之间的结果差异：有趣的是，PipeDream和DAPPLE之间的模型规模和性能差距较大，这也影响了它们各自的MPress变体。模型规模的差异在前面已经解释过。至于性能差距，DAPPLE显著优于PipeDream，因为DAPPLE默认启用了FP16低精度训练功能。此外，DAPPLE比PipeDream晚了两年发布，并且吸收了来自深度学习社区的各种优化，例如更好的计算和通信重叠技术。
D. 敏感性分析
为了更好地理解MPress相较于基线系统的优势，我们进行了敏感性分析，以探讨设备映射、独立的三种优化方法（即GPU-CPU交换、重计算、D2D交换）的开销比较，以及内存压缩计划的选择。
D2D交换优化的影响
图9总结了通过逐步添加设备映射和数据分条优化后MPress的性能提升情况，数据是在前面的D2D交换默认设置基础上进行的。结果已被归一化为默认设置，其中阶段按DAPPLE的建议映射到设备，且D2D交换已启用但未进行数据分条。
对于DGX-1服务器，设备映射优化提高了默认设置下的性能17.4%，而数据分条优化使性能进一步提高了33.3%。这是因为前者优化使交换操作能够通过可达的NVLink链路传输数据，而后者则进一步最大化了数据传输带宽。相比之下，在DGX-2服务器上，设备映射对性能没有影响，这在预期之中，因为对称的全互联NVLink连接使得每个GPU拥有相同数量的NVLink连接。然而，数据分条使MPress性能比默认设置提高了11.0%，这是由于数据分条技术利用了多条NVLink的聚合带宽，加速了模型数据的交换。
我们还评估了运行设备映射算法的时间开销。首先，我们设计了一个极端案例，其复杂性远高于图7和图8中的所有实验，以此来测试我们的搜索算法。即使在单线程实现下，MPress也能在47秒内找到最优映射。此外，在我们的所有评估案例中，设备映射搜索仅需几秒钟即可完成。因此，我们得出结论：我们的搜索算法不会带来高额的开销。如果有必要，我们还可以将其优化为多线程版本。
内存压缩方法的成本比较
表III报告了在BERT和GPT模型中，针对不同大小张量，三种内存优化方法的时间成本。显然，三种方法的性能差异显著，这在决定如何组合这些方法时起到了关键作用。
首先，在BERT模型中，张量t1具有最长的生命周期，足以涵盖GPU-CPU交换和D2D交换的时间成本。因此，MPress会优先选择GPU-CPU交换，因为它的成本可以被隐藏起来，而D2D交换则可以用于其他对性能更关键的任务。对于张量t2，GPU-CPU交换和重计算都会带来额外开销，并会延迟操作间并行训练，因为前者的时间成本超过了t2的生命周期，而重计算引入了3毫秒的额外前向计算。因此，MPress会选择D2D交换，因为它的时间成本只有3毫秒，且能够被轻松隐藏。最后，对于t3，MPress会舍弃GPU-CPU交换，因为它过于缓慢，并优先选择重计算而不是D2D交换，尽管两者的额外开销相同（4毫秒），但重计算不会消耗有限的GPU空闲内存，这些内存或许能更好地用于其他张量。
其次，GPT模型中的情况与BERT类似。对于张量t4，MPress会分配GPU-CPU交换，因为其生命周期较长。而对于t5，MPress优先选择D2D交换，因为它的性能比GPU-CPU交换快了7.6倍。最后，针对t6，MPress选择重计算，因为在三种方法中，重计算的额外开销最小。
MPress选择的策略
表IV展示了MPress在高GPU内存压力下，为四个操作间并行训练任务（BERT-1.67B、BERT-6.2B、GPT-10.3B、GPT-20.4B）生成的最佳策略。我们还报告了每种优化方法对GPU内存减少的百分比贡献。在三种方法中，GPU-CPU交换对内存节省的贡献为0%-42.2%，而重计算的贡献最大，达到51.2%-90.6%。D2D交换的内存节省量为3.9%-23.4%，虽然通常少于重计算和GPU-CPU交换，但在避免GPU-CPU交换和重计算带来的额外开销或GPU资源争用方面起到了关键作用。
对于BERT-1.67B，MPress没有使用GPU-CPU交换，因为其时间成本过高。相反，D2D交换节省了23.4%的GPU内存，并被用于将早期阶段的张量（阶段0-3）传输到后续阶段（阶段4-7）。重计算节省了阶段0-6的GPU内存，共减少了76.6%的内存。D2D交换和重计算的组合带来了最佳性能，D2D交换提升了19.5%的速度（见图7）。GPT-20.4B的表现则有所不同。首先，MPress选择对阶段0-7的242GB模型数据应用GPU-CPU交换，实现了42.2%的内存节省。其次，重计算减少了51.2%的GPU内存占用。最后，D2D交换节省了38GB的GPU内存，是这四个训练任务中D2D交换节省内存最多的一次。
V. 硬件见解
尽管GPU的HBM（高带宽内存）提供了极高的带宽，但要满足未来不断增长的模型需求仍然是一个挑战，因为HBM成本昂贵（例如，最新的GPU只有80GB的HBM）。相比之下，CPU内存更便宜，扩展性更好。两者的价格和容量差异主要是由不同的制造工艺导致的。实际上，NVIDIA的Grace-Hopper架构已经为每个GPU支持专用的CPU侧内存，并通过高带宽（NVLink C2C）连接到GPU。因此，MPress展示了此类架构的潜在优势及其在低硬件成本下解决内存墙问题的示例应用。
为了理解MPress在这种新架构下的优势，我们进行了简单的分析，预测了其理想的性能。我们发现，即使每个GPU拥有96GB（HBM）+ 512GB（Grace CPU内存）的设备内存，训练参数规模为175B的GPT-3模型依然会遇到OOM（内存不足）问题。然而，MPress可以解决这个问题。为了完全隐藏GPU-CPU交换的成本，我们预计每个GPU的PCIe带宽需要超过140GB/s，是Grace-Hopper目前64GB/s带宽的两倍多。因此，MPress中的D2D交换技术在这种情况下仍然有效，既可以节省重计算浪费的25%的资源，又可以避免通过超级芯片与PCIe连接的内存之间的交换所带来的13%的训练时间延长。此外，广泛采用新硬件技术需要时间。例如，AWS EC2上最新的GPU实例仍然使用DGX-2 A100，而且供应量非常有限。因此，MPress可以部分弥补当前硬件的局限性。
最后，MPress为重新思考内存架构提供了帮助。考虑到张量的生命周期以及各种内存技术在成本、容量和带宽方面的权衡，扩展内存层次结构是有益的。最快的内存层存储的是计算立即需要的数据，而速度较慢的层存储生命周期较长的数据。每个内存层都可以有不同的访问带宽，以进一步降低成本。在这种情况下，MPress的规划算法可以扩展，以智能地将不同模型张量分配到合适的内存层。
VI. 相关工作
我们将本文的相关工作分为以下几类进行讨论：模型并行化、流水线并行、内存优化、GPU-CPU交换和深度学习加速器。
模型并行化：当模型规模超过单个设备的内存容量时，模型并行化将模型的不同部分分配给不同的设备进行计算。近年来，研究者们提出了多种形式的模型并行化。Megatron-LM、GShard等系统通过将每个层的矩阵乘法分块，在多个GPU之间分配模型参数，这种方式被称为张量并行。此外，PipeDream、GPipe和DAPPLE等系统采用了流水线并行，它将模型的不同层分配给不同的设备，以支持更大的模型。MPress与这些工作是互补的，它不仅能够利用现有的流水线并行，还通过进一步减少内存消耗来训练更大规模的模型。
流水线并行：流水线并行是一种常见的训练大规模模型的技术，它通过将模型的层分成多个阶段，分别在不同的GPU上运行。GPipe首先引入了这种技术，随后PipeDream通过引入异步调度来改进GPipe的性能。DAPPLE进一步在流水线并行中添加了高效的调度和通信优化。与这些系统不同，MPress不局限于特定的调度策略，而是将D2D交换、GPU-CPU交换和重计算相结合，优化操作间并行的内存使用。
内存优化：为了应对GPU内存的限制，研究者们提出了多种内存优化技术。重计算（recomputation）是最常见的技术，它通过丢弃中间激活值并在反向传播中重新计算它们，来节省内存。近年来，ZeRO系列通过分片优化器状态、梯度和参数，进一步减少了内存消耗。ZeRO-Offload和ZeRO-Infinity利用CPU内存和NVMe存储来扩展模型规模。然而，MPress通过D2D交换减少了GPU内存需求，并且只在必要时使用GPU-CPU交换，从而避免了过多的性能损失。
GPU-CPU交换：一些研究工作探索了在GPU内存不足时，将张量交换到CPU内存的可能性。vDNN和SuperNeurons等系统在运行时根据内存使用情况，决定哪些张量需要交换到CPU。Tesseract进一步结合了流水线并行和GPU-CPU交换。然而，这些系统的性能往往受限于PCIe带宽，MPress通过优先使用NVLink进行D2D交换，避免了GPU-CPU交换的高延迟，进而提高了训练效率。
深度学习加速器：近年来，Google TPU和Graphcore IPU等专用加速器的兴起，提供了替代传统GPU的硬件选择。它们通过自定义的内存架构和高效的并行计算，提升了训练大规模模型的效率。然而，这些硬件需要定制的软件支持，且不一定适用于现有的GPU计算框架。相比之下，MPress能够在现有的GPU硬件上运行，并且无需对硬件进行额外的修改或优化。
总之，MPress与上述工作不同，它通过结合多种内存优化技术，提供了一种高效的解决方案，能够在现有的多GPU系统上训练超大规模的模型。
VII. 结论
我们提出了MPress，这是一种新的单服务器多GPU系统，旨在通过节省内存的操作间并行化方法，解决十亿规模模型训练中的GPU内存壁垒问题。MPress结合了三种内存优化技术：重计算、GPU-CPU交换和我们新提出的D2D交换。通过综合运用这些技术，MPress能够高效地减少GPU内存消耗，同时保持较高的训练性能。
我们通过在PipeDream和DAPPLE系统中集成MPress，评估了它在训练BERT和GPT等大规模DNN模型时的表现。实验结果表明，MPress显著提高了训练吞吐量，能够在同样的内存优化下，比现有的系统支持更大规模的模型训练。此外，MPress为未来高带宽异构计算环境中的大规模模型训练提供了设计思路。我们相信，MPress展示了在现有硬件约束下突破模型训练规模限制的潜力，为操作间并行训练和GPU内存优化提供了新的方向。