【论文阅读】TBA Faster Large Language Model Training Using SSD Based Activation Offloading

整体工作

这篇文章针对当前大模型训练，激活数据在GPU内存中占用量高、主导内存使用，影响限制着模型训练性能的问题，提出了一种解决方案：将激活卸载到比GPU、CPU内存更高容量的NVMe ssd上。通过设计使I/O与计算重叠，即保留了相比存储在CPU内存中近似的性能，又比分层完全重计算获得更好的内存节省。

文中部分术语解释

激活：“激活”（activations）通常指的是神经网络各层在前向传播（forward pass）过程中计算出的中间结果。这些激活值在反向传播（backward pass）过程中用于计算梯度，因此需要暂时存储在内存中

激活检查点（Activation Checkpointing）：在前向传播过程中有选择性地存储一部分激活值，在反向传播过程中需要计算梯度时再重新计算这些激活值。（时间换空间）

梯度累积（Gradient Accumulation）：通过多次前向和反向传播累积梯度，然后再进行一次权重更新。可以通过较小的批量大小减少激活的内存占用。

研究背景

CPU内存已成为llm持续成长的瓶颈。如下图所示

1. GPU内存容量的增长速度比LLM尺寸扩展速度和GPU FP16吞吐量的提高速度慢60%左右。
2. 最近的llm训练中，80%的GPU内存使用由激活组成。
3. 激活所需的内存比任何其他内存使用增长得更快
   
   FP 16吞吐量增长用于深度学习训练的GPU（右纵轴）与LLM的模型大小对齐（左纵轴），但GPU内存容量（左纵轴）落在[84]后面。横轴显示发布日期。点代表Nvidia 100-辅助并行绿色虚线以FP 16吞吐量增长率（黄色虚线）的50%增长，前者的增长速度快于内存容量增长率（红色虚线）。

常见的缓解GPU内存不足措施：

1. 减小批大小：如果原来每次迭代你传递给模型的是64个样本（批大小为64），那么可以减小这个数字，比如改成32或者16。
2. 梯度累积：一个批被分为micro-batch，希望模拟一个大批大小为64的效果，但由于GPU内存限制，每次只能处理16个样本。设定累积步数为4。经过4次迭代，在一次性使用累计的梯度来更新模型参数。
3. 使用激活检查点：将一些激活保留在GPU内存中，而其他激活被刷新，然后在反向传播期间重新计算。

但GPU的设计不是为了小输入而设计，上述1、2措施导致设备利用率不足和数学库性能降低。措施3虽然减小了内存需求，但未来依然会成为瓶颈，跟不上llm规模增长。

如图2所示，

1. 集群和云实例通常受限于CPU内存容量，而ssd提供的容量要高得多。
2. 有限的主机内存容量也被输入数据、检查点缓冲区和其他训练管理缓冲区所消耗，更少的内存可以用于激活。
3. 主机内存带宽是在主机CPU上运行的训练管理任务和卸载计算之间共享
4. SSD扩展增加更容易，本地+远程存储
   

本文设计并实现了TAB框架，将LLM训练中的激活卸载到NVMe ssd上；对SSD的寿命和 GPU PCIe带宽进行分析，证明了大规模系统中TAB的可行性和性能优势；引入3种优化技术：张量重复数据删除、张量转发和自适应卸载算法来实现数据传输和计算完全重叠，几乎不引入额外的性能开销。

对使用SSD的必要性和可行性进一步分析：

定性比较，激活的内存将继续主导GPU内存的使用。
对于具有𝐿层的模型，激活检查点可以将内存需求从𝑂(𝐿)减少到𝑂(√𝐿)，但也跟不上模型增长的需求。

每个被卸载的张量的大小很容易达到数百mb，读写较密集，

读写速度：闪存的随机写延迟已经降低到几十微秒，NVMe SSD的数据读写速率现在是几GB/s。

耐用性：激活卸载之类的写密集型场景下，单元的类型和数量、写入放大因子(WAF)和过度配置等因素影响实际寿命。但NVMe SSD的寿命已经足够支持大规模训练。（CPU基于电容器和晶体管存储，寿命基本无限）

PBW就是指寿命为写入多少PB的数据

GDS(GPU Direct Storage)支持：支持GPU与本地或远端NVMe ssd之间的数据直接连接。由于不需要让CPU参与缓冲，这种方法增强了带宽，减少了延迟和CPU负载。

异步数据传输。 这些系统[1,67,75]要么在加载卸载数据时阻塞训练计算，要么在每层同步。因此，在关键路径中暴露了I/O延迟。TBA通过将I/O与GPU计算重叠来隐藏I/O延迟。

互操作性。 由于LLM培训需要Python包的协同作用，并且生态系统正在迅速发展，因此卸载功能与同一库或其他库中的其他组件具有良好的互操作性至关重要。TBA依赖于PyTorch执行的进程本地交替，并且可以与分布式框架(如Megatron和DeepSpeed)一起工作。相比之下，DeepSpeed的卸载功能(例如ZeRO- infinity)仅在某些ZeRO阶段可用。Flexgen和LLM在Flash中有自己的运行时，不与分布式框架一起工作。

定量分析：

SSD寿命（左侧粉色垂直轴）、PCIe写入带宽（左侧蓝色垂直轴）和每个GPU的最大激活大小（右侧垂直轴）的估计。超过5年的寿命显示在粉红色条的顶部。水平轴显示GPU的数量，框架和模型大小。

（有趣点：当系统尺寸和模型尺寸扩大时，所需的PCIe写带宽减少，预计寿命增加。这种影响的发生是因为较大的系统意味着增加的通信开销和降低的计算效率，从而减慢了每个GPU上的训练迭代。）
同时，如菱形标记所示，每个GPU的最大激活大小范围从0.4 TB到1.8 TB，CPU内存无法保存，SSD是当前唯一选择。

具体设计和实现

图4以PyTorch为例演示了TBA是如何工作的。TBA启动自己的线程(与PyTorch的执行线程分开)来存储张量(1)并将它们加载回来(5)。在前向传播(F)中，一旦产生激活的操作完成(1)，就开始卸载激活。当激活在反向传播(B)中被重用时，预取(5)以与前向传播(2)中记录的层顺序相反的顺序发生。当最后一层(在示例中是L3)开始反向传播时，保留该层的激活而不是卸载它们(4)。TBA为每个微批保存单独的记录。当微批更改(2)时，TBA将自己的记录切换到与新微批对应的记录。

图5显示了TBA软件组件。
张量缓存，进行有效的卸载和重新加载张量，促进内存的释放，以及在需要反向传播之前将张量预取回内存。。为了实现这一点，它使用PyTorch钩子来改变PyTorch的执行；

针对同一节点内NVMe SSD的SSD卸载器：每个卸载器封装了将CUDA张量传输到卸载目标或从卸载目标传输的逻辑。SSD卸载程序利用GDS python绑定，使用LD_PRELOAD机制，CUDA malloc钩子是一个共享库，可以改变CUDA内存分配和释放API调用，以注册和释放以获得最佳GDS性能。

CPU offloader：为将来在具有大量远程SSD存储的集群上工作而设计。它由一个带有预分配的固定CPU内存的分配器提供支持。池大小是通过分析第一个训练步骤来确定的。新的API调用被添加到Megatron和DeepSpeed的调度器中，这样张量缓存可以得到关于阶段变化和micro-batch处理变化的提示，例如图4中的3和4。

上述算法1中的configure_tensor_cache()显示了在训练之前配置张量缓存的逻辑，注册hook，记录参数避免被卸载。

deepspeed_exec_schedule()显示了添加到DeepSpeed的管道调度程序的提示。在执行之前和之后，调用api来通知张量缓存即将到来的阶段(第13行)和操作的完成(第15行)，实现张量缓存预取数据。

Tensor Cache 设计

张量缓存是一个内存结构，它管理对所有激活的引用，并跟踪激活的状态，包括它们是否正在被卸载，文件系统中的路径等。

自适应卸载算法

为了从张量卸载中受益，卸载的张量拥有的GPU内存必须在张量不使用时释放。然而，PyTorch默认存储对计算图上所有激活的引用，不允许回收GPU内存。张量缓存改变了PyTorch的执行，使得激活的标识符（而不是引用）在计算图上注册;

在PyTorch重用激活张量时，张量缓存从计算图中获取标识符，并将其用作返回所请求张量的键。在前向传播中，当张量完成卸载时，张量缓存不再持有对它的引用，一旦Python控制流离开使用张量对象的函数作用域，它的内存就可以通过Python垃圾回收来回收。在反向传播中，张量缓存通过在使用前从SSD加载张量来保存对张量的引用;当张量引用的所有模块范围都完成时，不再保存引用，允许回收其内存。

算法中具体说明如何在开始和结束时标记，进行回收、预取等操作，通过维护栈堆判断是否进入激活张量作用域。

也通过作用域、训练阶段判断决定是否将当前张量保存在GPU或CPU，还是卸载到SSD。

重复张量数据删除

张量缓存有一个get_id()函数来为每个张量分配一个唯一的标识符。pytorch原生id()的缺点是它的返回值与GPU内存地址相关。当TBA卸载激活时，一旦控制流超出其使用范围，后者将被垃圾收集清除。gpu mmemory地址可能被重复使用，导致标识符冲突。为了解决这个问题，get_id()将第一次处理张量时的时间戳与张量形状结合起来作为唯一标识符:使用t.untyped_storage()而不是t。
在后向传播计算中也直接使用这个标识符，在各步骤中保持一致。

张量数据卸载和转发

张量缓存有两个线程池——一个用于存储张量，另一个用于加载张量。提交给每个线程池的任务按照先进先出(FIFO)的顺序执行。

1. 为了隐藏I/O延迟，张量缓存在相应模块的向后传播之前开始预取每个激活。
2. 数据转发：在加载一个张量时，如果它仍然被存储，张量缓存将返回其原始的内存引用，以跳过从SSD加载。直接加载引用使用或者将其存储在张量缓存中，供后续使用。

整体评估

q1：TBA在多大程度上隐藏了I/O延迟?

q2：TBA减少了多少峰值内存使用?

q3:给定相同的每个gpu用于激活内存预算，TBA可以提供多少吞吐量提升?

测试：使用一台带有2× A100 PCIe gpu和7× Intel P5800X ssd的机器；在BERT（仅解码）、GPT（仅解码）和T5（编解码）上进行测试。

Q1&Q2:尽管TBA及其优化引入了额外的CPU执行逻辑，但性能比较表明，这些逻辑并不在关键路径上，而是由GPU计算定义关键路径，CPU的作用主要在于在当前GPU操作完成之前启动新的GPU作业，非瓶颈，TBA的额外工作不会导致新任务延迟到达GPU。就激活的内存使用而言，TBA有效地将这些情况下的峰值降低了28%-40%。

重新计算-卸载-保持(ROK)曲线上的运行来比较三种不同的策略。在这两种情况下，TBA降低了GPU激活的内存峰值，允许更大的批处理大小来获得更高的吞吐量，利用TBA来运行更大的模型，或者使用更少的gpu。

分析SSD可能性中定量分析准确性，们将TBA的卸载量与模型估计值进行比较如下，同时隐藏维数越大，PCIe写带宽越小。

相关工作

交换和卸载。许多具有卸载能力的LLM系统是仅推理的（LLM in a flash、PagedAttention、FlexGen）。在推理中，权重和KV缓存永远不会改变，并在迭代中重复使用;研究人员利用这一点来提高局部性和内存效率。然而，在LLM训练中，权重在每次迭代中更新，并且所有张量在迭代中都会发生变化。一些工作利用卸载功能ZeRO-infinity进行训练，但大多数都是为了在较小的系统中容纳较大的模型而牺牲性能。它们缺乏维持性能所需的异步数据传输能力。

另一个方向是将数据和相关计算卸载到CPU[Fiddler,ZeRO
-Offload,Power Infer]。卸载的计算相对较轻，卸载的数据包括梯度、权重中的稀疏元素等。认识到这个方向，我们的工作是正交的，因为我们通过GDS将激活卸载到SSD，以最大限度地减少对CPU的干扰。激活用于梯度计算，这是计算密集型的，最好只在GPU上完成。

在LLM被大规模采用之前，人们一直在为深度学习卸载数据[5，24，61，70，87]。它们中的大多数将数据卸载到主内存，而一些(FlashNeuron)(针对DDN设计，将中间数据压缩存储至SSD)启用GPU-SSD数据路径。LLM训练是独特的，因为大规模并行及其对优化器状态，梯度和权重的内存使用的影响是设计空间的基础。TBA自然支持多个GPU。此外，我们证明了它的可行性集群，并介绍了ROK曲线，以帮助设计选择。另一方面，LLM对计算能力的需求如此之高，以至于它刺激了专用硬件的快速发展，例如，Transformer引擎[55]和分布式框架。这就是为什么我们要确保良好的互操作性。相比之下，这方面的大多数早期工作都绑定到特定的PyTorch版本或支持选择层的自定义运行时。

量化和稀疏性。一些关于卸载的工作使用量化和/或稀疏性来减少I/O大小[1，5，75]。为了减少计算，已经提出了算法来调整参数并将稀疏性引入模型[14，18，32，42，92]。混合专家（MoE）[74]是在这个方向上，因为它稀疏了MLP中的Token到神经元的连接到Token到专家Experts的连接。一些算法引入结构化稀疏性，例如，N：M [94]稀疏性和2：4 [62]稀疏性。另一方面，存在框架和专用内核来加速具有量化和/或稀疏性的模型[19，20，76，93]。一些内核利用专用硬件，安培张量核心[10，50]。这些技术与我们的工作是正交的，可以用来在使用TBA时替换模型和加速计算。值得注意的是，给定硬件，将计算与PCIe传输完全重叠的重用因子将根据新的数值格式或稀疏访问模式而改变。我们相信，TBA的自适应卸载算法有助于优化这些情况下的卸载量。

最佳化核心。之前的工作开发了优化的内核来加速LLM [11，12，56]。一些内核使用特殊的硬件[57]。TBA的互操作性确保了它可以轻松地与这些技术和即将到来的技术结合使用。

结论

在LLM训练系统中，激活在日益有限的GPU内存中占主导地位，TBA通过将激活卸载到SSD来解决这个问题。文中证明了它的可行性，在大规模系统的建模，将直接的GPU-SSD数据路径和良好的互操作性纳入TBA。为了将计算与数据传输完全重叠，TBA具有异步数据传输、张量重复数据删除、转发和自适应卸载功能。评估表明，TBA将激活峰值内存使用减少了47%，开销可以忽略不计，引入了ROK曲线来展示TBA卸载相对于重新计算和将激活保留在CPU内存中的优势。