【论文阅读】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. 激活所需的内存比任何其他内存使用增长得更快
   

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比分层完全重计算获得更好的内存节省，同时保留了将激活保存在内存中的性能。