【论文阅读】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持续成长的瓶颈。如下图所示
- GPU内存容量的增长速度比LLM尺寸扩展速度和GPU FP16吞吐量的提高速度慢60%左右。
- 最近的llm训练中,80%的GPU内存使用由激活组成。
- 激活所需的内存比任何其他内存使用增长得更快
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比分层完全重计算获得更好的内存节省,同时保留了将激活保存在内存中的性能。
本文来自博客园,作者:O_fly_O,转载请注明原文链接:https://www.cnblogs.com/world-explorer/p/18378703