Smart-Infinity: Fast Large Language Model Training using Near-Storage Processing on a Real System

最近大型语言模型（LLMs）的巨大进步主要是由于参数数量的增加。这导致了大量的内存容量需求，需要使用数十个GPU才能满足容量需求。其中一个流行的解决方案是存储卸载训练，它使用主机内存和存储作为扩展内存层次结构。然而，这显然会以存储带宽瓶颈为代价，因为存储设备的带宽比GPU设备内存低几个数量级。我们的工作Smart-Infinity，使用真实系统上的近存储处理设备解决了存储卸载LLM训练的存储带宽瓶颈。Smart-Infinity的主要组件是SmartUpdate，它在自定义的近存储加速器上执行参数更新。我们确定将参数更新移动到存储侧可以消除大部分存储流量。此外，我们提出了一种高效的数据传输处理器结构，以解决Smart-Infinity的系统集成问题。该处理器允许重复使用设备缓冲区，以重叠数据传输并固定内存消耗。最后，我们提出了加速器辅助的梯度压缩/解压缩来增强Smart-Infinity的可扩展性。当扩展到多个近存储处理设备时，共享通道上的写入流量成为瓶颈。为了减轻这种情况，我们在GPU上压缩梯度，并在加速器上解压缩。这提供了进一步的加速，减少了流量。因此，与基线相比，Smart-Infinity实现了显着的加速。值得注意的是，Smart-Infinity是一种可立即使用的方法，已经完全集成到PyTorch中的真实系统中。我们将开源Smart-Infinity以促进其使用。

论文旨在解决大规模语言模型训练中存储带宽瓶颈问题，提出了Smart-Infinity方法。

Smart-Infinity方法使用定制的近存储加速器在存储侧执行参数更新，以减少存储流量。同时提出了高效的数据传输处理程序结构和加速器辅助梯度压缩/解压缩来增强可扩展性。

论文在实验中使用了PyTorch框架，提供了开源代码。Smart-Infinity方法在解决存储带宽瓶颈问题上取得了显著的加速效果。

近期相关研究包括使用低精度计算进行模型训练和参数压缩，以及使用分布式训练等方法来提高训练效率。