摘要
GPU内存容量的增长速度跟不上大型语言模型(llm)的增长速度，阻碍了模型的训练过程。特别是，激活——在前向传播过程中产生的中间张量，并在后向传播中重用——主导着GPU内存的使用。为了应对这一挑战，我们建议TBA将激活有效地卸载到高容量NVMe ssd上。这种方法通过自适应地将数据传输与计算重叠来减少GPU内存的使用，而不会影响性能。TBA与流行的深度学习框架(如PyTorch、Megatron和DeepSpeed)兼容，并采用了张量重复数据删除、转发和自适应卸载等技术来进一步提高效率。我们在流行的LLMs如GPT、BERT和T5上进行了广泛的实验。结果表明，TBA有效地减少了47%的激活峰值内存使用。同时，TBA完美地将I/O与计算重叠，并且产生的性能开销可以忽略不计。我们引入了recompute-offloadkeep (ROK)曲线来比较TBA卸载与其他两种张量放置策略：将激活保持在GPU内存中和分层完全重新计算。我们发现TBA比分层完全重计算获得更好的内存节省，同时保留了将激活保存在内存中的性能。
1 介绍
GPU内存容量已成为llm持续成长的瓶颈。如图1所示，GPU内存容量的增长速度比LLM尺寸扩展速度和GPU FP16吞吐量的提高速度慢60%左右。大约80%用于训练最近llm的GPU内存由激活组成[35,41]，激活是由前向传播产生的中间张量，并在后向传播中重用。此外，激活所需的内存比任何其他内存使用增长得更快，使GPU内存成为未来LLM训练的更严重约束(详见2.2节)。
图1：用于深度学习训练的GPU的FP16吞吐量(右纵轴)增长与llm的模型大小(左纵轴)一致，但GPU内存容量(左纵轴)落后[84]。横轴显示发布日期。点代表Nvidia 100级gpu自K100和谷歌tpu。辅助并行的绿色虚线增长速度为FP16吞吐量增长率(黄色虚线)的50%，前者增长速度快于内存容量增长率(红色虚线)。
常见的缓解措施是减少批大小或通过梯度累积。通过梯度累积 （gradient accumulation），一个批被分成微批micro-batch，在梯度更新之间分别处理。虽然梯度积累已经被许多llm采用[28,77,90]，但GPU计算堆栈并不是为小输入而设计的，这两种缓解措施都会导致设备利用率不足[4,8]和数学库性能次优[2]。直观地说，较小的批大小可以通过更快的收敛来减少总训练计算。然而，LLM训练者已经为每个模型确定了一个关键批大小，低于该批大小收敛速度可以忽略不计甚至降低[31,45]。值得注意的是，随着训练损失的减少，关键批大小在训练期间会增加。
另一种减少GPU内存使用的常见方法是激活检查点active checkpointing。使用此策略，只有一些激活保留在GPU内存中，而其他激活被刷新，然后在反向传播期间重新计算。对于具有𝐿层的模型，激活检查点可以将内存需求从𝑂(𝐿)减少到𝑂(√𝐿)[9]。然而，正如我们在2.2节中所示，即使这种减少也不足以消除GPU内存限制对未来llm造成的瓶颈。
这项工作提出了TBA，这是一个软件框架，可以将激活卸载到非易失性内存快速(NVMe) ssd，并在反向传播需要之前重新加载激活。TBA能够将激活传输与计算完全重叠，从而减少激活内存使用，而不会产生显著的性能开销。ssd是一个比主(CPU)内存更有吸引力的目标，原因如下。首先，如图2所示，集群和云实例[21,48,51]通常受限于主机内存容量，而ssd提供的容量要高得多。有限的主机内存容量也被输入数据、检查点缓冲区和其他训练管理缓冲区所消耗，这进一步减少了可用于激活卸载的内存量。其次，主机内存带宽是在主机CPU上运行的训练管理任务和卸载计算之间共享的[30,69,82]，对于保存和恢复激活来说，带宽可能非常有限，甚至不可预测。首先，如图2所示，集群和云实例[21,48,51]通常受限于主机内存容量，而ssd提供的容量要高得多。有限的主机内存容量也被输入数据、检查点缓冲区和其他训练管理缓冲区所消耗，这进一步减少了可用于激活卸载的内存量。其次，主机内存带宽是在主机CPU上运行的训练管理任务和卸载计算之间共享的[30,69,82]，对于保存和恢复激活来说，带宽可能非常有限，甚至不可预测。相反，SSD带宽可以专门用于训练期间的激活卸载。第三，ssd更具弹性，可以通过添加更多的ssd甚至必要时的PCIe交换机，也可以通过使用可选的远程高吞吐量存储[22,43]。这种弹性使数据中心能够跟上快速增长的激活规模。相比之下，GPU云实例和集群节点的内存容量很难扩展。
[图片]
本工作的主要贡献如下:
- 为了解决GPU内存容量问题以及在LLM模型训练期间导致GPU利用率不足的问题，我们设计并实现了TBA框架，将LLM训练中的激活卸载到NVMe ssd上。我们通过模拟性能、估计SSD寿命和所需的每个gpu PCIe带宽来证明TBA在大规模系统上的可行性。
- 除了一个很小的CUDA内存分配API钩子库外，所有代码都用Python编写，TBA可以与最新的PyTorch和分布式框架(包括Megatron[77]和DeepSpeed[68])一起工作。我们使用Megatron-DeepSpeed[46]在带有7× Intel Optane ssd的双gpu节点上开发和测试了TBA。
- TBA几乎不会带来任何性能开销，因为它将数据传输与计算完全重叠。为了实现这一点，我们引入了几种优化技术，包括张量重复数据删除、张量转发和自适应卸载算法。
- 评估显示，TBA与没有TBA的原始系统在每步训练时间上几乎相同，同时将激活峰值内存使用减少了47%。我们引入了重计算-卸载-保持(ROK)曲线来比较TBA卸载与其他两种张量放置策略，将激活保持在内存中和分层完全重新计算。TBA与将激活保存在内存中具有相同的性能，并且与激活检查点相比具有更低的内存峰值。工件将随着本文的发表而发布。
2 背景
2.1 基于Transformer的LLM
[图片]
大多数LLM架构，包括GPT[64]，都是基于Transformer的[86]。如图3(a)所示，GPT模型主要由多个变压器层组成。在转换层之前，GPT接受标记化的文本，并将标记映射到具有位置信息的密集向量。任务决定了模型体系结构的最后一部分。例如，可以为文本分类任务添加分类器。图3(b)显示，每个transformer层主要由注意块和多层感知(MLP)块组成。注意块(图3(c))为每个令牌对计算一个权重，称为注意力，并通过加权求和为每个令牌生成密集向量。MLP块将每个令牌的向量转换为一个新向量。
GPT是一个纯解码器模型，因为它只涉及transformer解码器层。transformer编码器层与transformer解码器层具有相同的结构，除了后者对图3(d)中的注意掩码施加因果关系:因果掩码确保注意块为每个令牌产生的新向量仅依赖于令牌的向量，而不是依赖于该令牌。通过这种分类，transformer模型被分为(1)仅编码器，例如BERT b[15];(2)仅解码器，例如GPT, Llama[85];(3)编码器-解码器，例如T5[65]。在编码器-解码器模型中，transformer解码器层接受编码器和另一个文本的输出，并应用两个注意块——自注意块应用于新文本，交叉注意块应用于编码器序列中的标记和新文本中的标记之间。
并行化LLM训练包括将模型和数据分区和/或复制到不同的gpu中[91]。管道并行、数据并行和模型并行是所有LLM模型可用的三种并行级别，在框架中被广泛采用，例如Megatron、DeepSpeed和PyTorch 2.0[3,68,77]。流水线并行将模型划分为几个层块，并将它们放在不同的gpu上。在一个步骤中，当gpu完成其层时，输出将传递给拥有下一层的gpu。数据并行性在不同的gpu组中复制模型，并为每个组分配单独的微批。在一个步骤的最后，聚合每个组中的梯度以更新所有模型副本。模型并行对权重张量进行切分，并将切分放到不同的gpu上。每个GPU使用对应的运算符的分片执行一部分计算。考虑到系统的规模和互连性，可以使用三个级别中的全部或几个。零冗余优化器(Zero)[66]通过对优化器状态和/或可选的梯度和参数进行分片，并在这些gpu上存储分片，进一步减少了数据并行性的内存使用。
2.2 GPU 内存和模型吞吐量
正如第4节的图11所示，GPU内存容量限制了模型吞吐量。通过将激活卸载到ssd, TBA可以减轻这种限制并提高每个gpu模型的吞吐量。一个重要的问题是，根据LLM扩展的趋势，GPU内存容量是否会继续成为每个GPU模型吞吐量的限制因素。本节显示，历史趋势将使GPU内存容量成为每个GPU模型吞吐量的更重要限制因素。
神经尺度定律[29,31,45]指导LLM随着计算能力的增加而缩放。我们在推理时遵循这些规律。整个系统的GPU计算吞吐率为:<s:1> <s:1>𝐷𝑏𝑎𝑡𝑐<e:1>，其中:<s:1>是参数的个数，𝐷𝑏𝑎𝑡𝑐<e:1>是批处理中令牌的个数。根据Chinchilla缩放定律[29]得出的结论是，最优模型设计遵循的是:操作操作(操作操作)(操作操作)(操作操作)(操作操作)(操作操作)(操作操作)(操作操作)(操作操作)(操作操作)(操作操作)。整个系统GPU内存的使用由两部分组成:激活，它需要𝑆𝑎𝑐𝑡𝑣𝑎𝑡𝑜𝑛𝑠∝聚合聚合𝐷𝑏𝑎𝑡𝑐，其中，是层中的隐藏维度，是一个缓慢增长的函数，例如，∝𝑁1/3，以及所有其他的内存使用，𝑆𝑜𝑡𝑒𝑟𝑠∝聚合聚合，包括参数、梯度和优化器状态。比较这些因素，我们可以推断出(1)𝑆𝑎𝑐𝑡𝑣𝑎𝑡𝑜𝑛𝑠的增长速度快于𝑆𝑜𝑡𝑒𝑟𝑠，(2)由激活主导的全系统内存使用的增长速度略慢于计算吞吐量的增长速度(约为𝐶5/6)。然而，图1显示GPU内存容量的历史增长(红色虚线)甚至比计算吞吐量的平方根(绿色虚线)还要慢。因此，GPU内存容量将变得越来越不足以饱和计算吞吐量，而用于激活的内存将继续主导GPU内存的使用。
（查看原文）分析证明了使用activation checkpoint，也不够跟上增长速度。
2.3 SSD续航时间
价格、延迟和带宽的趋势导致ssd被广泛采用并集成到云实例和集群中[21,48,51]。闪存的随机写延迟已经降低到几十微秒[71]，NVMe SSD的数据速率现在是几GB/s。
SSD的耐用性仍然是一个问题:在诸如激活卸载之类的写密集型场景中，SSD能坚持多久?SSD的耐用性取决于单元的类型和数量、写入放大因子(WAF)和过度配置。SSD单元可以用于存储一个比特，即单级单元(slc)，也可以用于存储多个级别，例如三级单元(tlc)。一般来说，一个单元存储的比特越多，它在程序擦除(PE)周期中的寿命就越短。WAF是媒体写量与主机写量的比率——ssd一次写页面，但擦除页面块，粒度更粗。擦除部分空的块需要重新定位剩余的有效页，从而导致写入放大。反过来，供应商采用过度配置来为损耗均衡保留一些块，均匀分配跨块的写操作。
表1以当前SSD型号为例。D7-P5620代表了主流数据中心模型，具有144层(L) TLC单元和每日3次磁盘写入(DWPD)的评级。FL6和D7-P5810 ssd是为写密集型场景而设计的，具有更高的耐用性。值得注意的是，SSD耐久性评级使用JESD测试方法[27]，该方法在严格的预处理后执行随机写入。在我们的场景中，写入量很大，顺序写入，因为每个被卸载的张量的大小很容易达到数百mb。与用于确定JESD等级的写相比，这样的写对耐力更友好。例如，3-DWPD ssd通常允许的顺序写次数是JESD评级的2.5倍[38,63,78]。供应商指南[25,72,79]和经验数据[44]证实了这一差异。3.4节通过建模来说明为什么D7-P5620等主流数据中心ssd是支持大规模LLM培训系统部署tba的可行选择。
（查询内存和SSD耐用读写次数的差异）
[图片]
具有高PB写入耐久性(PBW)的批量生产SSD型号样本 PBW就是指寿命为写入多少PB的数据
DWPD：每天写全盘次数
2.4 LLM的SSD卸载系统
GDS (GPU Direct Storage)支持GPU与本地或远端NVMe ssd硬盘[26]之间的数据直接连接。由于不需要让CPU参与弹跳缓冲，这种方法增强了带宽，减少了延迟和CPU负载。
表2说明了支持SSD卸载的早期LLM系统与TBA特性之间的主要区别:正如第1节提到的，通过CPU进行传输会干扰CPU工作负载，从而影响效率。
[图片]
比较TBA与其他提供SSD卸载功能的LLM系统[1,67,75]。如果没有反向传播，推理系统可能会在一层完成后丢弃大多数中间张量。我们将“激活”概括为指推理系统中的键值(KV)缓存，因为它可以跨步骤重用。
[图片]
图4。2-microbatch-3layer(L)模型一步TBA时间线。PyTorch钩子用于触发张量缓存簿记，张量卸载(1)和张量加载(5)。在前向(F)传播中，TBA记录范围的顺序(2)，并在阶段结束时在微批之间切换(3)。TBA开始加载时，它被切换到向后(B)传播(4)。
异步数据传输。这些系统要么在加载卸载数据时阻塞训练计算，要么在每层同步。因此，在关键路径中暴露了I/O延迟。TBA通过将I/O与GPU计算重叠来隐藏I/O延迟。
互操作性。由于LLM培训需要Python包的协同作用，并且生态系统正在迅速发展，因此卸载功能与同一库或其他库中的其他组件具有良好的互操作性至关重要。TBA依赖于PyTorch执行的进程本地交替，并且可以与分布式框架(如Megatron和DeepSpeed)一起工作。相比之下，DeepSpeed的卸载功能(例如ZeRO- infinity)仅在某些ZeRO阶段可用。Flexgen和LLM在Flash中有自己的运行时，不与分布式框架一起工作。
3 设计和实现
3.1 整体概览
TBA实现了一个张量缓存，以促进有效的卸载和重新加载张量，促进内存的释放，以及在需要反向传播之前将张量预取回内存。
图4以PyTorch为例演示了TBA是如何工作的。TBA启动自己的线程(与PyTorch的执行线程分开)来存储张量(1)并将它们加载回来(5)。在前向传播(F)中，一旦产生激活的操作完成(1)，就开始卸载激活。当激活在反向传播(B)中被重用时，预取(5)以与前向传播(2)中记录的层顺序相反的顺序发生。当最后一层(在示例中是L3)开始反向传播时，保留该层的激活而不是卸载它们(4)。TBA为每个微批保存单独的记录。当微批更改(2)时，TBA将自己的记录切换到与新微批对应的记录。
[图片]
图5显示了TBA软件组件。张量缓存管理激活并执行张量卸载和加载。为了实现这一点，它使用PyTorch钩子来改变PyTorch的执行。第3.2节详细介绍了张量缓存的设计和实现。TBA具有针对同一节点内NVMe SSD的SSD卸载器和针对主机内存的CPU卸载器。每个卸载器封装了将CUDA张量传输到卸载目标或从卸载目标传输的逻辑。SSD卸载程序利用GDS python绑定kvikio[54]。使用LD_PRELOAD机制，CUDA malloc钩子是一个共享库，可以改变CUDA内存分配和免费API调用，以便正确注册和取消内存以获得最佳GDS性能。这允许我们保留PyTorch CUDA缓存内存分配器，以便与基线进行比较，而无需在PyTorch可插拔内存分配器中复制其实现或修改PyTorch运行时c++代码。CPU offloader是为将来在具有大量远程SSD存储的集群上工作而设计的。它由一个带有预分配的主机固定内存的分配器提供支持。池大小是通过分析第一个训练步骤来确定的。新的API调用被添加到Megatron和DeepSpeed的调度器中，这样张量缓存就可以得到关于阶段变化和微批处理变化的提示，例如图4中的3和4。下一段的以DeepSpeed的调度程序为例详细解释。
要使用TBA，需要在现有脚本中添加适度的代码:算法1中的configure_tensor_cache()显示了在训练之前配置张量缓存的逻辑。逻辑注册PyTorch钩子，记录参数，以便在它们注册到计算图时不卸载它们，并对调度器进行monkey-patch[89]。使用PyTorch的动态性，monkey-patch通过将自定义实现分配给包中的已定义方法来覆盖已定义的函数。deepspeed_exec_schedule()显示了添加到DeepSpeed的管道调度程序的提示。在执行之前和之后，调用api来通知张量缓存即将到来的阶段(第13行)和操作的完成(第15行)。相应地，张量缓存可以预取数据，或者等待I/O直到它完成。Megatron的调度程序也打了类似的补丁。
TBA自然地扩展到分布式框架，比如与ZeRO一起使用，因为像DeepSpeed和Megatron这样的框架将工作负载划分为构建在PyTorch内置张量功能之上的进程。通过在PyTorch下工作并保持每个进程的活动在本地，TBA直接应用于分布式启动。
3.2 基于钩子的张量缓存 Tensor Cache实现
为了从张量卸载中受益，必须在不使用张量时释放已卸载张量所拥有的GPU内存。然而，PyTorch在默认情况下存储对计算图上所有激活的引用，不允许回收GPU内存。张量缓存改变PyTorch的执行，以便在计算图上注册激活的标识符，而不是引用;在PyTorch重用激活张量时，张量缓存从计算图中获取标识符，并将其用作返回所请求张量的键。在正向传播中，当张量完成卸载时，张量缓存不再持有对它的引用，一旦Python控制流离开使用张量对象的函数作用域，它的内存就可以被Python垃圾收集回收。在反向传播中，张量缓存通过在使用张量之前从SSD加载张量来保存对它的引用;当张量所引用的所有模块作用域都已完成时，将不再持有该引用，从而允许回收其内存。
简而言之，张量缓存是一个内存结构，它管理对所有激活的引用，并跟踪激活的状态，包括它们是否正在被卸载，文件系统中的路径等。
正如算法2所示，张量缓存依赖于三个PyTorch钩子对来改变其执行行为。
前向钩子对在前向传播中工作:模块的开始触发前向预钩子，模块的结束触发前向钩子。张量缓存使用前向钩子对来维护当前作用域堆栈:当进入一个模块时，该模块被推入堆栈;当模块退出时，它被弹出。
向后的钩子对是类似的。当进入一个模块时，张量缓存预取即将到来的模块中的激活。第3.3.2节详细介绍预取。退出模块时，张量缓存将其从所有激活的作用域列表中删除。不再使用的激活将被删除，其内存将由垃圾收集释放。
当一个张量要注册到计算图上时，调用包钩子来生成一个要注册的值。当张量被重用时，调用unpack钩子来接收计算图上的对象并返回原始张量。图6显示了当pack或unpack钩子被触发时张量缓存的活动。当乘法运算符x·w完成(1〇)时，对输入x和参数w调用包钩子(2〇)。
张量缓存有一个参数记录，并相应地返回w，让它按原样注册在图上。如果张量在CPU上或者它太小(算法2中的第12行)，张量也将原样返回。正如算法2中的第16行所示，张量缓存不会卸载张量，而是只在模块要保存在内存中或向后传播时保留记录。当自适应卸载算法决定将最后几个模块保留在GPU内存中时，第一个条件成立(章节3.3.3)。当启用激活检查点的函数在反向传播中重新计算以再现激活时，第二个条件为真。对于图6中的张量x，张量缓存将其存储到ssd(3〇)中，并返回一个张量标识符。当解包钩子被触发(B〇)时，在反向传播(A〇)中，张量缓存要么等待直到预取完成(C〇)，并最终返回张量。
3.3 张量缓存机制和优化
3.3.1重删张量和排除参数。张量缓存有一个get_id()函数来为每个张量分配一个唯一的标识符。pytorch原生id()的缺点是它的返回值与GPU内存地址相关。当TBA卸载激活时，一旦控制流超出其使用范围，后者将被垃圾收集清除。gpu mmemory地址可能被重复使用，导致标识符冲突。为了解决这个问题，get_id()将第一次处理张量时的时间戳与张量形状结合起来作为唯一标识符:当get_id()第一次处理一个张量t时，get_id()将当前时间戳作为一个附加属性添加到张量的底层存储t.untyped_storage()而不是t。这是因为有时PyTorch会创建新的torch。张量对象表示相同张量。以后所有的get_id()调用都会获取属性值。此重复数据删除方案可以避免冗余I/ o。
PyTorch将向后传播中所需的所有张量注册到计算图中，包括激活和参数。由于这项工作的重点是卸载激活，张量缓存不包括模型参数。为了实现这一点，在训练之前，张量缓存记录所有模型参数的标识符(算法1中的第4行)。由于线性层存储参数张量的转置进行反向传播，因此记录转置的唯一标识符。我们的get_id()方案的一个好处是，相同参数张量的转置的标识符在各个步骤中保持一致。这是因为转置使用了原始张量的底层存储，我们在训练之前已经为其分配了时间戳。
3.3.2卸载和转发张量。张量缓存有两个线程池——一个用于存储张量，另一个用于加载张量。提交给每个线程池的任务按照先进先出(FIFO)的顺序执行。
为了隐藏I/O延迟，张量缓存在相应模块的向后传播之前开始预取每个激活。最后一个模块中的激活保存在GPU内存中，因此它们不需要预取。这个简单的方案就足够了，因为在PyTorch中，CPU在GPU执行之前提交GPU内核启动和内存操作。只要GPU作业队列中总是有I/O任务使PCIe保持忙碌，那么预取方案是等效的。
在加载一个张量时，如果它仍然被存储，张量缓存将返回其原始的内存引用，以跳过从SSD加载。我们称之为数据转发。例如，在图6中，当PyTorch引擎从MulBWD节点检索张量x时，如果它仍然被存储到ssd中，则它在内存中。张量缓存不加载张量，而是通过将弱引用转换为引用返回其内存中的引用，并将获得的引用存储在张量缓存中，以备将来在其他作用域中使用时使用。
[图片]
图6。张量缓存寄存器pack-unpack钩子对用于卸载张量和重新加载张量。(a)显示PyTorch计算图。(b)显示硬件数据路径。(c)和(d)显示了当包或解包钩子被触发时张量缓存的状态。在运算符(1)期间，PyTorch调用包钩子，其中包含要保存用于向后传播的张量，并在计算图(2)上注册返回值。张量缓存跟踪张量，卸载它们(3)，并返回张量的标识符。在反向传播的操作符(A)中，PyTorch调用带有标识符的unpack钩子来获取张量(B)。张量缓存阻塞，直到请求的张量被加载到GPU内存(C)。
3.3.3自适应卸载。在这项工作中，我们得到的一个见解是，激活卸载应该以最小化峰值内存使用为目标，以便同一个系统可以容纳具有较大激活的配置，而不会触发内存不足(OOM)错误。在峰值之后卸载张量是没有帮助的。在图7中，黑色曲线是未卸载的内存占用:它说明GPU内存使用在反向传播开始时达到峰值。蓝色曲线显示了卸载时的内存占用，由于正在进行的卸载任务和在反向传播中创建的新中间张量，峰值被延迟。过度的张量卸载可能会保留张量引用，即使在它最后一次在反向传播中使用之后，延迟其内存的回收。为了减少峰值后不必要的卸载，我们设计了具有两个特征的自适应卸载。
[图片]
在表3的系统中，一个A100在有卸载(黑色)和没有卸载(蓝色)的BERT训练步骤中的内存占用。由于张量卸载和重新加载引起的内存释放和分配，使用卸载运行会导致更多的分配器事件。TBA在反向传播开始时减少了45%的内存占用，端到端峰值内存占用减少了25%。
[图片]
首先，当一个线程被分配存储任务时，线程将检查张量是否被转发。如果是，作业将被取消。其次，如图8所示，我们设计了一种算法来选择一个暂停卸载的模块:我们设计了一个步骤来收集(1)每个MLP块和注意力块的数据传输大小和计算时间，(2)前向传播的计算时间、数据传输时间和总数据传输量。假设模块𝑚是一个步骤中要卸载的最后一个模块。所需的数据传输带宽是在模块𝑚开始反向传播之前完成𝑚之前所有模块的卸载，以及模块𝑚的卸载和重新加载。通过估计反向传播时间是正向传播时间的两倍，可以通过收集到的数据数计算出所需的数据传输带宽，并且不应大于所测量的正向传播中的写带宽。
3.4 SSD写容量、带宽、寿命
为了确认我们的设计在大规模训练系统中是否可行，特别是在SSD耐用性和所需带宽方面，我们进行了性能建模，以获得每个训练步骤的前向传播时间和过程中产生的激活量。
我们扩展了性能模型包llm-analysis[39]。为了估计正向传播时间，llm分析将每个transformer层建模为一个简单的管道，𝑡= max´Í𝑙max´𝑡𝑙，𝑐𝑜𝑚𝑝𝑢𝑡𝑒，𝑡𝑙，𝑚𝑒𝑚𝑜𝑟至极，𝑡𝑍𝑒𝑅𝑂，𝑐𝑜𝑚𝑚𝑢𝑛至极，其中，至极表示transformer层内的任何层。当ZeRO使能时，假定ZeRO通信时间与transformer层的非ZeRO计算和内存操作完美地流水线化。
我们将每个GPU所需的PCIe写入带宽建模为激活总量除以训练时间的一半:正如第3.3.3节所解释的，一些激活可能在反向传播的早期阶段写入，以减少所需的PCIe带宽。我们还假设训练步长𝑡𝑠𝑡𝑒𝑝是前向传播时间的三倍。然后将寿命预测为𝑡𝑙𝑆𝑓𝑒=𝑆𝑒𝑛𝑑𝑢𝑟𝑎𝑛𝑐𝑒·𝑡𝑠𝑡𝑒𝑝/𝑆𝑎𝑐𝑡𝑎𝑣𝑎𝑡𝑐𝑜𝑛𝑠，其中𝑆𝑒𝑛𝑑𝑢𝑟𝑎𝑛𝑐𝑒为SSD持久度评级允许的寿命，𝑡𝑣𝑎𝑡𝑜𝑛𝑠是每个训练步骤的激活量。我们在第4节的实验中验证了具有剖面激活尺寸的𝑆𝑎𝑐𝑡‑𝑣𝑎𝑡‑𝑜𝑛𝑠公式。我们假设每个GPU有四个Solidigm D7-P5810 12.8TB(表1)，并假设WAF在JESD评级中为2.5，在我们的场景中为1。
有了这些，我们得到了图9。我们使用Megatron-LM的系统配置和测量的浮点吞吐量[77]。图形处理器为A100 PCIe。其中，预计寿命在3年以上，每个GPU的PCIe写带宽不大于12.1GB/s。此外，当系统尺寸和/或模型尺寸扩大时，所需的PCIe写带宽减少，预计寿命增加。这种影响的发生是因为较大的系统意味着增加的通信开销和降低的计算效率，从而减慢了每个GPU上的训练迭代。
我们还估计了每个GPU在一个步骤中产生的最大激活大小:我们计算最大微批大小，假设同时在GPU内存中的一个软件中只有两层，而所有其他激活都被卸载。那么，在一个步骤中产生的最大激活微批就是卸载可以打开的最大激活，如图9中的菱形标记所示。每个GPU的最大激活大小范围从0.4 TB到1.8 TB，而微批大小范围从8到32。如此大的激活不能再由主内存保存(图2)，因此SSD是卸载目标的唯一选择。
为了进一步提高SSD的耐用性，可以放宽数据保留期:当数据保留期从3年放宽到3天时，NAND闪存将获得50× PE周期[7,33,40]。在本小节的推理中没有使用这种技术，但我们将在第4.4节中讨论它对成本的影响。
4 Evaluation
我们评估TBA并回答以下问题。
q1。TBA在多大程度上隐藏了I/O延迟?
q2。TBA减少了多少峰值内存使用?
q3:给定相同的激活每个gpu内存预算，TBA可以提供多少吞吐量提升?4.2节通过比较TBA和没有TBA的执行来回答Q1和Q2。我们将在4.4节中讨论性能分析、扩展的影响和成本。
我们使用一台带有2× A100 PCIe gpu和7× Intel P5800X ssd的机器，如表3所示。ssd盘被组织成两个RAID0阵列:一个有3个ssd盘，另一个有4个ssd盘。在评估期间，我们测量了带有4个ssd的A100的内存使用情况。每个阵列是其中一个A100 gpu的专用卸载目标。为了保持一致的性能，gpu被锁定在基频。安装了最新的Megatron DeepSpeed[46]，它将DeepSpeed技术整合到Megatron中，并确保互操作性。
我们在三个模型上测量系统预训练性能，BERT[15]作为纯编码器模型，GPT[64]作为纯解码器模型，T5[65]作为编码器-解码器模型。我们使用OSCAR语料库[59,60]作为数据集。
我们使用两个A100 gpu来实现张量并行性。每一步的微批数固定为1，因为没有管道并行性，在每次训练迭代中，在前一个微批的前向传播和后向传播完成之前，MegatronDeepSpeed不会启动一个新的微批。大于1的微批数只会带来梯度积累，而不会对活化卸载模式产生影响。在我们的实验中，隐藏维度从8192到16384，我们使用典型的超参数[15,65,85]来处理这个范围内的隐藏维度。注意头维度为128。文本序列长度为1024。对于T5，解码器的数量是总层数的一半，向下四舍五入。FlashAttention-2[11]是使用或不使用TBA优化的注意力计算。
由于每个A100只有40GB的设备内存，为了探索更接近现实世界中使用A100 80GB及更高版本gpu的训练系统的设计空间[41,77]，我们做了几个缓解措施。首先，我们使用FP16精度代替混合精度，消除了FP32的权重拷贝。其次，我们使用SGD代替adam作为优化器，以减少优化器状态对内存的使用。这两个度量只影响累积操作和梯度更新，因此在执行或不执行TBA时，在训练步长时间和内存使用方面施加恒定的偏差。
4.2 性能和峰值内存使用
[图片]
为了理解TBA对执行时间和峰值内存使用的影响，我们测量了BERT、T5和GPT的步长时间以及前向和后向传播期间的内存峰值。图10比较了使用TBA和不使用TBA的系统收集的度量。对于每个模型，我们收集了三种不同(隐藏维度，层数)的场景:(8192,4)，(12288,3)和(16384,2)。如图所示，TBA在所有情况下几乎没有性能开销。尽管TBA及其优化引入了额外的cpu执行逻辑，但性能比较表明该逻辑不在关键路径上。相反，GPU计算定义了关键路径，CPU的作用主要在于在当前GPU操作完成之前启动新的GPU任务。因此，CPU没有得到充分利用，TBA的额外工作不会导致新任务到达gpu的延迟。在激活的内存使用方面，TBA在这些情况下有效地降低了28%-40%的峰值。
4.3 通过重新计算-卸载-保持(ROK)曲线比较激活放置策略
[图片]
除了将激活保存在GPU内存和激活检查点中之外，TBA还将卸载激活作为一种选项打开到ssd。我们在这里通过绘制重新计算-卸载-保持(ROK)曲线上的运行来比较三种不同的策略。图11显示了两个3层BERT模型训练的ROK曲线，其中一个隐藏维数为12288，另一个隐藏维数为14336。在ROK曲线中，每次训练用一个点表示。x轴是激活内存峰值，y轴是模型吞吐量。模型吞吐量[77]是指在不考虑软件和硬件实现的情况下，训练步骤所涉及的算法计算次数，例如，激活是否重新计算，除以训练步骤时间。在这两种情况下，TBA降低了GPU激活的内存峰值，允许更大的批处理大小来获得更高的吞吐量。给定相同的批大小，TBA卸载获得的吞吐量与激活保存在内存中时的吞吐量相同。与重计算相比，TBA的激活内存峰值更低。与将激活保存在内存中相比，TBA能够在相同的激活内存预算下将批大小增加一倍。或者，人们可以利用TBA来运行更大的模型，或者使用更少的gpu。
除了这三种策略外，在FlashAttention[12]之前，Megatron[35]提出了选择性重计算:注意到在transformer层中，核心注意力模块(图3中整个灰色框)执行的操作与MLP块相比，计算量较少，但产生了较大的中间张量，工作只重新计算了核心注意力模块。当我们采用FlashAttention时，核心注意模块在一个内核中完成，消除了这些中间张量。选择性重计算对激活的性能和峰值内存使用的影响可以忽略不计。
4.4 Discussion
检查建模。为了理解第3.4节中性能模型的准确性，我们将TBA的卸载量与模型估计值进行比较。如表4所示，这两个数字很接近。我们还使用测量训练时间的一半计算所需的PCIe写带宽。如图所示，隐藏维数越大，PCIe写带宽越小。通常情况下，超过60b个参数的模型隐含维数不小于8k[29,85]。BERT模型的PCIe写带宽与3.4节的估计一致。
因此，减少了与SSD访问完全重叠的计算所需的带宽。简而言之，LLM的扩展本质上是一个弱扩展场景，SSD IO延迟在扩展时更容易隐藏。在大规模系统中，TBA允许容纳的较大活化量可以分配以扩大微批数量和/或扩大批大小。例如，管道的并行性会带来装置的空转气泡，这可以通过增加微批次数量来缓解[77]。增加微批大小和增加微批数量所带来的吞吐量提升在某一点上都达到饱和，这就留下了(1)在给定并行配置的情况下分配激活内存的优化策略，以及(2)并行配置和激活内存分配的联合优化问题。我们将不同并行策略和微批大小的详细吞吐量建模留给未来的工作。成本分析。我们研究了在LLM系统中采用TBA卸载相关的SSD成本。为了获得图9所示的续航能力，每个售价为1万美元的A100都与价值6.4万美元的ssd配对。在评估中，我们为2台a100分配了7台英特尔P5800X。虽然P5800X比表1中的型号更贵，但每PBW的价格为10.27亿美元。我们可以通过放宽数据保留期来进一步降低成本到几个百分点。为了为其他数据提供更持久的存储，系统可以将激活卸载限制在专用SSD上，或者使用配备了ZNS标准的硬件[23,83]，将磨损限制在同一SSD上物理块的指定区域内。
5 Related Work
交换和卸载。许多具有卸载能力的LLM系统都是纯推理的[1,36,75]。在推理中，权重和KV-cache永远不会改变，并且可以跨迭代重用;研究人员利用这一点来提高局部性和内存效率。然而，在LLM训练中，权重在每次迭代中更新，并且所有张量在迭代中都是变化的。一些工作利用卸载特性[67]进行训练，但主要是为了在较小的系统中适应较大的模型，以牺牲性能为代价。它们缺乏保持性能的异步数据传输能力。
另一个方向是将数据和相关计算卸载给CPU[30,69,82]。卸载的计算量相对较轻，卸载的数据包括梯度、权重中的稀疏元素等。认识到这个方向，我们的工作是正交的，因为我们通过GDS将激活卸载到ssd，以尽量减少对CPU的干扰。激活用于梯度计算，这是计算密集型的，最好只在gpu上完成。
在大规模采用llm之前，有关于深度学习卸载数据的工作[5,24,61,70,87]。大多数[5]将数据卸载到主内存，而有些[5]启用GPU-SSD数据路径。LLM训练是独一无二的，因为大规模并行性及其对优化器状态、梯度和权重的内存使用的影响是设计空间的基础。TBA自然支持多个gpu。此外，我们还证明了它在集群上的可行性，并引入了ROK曲线来帮助设计选择。另一方面，LLM对计算能力的要求非常高，它刺激了专用硬件的快速发展，例如transformer引擎[55]和分布式框架。这就是为什么我们要确保良好的互操作性。相反，在这个方向上的大多数早期工作都绑定到特定的PyTorch版本或支持选择层的自定义运行时。量化和稀疏性。一些卸载工作使用量化和/或稀疏性来减少I/O大小[1,5,75]。为了减少计算量，已经提出了将参数量化并将稀疏性引入模型的算法[14,18,32,42,92]。混合专家(MoE)[74]就是在这个方向上，因为它将MLP中的令牌到神经元的连接简化为令牌到专家的连接。一些算法引入了结构化稀疏性，如N:M[94]稀疏性和2:4[62]稀疏性。另一方面，有一些框架和专门的内核可以通过量化和/或稀疏性来加速模型[19,20,76,93]。一些核利用专门的硬件，例如，安培张量核[10,50]。这些技术与我们的工作是正交的，可以用来替换模型，并在使用TBA时加速计算。值得注意的是，给定硬件，将计算与PCIe传输完全重叠的重用因子将根据新的数字格式或稀疏访问模式而改变。我们相信TBA的自适应卸载算法有助于优化这些情况下的卸载量。
优化的内核。先前的工作开发了优化的内核来加速LLM[11,12,56]。一些内核使用特殊的硬件[57]。TBA的互操作性确保了它可以轻松地与这些和即将到来的技术一起使用。
6 Conclusion
在LLM训练系统中，激活占据了越来越有限的GPU内存。我们建议TBA通过将激活卸载到ssd来解决这个问题。我们通过建模证明了它在大型系统中的可行性。我们在TBA中加入了直接的GPU-SSD数据路径和良好的互操作性。为了使计算与数据传输完全重叠，TBA具有异步数据传输、张量重复数据删除、转发和自适应卸载功能。评估显示，TBA将激活峰值内存使用减少了47%，开销可以忽略不计。我们引入ROK曲线来展示TBA的卸载相对于重新计算和将激活保存在内存中的优势。
