摘要

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内存中，而其他激活被刷新，然后在反向传播期间重新计算。对于具有