ProTrain: Efficient LLM Training via Adaptive Memory Management

文章标题：ProTrain：通过自适应内存管理实现高效的 LLM 训练
作者/机构：
- Hanmei Yang, Jin Zhou, Hui Guan, Tongping Liu (马萨诸塞大学阿默斯特分校)
- Yao Fu, Xiaoqun Wang, Ramine Roane (Advanced Micro Devices, Inc.)

A1 主要贡献

本文旨在解决训练大型语言模型（LLM）时遇到的极端内存消耗问题。现有的解决方案（如ZeRO-Offload）虽然通过结合CPU和GPU内存降低了训练门槛，但它们通常提供的是粗粒度的内存管理，并且需要专家进行繁琐的配置调优，导致硬件利用率和性能不佳。

核心问题与研究目标：
1. 粗粒度控制：现有框架（如FSDP）对内存管理技术的控制是二元化的，例如，所有模型状态要么完全卸载到CPU，要么完全保留在GPU；所有Transformer块要么全部使用梯度检查点，要么都不使用。
2. 手动配置繁琐：用户需要手动设置大量复杂的配置项，如ZeRO的优化阶段、卸载选项、参数获取和集合通信的阈值等，配置不当会导致内存溢出或效率下降。

为了解决这些问题，本文提出了一个名为 ProTrain 的自适应内存管理训练系统，它能根据LLM的结构和可用内存资源智能地协调内存、计算和IO，无需任何手动干预，从而在不影响训练精度的前提下，平衡内存使用和性能。

创新点与核心贡献：
1. 自适应内存管理：ProTrain能够自动决定是否使用卸载或梯度检查点，并确定需要卸载的模型状态和激活值的数量，以及应用梯度检查点的Transformer块数量，整个过程无需用户输入。
2. 计算与IO的深度重叠：
* 计算：ProTrain将前向/后向计算保留在GPU上以保证效率，同时动态决定在CPU和GPU上执行参数更新的比例。此外，它将CPU上的参数更新与GPU上的后向计算并行执行，以隐藏CPU更新的开销。
* IO：通过在前向/后向计算期间主动预取未来参数、将梯度卸载与后向计算并行化、以及仅在开销可被计算隐藏时交换激活值，ProTrain实现了IO通信与计算的重叠。

1. 核心技术组件：
   • 基于块的模型状态管理（Chunk-Based Model State Management）：将模型状态组织成统一大小的块（chunk），并引入持久化块和块缓冲区，以最小化不必要的数据复制和动态内存分配。
   • 逐块的激活管理（Block-Wise Activation Management）：在Transformer块的层级上处理激活，根据需要对每个块执行交换（swapping）或梯度检查点（gradient checkpointing）。通过交错式交换和检查点策略，确保在降低内存使用的同时不损害性能。
   • 内存感知的运行时分析器（Memory-Aware Runtime Profiler）：精确估计运行时开销和内存需求，为自适应管理提供决策依据，指导计算、内存和IO的有效编排。

实验结果表明，与当前最先进的训练系统（SOTA）相比，ProTrain的训练吞吐量提高了1.43倍至2.71倍，并且能够训练更大规模的模型。

A3 背景知识

2.1 深度学习模型训练

深度学习模型训练过程。训练深度学习模型涉及一个在多个迭代和周期中重复的三阶段过程。这些阶段包括：前向传播（FWD），将一批训练样本传递给模型以计算损失；后向传播（BWD），通过反向传播损失来计算梯度；以及参数更新（OPTIM），使用梯度通过优化器更新模型参数。对于大型模型的训练，通常采用混合精度训练【索引29, Mixed precision training, 2017, arXiv preprint arXiv:1710.03740】，即FWD和BWD使用较低精度的数据类型，而OPTIM则保持较高精度以确保准确性。

训练中的内存消耗。训练期间的内存消耗主要来自两个来源：模型状态和残差状态。模型状态包括参数、梯度和优化器状态，而残差状态则由激活值和临时缓冲区组成。FWD和BWD阶段的计算复杂度随模型大小和批量大小而扩展，因此由于其密集的计算需求，必须在GPU上执行。相比之下，OPTIM阶段涉及的操作较为简单，可以有效地卸载到CPU【索引40, Zero-offload: Democratizing billion-scale model training, 2021, 2021 USENIX Annual Technical Conference (USENIX ATC 21)】，通过在CPU上分配内存密集型的优化器状态，可以显著节省GPU内存。

2.2 ZeRO 技术

ZeRO技术介绍。零冗余优化器（ZeRO）【索引37, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】是一种用于大规模分布式模型训练的内存优化技术。它通过在多个GPU之间分配模型状态来增强传统的数据并行性，从而缓解内存瓶颈。ZeRO分三个阶段运行：ZeRO-1在GPU之间划分优化器状态；ZeRO-2通过同时分配梯度来扩展这一点；ZeRO-3进一步划分了参数，这些参数在进行前向/后向计算之前需要被收集起来。

ZeRO在主流框架中的应用。ZeRO技术已被集成到最先进的框架中，如DeepSpeed【索引39, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】、FSDP【索引51, Pytorch fsdp: experiences on scaling fully sharded data parallel, 2023, arXiv preprint arXiv:2304.11277】和Colossal-AI【索引24, Colossal-ai: A unified deep learning system for large-scale parallel training, 2023, Proceedings of the 52nd International Conference on Parallel Processing】，每个框架在参数组织上有所不同，以优化带宽利用率。与需要手动配置参数分组的DeepSpeed和FSDP不同，Colossal-AI根据模型规模自动将参数分组为块并动态调整其大小。这种基于块的方法，灵感来源于PatrickStar【索引9, Parallel training of pre-trained models via chunk-based dynamic memory management, 2022, IEEE Transactions on Parallel and Distributed Systems】，也被ProTrain采用，因为它不仅提高了带宽效率，还支持了更高级的模型状态管理，如3.1节所述。

A2 方法细节

ProTrain由三个核心组件构成：一个包含用于模型状态的基于块的管理（§ 3.1）和用于激活的逐块管理（§ 3.2）的细粒度内存管理模块，一个收集模型运行时和内存信息的内存感知运行时分析器（§ 3.3），以及一个识别并应用最优配置以提升性能和硬件利用率的自适应内存管理模块（§ 3.4）。

3.1 基于块的模型状态管理

基于分块管理的关键操作。ProTrain提出了一种新的基于块的管理方法，将模型状态组织成统一大小的块。图1概述了涉及块的五个关键操作。在分布式训练中，每个块被均匀划分并分布在多个GPU上，每个GPU持有一个分片。在每次前向或后向传播之前，会执行一次all-gather操作，从所有GPU收集必要的分片，将它们组装成一个完整的参数块以用于接下来的计算（❷）。在后向传播中，ProTrain重用参数块来存储计算出的梯度，从而减少内存使用。一旦梯度完全替换了块内的参数，就会执行一次reduce-scatter操作来平均所有GPU上的梯度（❸）。当GPU内存不足时，ProTrain会将模型状态卸载到CPU。因此，参数块必须在收集前上传到GPU（❶），而梯度块在梯度归约后卸载到CPU（❹）。最后，参数更新根据模型状态的位置在GPU或CPU上进行（❺）。

持久化和非持久化块的双块系统。像FSDP那样完全卸载所有参数，通常会导致GPU内存利用效率低下和高昂的数据传输开销。为了解决这个问题，ProTrain引入了一个由持久化和非持久化块组成的双块系统。持久化块永久地将模型状态存储在GPU上，消除了数据传输并支持直接在GPU上进行参数更新。相反，非持久化块存储在CPU内存中，其参数在计算时上传到GPU，梯度则卸载回CPU进行参数更新。尽管CPU更新通常比GPU更新慢，但ProTrain利用了在GPU进行后向计算时CPU通常处于空闲状态的观察，并发地执行CPU参数更新，从而有效地隐藏了高昂的CPU更新开销，并提高了整体硬件利用率。

持久化块优先策略与块缓冲区。基于双块系统，ProTrain设计了一种持久化块优先策略，优先为早期执行的层使用持久化块。该策略不仅通过消除参数上传降低了启动延迟，还通过启用GPU参数更新提高了效率，因为这些早期执行的层在后向传播中是最后处理的，留下的计算量较少，难以与CPU更新重叠。此外，ProTrain为非持久化块引入了预分配的块缓冲区，以避免频繁的内存分配和释放。至少需要三个缓冲区：一个用于预取参数，一个用于参与计算，一个用于卸载梯度。这些缓冲区还充当缓存，使得在前向传播中加载的参数可以在后向传播开始时被重用，从而消除了数据传输的需要。

优化的块组织方式。此外，ProDtrain根据模型参数在运行时的执行顺序来组织它们，而不是像Colossal-AI中使用的初始化顺序。这种安排减少了因内存限制而频繁加载和卸载块的需求，从而最大限度地减少了可能降低性能的不必要的来回访问。对于Transformer模型，ProDtrain将来自同一块的参数分组到一个块中，这最大限度地减少了内存访问，尤其是在使用梯度检查点时，因为梯度检查点需要在后向传播过程中反向重新访问参数。通过优化块的组织方式，ProDtrain简化了性能和内存使用之间的平衡问题，将其归结为确定持久化块和块缓冲区的最佳数量，为第3.4节中详述的自适应内存管理铺平了道路。

3.2 逐块的激活管理

逐块激活管理的类型。ProTrain提出了一种新颖的逐块激活管理方法，它无缝集成了激活交换和梯度检查点，以在不影响性能的情况下优化内存使用。在ProTrain中，每个Transformer块的激活属于以下三种类型之一：交换、检查点或两者皆非，如图2所示。交换表示该块将被换出，而检查点意味着将通过保存块的输入张量来重新计算整个块。尽管之前已经提出了交换和梯度检查点的集成方案【索引32, Capuchin: Tensor-based gpu memory management for deep learning, 2020, ASPLOS ’20】【索引2, Efficient combination of rematerialization and offloading for training dnns, 2021, Advances in Neural Information Processing Systems】，ProTrain的激活管理有以下显著不同。

块级管理与现有机制的集成。首先，ProTrain在块级别管理激活，而不是张量粒度。尽管在张量级别管理提供了更大的灵活性，但由于Transformer的可预测执行模式，这种灵活性并非必需。现有的梯度检查点机制通常对整个Transformer块进行检查点，这需要重新计算该块内的所有激活。ProTrain的方法与这些机制更无缝地集成，但不同之处在于它能够选择性地将梯度检查点应用于块的特定部分，而不是统一应用于所有块，从而提供了更细粒度和高效的激活管理。

交错式交换与检查点策略。其次，ProTrain设计了一种交错式交换和检查点策略以提高训练效率。通常，ProTrain对第一个Transformer块使用激活交换，并以交错方式对后续块应用梯度检查点。交换间隔经过精心选择，以使计算时间与换出一个块所需的时间相匹配。如果有足够的内存同时容纳多个块，这些块将在没有交换或检查点的情况下处理。如图2所示，这种交错策略最小化了峰值内存使用。在前向传播期间，由于交换间隔的存在，一次只有一个交换块的激活会累积。在后向传播期间，首先处理没有优化的块，消耗激活并为后续的检查点和交换释放内存。ProTrain还通过平衡早期预取以隐藏通信延迟与内存溢出风险来优化激活预取。只有在确认有足够内存时才开始预取，并将激活分组为可管理的块以提高带宽利用率。

3.3 内存感知的运行时分析

模型级运行时分析。ProTrain设计了一个内存感知的运行时分析器，即使在内存容量有限的情况下，也能提供对内存需求的精确洞察。该分析器采用模型级运行时分析，以解决静态分析【索引15, Elixir: Train a large language model on a small gpu cluster, 2022, arXiv preprint arXiv:2212.05339】和逐层运行时分析【索引2, Efficient combination of rematerialization and offloading for training dnns, 2021, Advances in Neural Information Processing Systems】中常见的内存需求低估问题，这些方法没有考虑到不可钩取（unhookable）的操作和临时缓冲区。具体来说，ProTrain为每个可钩取的操作注册钩子，并分析连续可钩取操作之间的内存变化，以推断不可钩取操作的内存使用情况。此外，通过在每个可钩取操作之前和之后注册钩子，ProTrain监控峰值内存变化，以了解临时缓冲区的内存使用情况。为了在不出现内存不足问题的情况下分析非常大的模型，ProTrain设计了一种“丢弃-再生”方法，该方法仅将当前层的数据保留在GPU内存中，丢弃其他数据（如参数、梯度、激活），并根据需要重新生成它们。虽然这种方法节省了内存，但它也改变了总内存使用量。ProTrain的分析器使用钩子来监控操作前和操作期间的内存变化和峰值使用情况，使其具备内存感知能力。通过将这些观察到的内存波动与已知丢弃张量的大小相结合，它可以精确预测各种内存管理技术下的内存需求，详见第3.4节。

性能指标分析。我们的分析器还跟踪每个操作的执行时间。与内存分析类似，我们利用可钩取操作之间的时间间隔来估计不可钩取操作的执行时间。此外，分析器还收集详细的硬件指标，包括在独立和重叠场景下的内存传输带宽和集体通信操作的持续时间。这种全面的数据收集准确反映了实际的系统性能，并有助于预测各种条件下的性能，从而实现针对特定模型和硬件量身定制的自适应内存管理。

3.4 自适应内存管理

ProTrain的自适应内存管理由三个组件组成：块感知运行时估算器、峰值内存使用估算器和最优配置搜索。这些组件协同工作，选择能在确保内存需求不超过硬件限制的前提下最小化运行时的配置。

块感知运行时估算器。ProTrain的运行时估算器在块（chunk）级别分析计算和通信时间，这与其主要基于块的操作设计相符，如图1所示。它将一个块内的各个操作符的运行时进行汇总，用于前向计算；在后向计算中，它还额外包括了检查点块的重计算时间，这得益于将每个块的参数合并到单个块中的设计。对于通信，ProTrain使用块大小和传输带宽来估算上传和卸载时间，并对块收集和归约操作的集体通信开销进行建模。通信开销的估算考虑了当前配置下的重叠情况。例如，当启用激活交换时，参数预取的传输吞吐量可能会受到影响。一个块的计算与下一个块的预取以及前一个块的卸载是重叠的。通过比较这些开销，估算器确定一个块是计算密集型还是通信密集型，并使用较大的值作为该块的运行时估算。每个块的前向和后向运行时估算被累加以计算一次迭代的总前向（TFWD）和后向运行时（TBWD）。对于使用GPU参数更新的持久化块（TGPU_OPTIM），我们利用FusedAdam优化器【索引31, Api documentation of apex optimizers, 2018, NVIDIA】，其运行时可根据参数大小预测。对于使用CPU参数更新的非持久化块（TCPU_OPTIM），估算器会检查更新是否能与剩余的计算完全重叠，如公式（1）所示。

$$T_{\text {Iteration }}=T_{\text {FWD }}+max\left\{T_{\text {BWD }}+T_{\text {GPU\_OPTIM }}, T_{\text {CPU\_OPTIM }}\right\}$$

峰值内存使用估算器。如第3.3节所述，我们的分析器捕获操作前（PriorOpi）和操作期间（Opi）的当前（C）和峰值（P）内存变化，重点关注通常出现峰值内存使用的后向传播过程。最初，我们排除模型状态，并将当前内存（$M_{Current}$）设置为所有操作的前向内存使用量和激活内存（$A_{Op}$）的总和，同时考虑到分析期间激活的丢弃。然后，我们迭代计算每一层的当前和峰值内存，以获得基线峰值内存（$M_{Peak}$），如公式（2）和（3）所示。对于模型状态，内存使用在基于块的管理下是可预测的，由块大小、持久化块数量和块缓冲区数量决定。逐块激活管理带来的内存节省是根据指定用于交换和检查点的块数计算的。最终的峰值内存需求是基线峰值内存与模型状态内存之和，再减去逐块激活管理带来的节省。这种方法提供了对内存需求的精确而全面的概览，解决了现有工作【索引48, H3t: Efficient integration of memory optimization and parallelism for large-scale transformer training, 2024, Advances in Neural Information Processing Systems】【索引9, Parallel training of pre-trained models via chunk-based dynamic memory management, 2022, IEEE Transactions on Parallel and Distributed Systems】中内存估算存在的挑战。

$$M_{\text{Current}} = M_{\text{Current}} + C_{\text{PriorOp}_{i}} + C_{\text{Op}_{i}} - A_{\text{Op}_{i}}$$ $$M_{\text {Peak}}=max\{M_{\text {Peak}}, M_{\text {Current }}+P_{\text {PriorOp}_{i}}, M_{\text {Current}}+C_{\text {PriorOp}_{i}}+P_{\mathrm{Op}_{i}}\}$$

最优配置搜索。在ProTrain中，可调整的配置包括持久化块的数量 $N_{persist}$、块缓冲区的数量 $N_{buffer}$、交换块的数量 $N_{swap}$ 和检查点块的数量 $N_{checkpoint}$。这些配置受到约束，即它们的总和不能超过总的块数或层数。增加持久化块和块缓冲区的数量通常会带来性能提升，因为减少了参数预取的需求，但这种好处会被增加的GPU内存使用所抵消。相反，更多的交换和检查点可以带来显著的内存节省，但交换可能会因带宽有限而干扰参数预取，而检查点会引入额外的重计算开销。孤立地优化这些设置是不够的，因为一个配置的改进可能会导致其他方面的性能次优。因此，我们需要对所有配置进行整体考虑，以评估它们对性能和内存使用的综合影响。

ProTrain中的配置空间是结构化且有限的，允许对所有可能的配置进行详尽搜索。即便如此，ProTrain仍采用特定的剪枝策略来进一步缩小搜索空间。例如，可交换块的最大数量受到交换间隔的限制，以确保它们与前向计算重叠。在后向阶段，系统会监控块预取的带宽使用情况，以确保为激活预取留有足够的带宽。此外，由于配置是从小到大遍历的，任何导致内存溢出的交换和检查点组合都会被立即丢弃，并且涉及此组合的后续迭代也会被跳过。对于每个可行的配置，ProTrain的运行时估算器会预测其运行时，并选择运行时最短的配置作为最终设置。

A4 实验环境

• 工作负载：

  • 模型：测试了包括GPT-2【索引36, Language models are unsupervised multitask learners, 2019, OpenAI blog】、OPT【索引49, Opt: Open pre-trained transformer language models, 2022, arXiv preprint arXiv:2205.01068】、Mistral【索引19, Mistral 7b, 2023, arXiv preprint arXiv:2310.06825】和LLaMA【索引45, Llama: Open and efficient foundation language models, 2023, arXiv preprint arXiv:2302.13971】在内的多种模型。
  • 规模：通过改变隐藏维度、Transformer块数量和注意力头数，生成了不同参数大小的模型。
  • 配置：默认序列长度设置为1024。详细模型配置见附录B.1。

• 硬件配置：

  • 平台1：1个节点，配备4块NVIDIA GeForce RTX 3090 24GB GPU，384GB DRAM。
  • 平台2：1个节点，配备4块NVIDIA A100 SXM4 80GB GPU，支持NVLink 3.0，1TB DRAM。
  • 详细硬件信息见附录B.2。

• 基线系统：

  • FSDP【索引51, Pytorch fsdp: experiences on scaling fully sharded data parallel, 2023, arXiv preprint arXiv:2304.11277】：PyTorch原生支持的ZeRO-3技术。
  • DeepSpeed【索引39, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】：广泛使用的分布式训练框架，采用ZeRO和卸载技术（测试中使用ZeRO-3）。
  • Colossal-AI【索引24, Colossal-ai: A unified deep learning system for large-scale parallel training, 2023, Proceedings of the 52nd International Conference on Parallel Processing】：采用与ZeRO-3兼容的基于块的内存管理。

A4 实验结果

4.2.1 模型规模

最大可训练模型尺寸：如表1所示，ProTrain在模型规模支持上表现出色。在单张RTX 3090 GPU上，ProTrain能够训练高达300亿参数的模型，并且在四卡配置下保持了这一能力。在性能更强的A100 GPU上，ProTrain可处理高达700亿参数的模型，在单卡和四卡配置中分别比Colossal-AI和DeepSpeed大1.75倍和2.06倍。FSDP在GPU数量较少时表现不佳。这些结果表明ProTrain能高效利用异构内存资源，推动了大型语言模型训练的普及。

表1：最大可训练模型尺寸（单位：十亿）

4.2.2 训练吞吐量

最大训练吞吐量对比：图3展示了在四卡RTX 3090和A100 GPU上，各框架通过调整批量大小所能达到的最大训练吞吐量。结果显示，ProTrain在不同硬件和模型配置下均持续领先。

• 在RTX 3090上，ProTrain的平均吞吐量比其他框架高出约1.77至2.71倍。
• 在A100上，ProTrain的吞吐量分别比DeepSpeed、Colossal-AI和FSDP高出1.85倍、1.43倍和2.25倍。

随着模型规模增大，ProTrain的性能优势更加明显，例如，在训练15B的GPT-2（RTX 3090）和34B的LLaMA（A100）时，其速度分别是FSDP的5.05倍和3.31倍。总体而言，ProTrain显著提升了大型模型训练的效率。

4.2.3 性能可扩展性

GPU数量扩展性：如图4(a)所示，在训练10B GPT-2模型时，ProTrain表现出良好的扩展性。当GPU数量从1增加到4时，其吞吐量从单卡增长了3.5倍，达到2493 token/s，优于其他框架。

批量大小扩展性：图4(b)展示了在四卡RTX 3090上训练10B GPT-2时，不同批量大小下单次迭代的时间分解。
- 在小批量下，ProTrain通过高效的计算与IO重叠（参数更新时间可忽略不计）以及自适应内存管理策略，性能远超其他框架。
- 随着批量增大，所有框架的单次迭代时间都增加。此时，ProTrain的性能优势主要来源于其更优的重叠策略。

4.2.4 自适应内存管理的效果

自适应与固定配置对比：图5(a)比较了在四卡RTX 3090上训练10B GPT-2时，ProTrain使用自适应配置与固定配置（所有块使用梯度检查点，仅用3个块缓冲区，无持久化块和交换块）的运行时。结果表明，自适应配置通常能实现更低的运行时，尤其是在内存较充裕的小批量场景下。随着批量增大，内存成为瓶颈，自适应配置的效果趋近于固定配置。

预测与实际运行时对比：图5(b)展示了10B GPT-2在不同配置下的预测与实际运行时。两者高度吻合，证明了ProTrain运行时估算器的准确性，它能有效找到在特定模型和硬件下平衡内存与运行时的最优配置，从而提升性能。

A5 结论

本文提出了ProTrain，一个旨在轻松协调内存、计算和IO的新型训练系统。ProTrain通过自适应内存管理简化了训练过程，使用户无需手动干预即可实现比现有最先进框架快5倍的速度。重要的是，ProTrain使得在单块A100 GPU上训练高达700亿参数的模型成为可能。我们希望我们的工作能帮助那些GPU资源有限的AI研究人员和从业者，让大型模型变得更加易于获取。

A6 附录

A 实现细节

ProTrain使用Python语言实现，总代码量为6000行。它提供了非常简单的API，方便使用，在PyTorch训练脚本的基础上仅需修改不到5行代码。与现有工作不同，由于其自适应内存管理设计，ProTrain几乎不需要任何手动配置工作。

A.1 自适应块大小

自适应块大小确定机制。ProTrain采用一种动态搜索机制来确定模型训练的最佳块大小。它根据参数的执行顺序来组织参数，并确保一个块内的所有参数都分在同一个块中。对于通常在层间共享参数的Transformer模型，ProTrain使用参数的首次出现作为排序标准。为了找到最高效的块大小，ProTrain进行网格搜索，模拟不同块大小下的内存浪费情况，以确定能使浪费最小化的大小。

A.2 内存优化

• 主动内存分配：ProTrain为那些持续到训练完成的张量预先分配内存，包括早期分配用于参数和优化器状态的持久化块，以及GPU块缓冲区。这种主动策略减少了内存分配次数，并通过将长生命周期的张量组织在一起，减轻了内存碎片，确保了更规整高效的内存布局。
• 单流内存分配：ProTrain将内存分配统一在默认流中以提高内存利用率。PyTorch的分配器采用多堆设计，每个流都有自己的堆，这限制了跨堆的内存重用，并需要使用record_stream()来确保正确性。通过使用单个流进行所有分配并由我们自己直接管理释放同步，我们有效地防止了误用和重分配冲突，从而提高了内存效率。
• 定制化的固定内存分配器：我们观察到默认的固定内存分配器（CUDAHostAllocator）通常会向上取整到最接近的2的幂次方进行过度分配，导致严重的内存浪费。为了解决这种低效问题，ProTrain开发了一个定制的固定内存分配器，它利用自适应内存管理的洞察力来精确确定固定内存需求，从而提供更精细的控制并避免默认分配器的过度内存预留。

B 实验设置

B.1 模型配置

实验中使用的模型配置如表2所示。

表2：模型配置

B.2 硬件配置

• 4× RTX 3090：该系统包含四块NVIDIA GeForce RTX 3090 GPU，每块24GB内存。CPU为Intel(R) Xeon(R) Silver 4214R @ 2.40GHz，24核。CPU DRAM大小为384GB。PCIe版本为3，带宽为15.8GB/s。该配置中无NVLink。
• 4× A100：该系统包含四块NVIDIA A100 GPU，每块80GB内存。CPU为Intel(R) Xeon(R) Platinum 8480+，112核。CPU DRAM大小为1TB。PCIe版本为4，带宽为31.5GB/s。GPU之间通过NVLink 3.0全连接，带宽为300GB/s。

B.3 基线配置

• DeepSpeed：我们使用了DeepSpeed-0.12.1版本，启用了ZeRO-3以及参数和优化器状态的卸载。为了获得最佳性能，我们对配置进行了微调，关键设置包括stage3_max_live_parameter, stage3_max_reuse_distance, stage3_prefetch_bucket_size和reduce_bucket_size。
• Colossal-AI：我们使用了0.3.3版本，并结合Gemini插件以支持基于块的内存管理。该设置采用静态放置策略，并启用了参数和优化器状态的卸载，以支持大型模型的训练。
• FSDP (Fully Sharded Data Parallel)：该功能集成在PyTorch-2.0.1中，我们使用了transformer_auto_wrap_policy来确保每个Transformer块被封装在单个FlatParameter中。我们还启用了CPU卸载以适应更大模型的训练。

C 完整实验结果

C.1 A100 GPU上的吞吐量可扩展性

图6展示了ProTrain在四块A100 GPU上训练LLaMA 34B时与其他框架相比的可扩展性性能。与RTX 3090 GPU上的结果类似，ProTrain展示了卓越的可扩展性，在所有配置中均实现了最佳性能。图6(b)突显了参数更新重叠的有效性。DeepSpeed和FSDP将所有参数更新卸载到CPU，导致成本更高，而ProTrain则将部分更新保留在GPU上，并将CPU更新与后向执行重叠。这种方法与其他框架相比，显著减少了总迭代时间。

C.2 使用与不使用卸载时的训练吞吐量

尽管ProTrain是为模型无法完全装入GPU内存（需要卸载）的场景设计的，但在非卸载场景中，与基线相比，它也表现出优异的性能。如表3所示，当DeepSpeed和Colossal-AI在不卸载的情况下运行时，它们对较小模型的训练吞吐量有所提高。然而，随着模型尺寸的增加，不使用卸载可训练的批量大小减小，从而削弱了性能优势。例如，Colossal-AI在LLaMA 13B上的性能在不卸载时比卸载时慢15%。总的来说，无论基线是否使用卸载，ProTrain始终能实现最佳性能，显示了其在不同训练场景下的多功能性和适应性。

表3：在四块A100 GPU上使用和不使用卸载时的最大训练吞吐量（单位：token/s）

C.3 运行时/峰值内存使用估算器的效果

图7比较了在四块RTX 3090 GPU上使用ProTrain选择的配置时，预测的运行时和峰值内存使用与实际值的对比。上图显示，运行时预测误差不超过5%，反映了运行时估算器在不同模型和批量大小下的高准确性。下图比较了预测和实际的峰值内存使用（使用max_memory_allocated测量）。随着批量大小的增加，预测误差略有增加，通常高估不超过10%。这种保守的估算通过考虑内存碎片，有助于降低内存不足错误的风险，从而确保在不同训练条件下的可靠性能。总的来说，这些结果验证了ProTrain的运行时和内存估算器的可靠性，证实了它们在自适应内存管理中的作用。

D 相关工作

交换与重计算。交换（Swapping）【索引41, vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design, 2016, 2016 49th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO)】【索引22, Tflms: Large model support in tensorflow by graph rewriting, 2018, arXiv preprint arXiv:1807.02037】【索引14, Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping, 2020, Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems】【索引40, Zero-offload: Democratizing billion-scale model training, 2021, 2021 USENIX Annual Technical Conference (USENIX ATC 21)】【索引44, Stronghold: fast and affordable billion-scale deep learning model training, 2022, SC22: International Conference for High Performance Computing, Networking, Storage and Analysis】是一种常用技术，它利用CPU内存等外部存储来卸载张量，从而扩展可用于训练的内存。传统的交换方法主要关注卸载激活值，SwapAdvisor【索引14, Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping, 2020, Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems】将其扩展到参数，而ZeRO-offload【索引40, Zero-offload: Democratizing billion-scale model training, 2021, 2021 USENIX Annual Technical Conference (USENIX ATC 21)】进一步将其扩展到优化器状态。重计算（Recomputation）【索引6, Training deep nets with sublinear memory cost, 2016, arXiv preprint arXiv:1604.06174】【索引17, Checkmate: Breaking the memory wall with optimal tensor rematerialization, 2020, Proceedings of Machine Learning and Systems】【索引13, Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory, 2019, arXiv preprint arXiv:1911.13214】【索引50, Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch, 2023, International Conference on Machine Learning】【索引21, Reducing activation recomputation in large transformer models, 2023, Proceedings of Machine Learning and Systems】，也称为梯度检查点，是另一种广泛使用的技术，它通过在后向传播期间增加额外的重计算时间来换取激活内存使用的减少。最初，Chen等人【索引6, Training deep nets with sublinear memory cost, 2016, arXiv preprint arXiv:1604.06174】的工作侧重于同构顺序网络，后续研究【索引17, Checkmate: Breaking the memory wall with optimal tensor rematerialization, 2020, Proceedings of Machine Learning and Systems】【索引13, Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory, 2019, arXiv preprint arXiv:1911.13214】将其适用性扩展到异构网络。考虑到Transformer的规模和复杂性，它通常包含大量层，以前的方法效率较低。因此，Rockmate【索引50, Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch, 2023, International Conference on Machine Learning】通过将模型划分为细粒度的块来优化计划生成。NVIDIA进一步提出了选择性激活重计算，该技术对部分层进行检查点和重计算【索引21, Reducing activation recomputation in large transformer models, 2023, Proceedings of Machine Learning and Systems】。为了两全其美，一些工作【索引32, Capuchin: Tensor-based gpu memory management for deep learning, 2020, ASPLOS ’20】【索引2, Efficient combination of rematerialization and offloading for training dnns, 2021, Advances in Neural Information Processing Systems】联合优化了交换和重计算，而ProTrain的独特之处在于其为适应Transformer的特定结构而量身定制。

并行化技术。训练大型模型通常需要多个GPU来有效分配内存和计算负载。这种方法利用了三种并行技术：数据并行（DP）【索引23, Pytorch distributed: Experiences on accelerating data parallel training, 2020, arXiv preprint arXiv:2006.15704】，它在设备间分配输入数据；张量并行（TP）【索引42, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv preprint arXiv:1909.08053】，它在一个操作符内将张量划分到多个GPU上；以及流水线并行（PP）【索引16, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019, Advances in neural information processing systems】，它将模型划分为分配给不同设备的子图。对于DP，引入了零冗余优化器（ZeRO）【索引37, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】，通过在不同设备间划分和分配优化器状态、梯度和参数来提高内存效率，显著减少了内存需求。在这项工作中，我们专注于使用ZeRO的数据并行，并将其他类型的并行留作未来工作。

计算与通信重叠。有大量关于重叠计算和通信的工作，许多研究【索引28, Better together: Jointly optimizing {ML} collective scheduling and execution planning using {SYNDICATE}, 2023, 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23)】【索引12, Tictac: Accelerating distributed deep learning with communication scheduling, 2019, Proceedings of Machine Learning and Systems】【索引33, A generic communication scheduler for distributed dnn training acceleration, 2019, Proceedings of the 27th ACM Symposium on Operating Systems Principles】【索引18, Breaking the computation and communication abstraction barrier in distributed machine learning workloads, 2022, Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems】【索引4, Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】侧重于替换、拆分和调度复杂操作符以实现细粒度重叠。CoCoNet【索引18, Breaking the computation and communication abstraction barrier in distributed machine learning workloads, 2022, Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems】增强了底层操作符的优化，而Centauri【索引4, Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】则将其扩展到图级调度，提供了更具层次化的抽象。尽管取得了这些进展，但大多数研究都集中在分布式情况下集体通信操作的优化上。然而，ProTrain还考虑了在GPU内存受限条件下CPU和GPU之间的通信，使其与现有研究是正交的。

Transformer训练框架。为响应对高效训练Transformer日益增长的需求，已开发出多种专用框架，各自提供独特的功能和优化。微软的DeepSpeed【索引39, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】通过ZeRO系列技术【索引37, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】【索引40, Zero-offload: Democratizing billion-scale model training, 2021, 2021 USENIX Annual Technical Conference (USENIX ATC 21)】【索引38, Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning, 2021, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】提高了训练效率，并支持各种并行策略、交换和重计算。来自HPC-AI Tech的Colossal-AI【索引24, Colossal-ai: A unified deep learning system for large-scale parallel training, 2023, Proceedings of the 52nd International Conference on Parallel Processing】提供类似功能，但其特色在于采用了基于块的内存管理方法【索引9, Parallel training of pre-trained models via chunk-based dynamic memory management, 2022, IEEE Transactions on Parallel and Distributed Systems】，我们的工作也采纳了这一点。另一方面，NVIDIA的Megatron-LM【索引42, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv preprint arXiv:1909.08053】专注于模型并行。这些框架专为大规模Transformer训练而设计，并得到学术界努力【索引44, Stronghold: fast and affordable billion-scale deep learning model training, 2022, SC22: International Conference for High Performance Computing, Networking, Storage and Analysis】【索引25, Harmony: Overcoming the hurdles of gpu memory capacity to train massive dnn models on commodity servers, 2022, arXiv preprint arXiv:2202.01306】【索引10, Mobius: Fine tuning large-scale models on commodity gpu servers, 2023, Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】的支持，以便在较小系统上进行训练。

E 讨论

局限性与未来工作。ProTrain是为小型集群设计的，这可能会限制其在GPU间通信更为普遍的大规模训练场景中的性能。在这种情况下，网络拓扑和通信模式成为关键因素，需要考虑多样的并行策略以获得最佳性能。目前，ProTrain仅支持ZeRO风格的数据并行，我们旨在在未来的工作中探索各种并行技术对性能的影响。