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

📄 arXiv: 2403.06664v1 📥 PDF

作者: Hongsun Jang, Jaeyong Song, Jaewon Jung, Jaeyoung Park, Youngsok Kim, Jinho Lee

分类: cs.AR, cs.LG

发布日期: 2024-03-11

备注: Published at HPCA 2024 (Best Paper Award Honorable Mention)

💡 一句话要点

提出Smart-Infinity以解决存储带宽瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 存储卸载训练 近存储处理 参数更新 梯度压缩 PyTorch集成 高效数据传输

📋 核心要点

  1. 现有的存储卸载训练方法面临存储带宽瓶颈，限制了大型语言模型的训练效率。
  2. Smart-Infinity通过在近存储处理设备上执行参数更新，减少了存储流量并提高了训练速度。
  3. 实验结果表明，Smart-Infinity相比基线实现了显著的加速，且已集成到PyTorch中，便于实际应用。

📝 摘要（中文）

近年来，大型语言模型（LLMs）的快速发展主要得益于参数数量的增加，这导致了巨大的内存容量需求，通常需要数十个GPU来满足。然而，存储卸载训练的方案虽然扩展了内存层次，但也带来了存储带宽瓶颈。本文提出的Smart-Infinity通过在真实系统中使用近存储处理设备来解决这一瓶颈。Smart-Infinity的核心组件SmartUpdate在定制的近存储加速器上执行参数更新，从而减少了大部分存储流量。此外，本文还提出了一种高效的数据传输处理结构，允许在固定内存消耗下重用设备缓冲区，实现数据传输的重叠。最后，通过加速器辅助的梯度压缩/解压缩，进一步提升了系统的可扩展性。Smart-Infinity在与基线比较时实现了显著的加速，并已完全集成到PyTorch中，未来将开源以促进使用。

🔬 方法详解

问题定义：本文旨在解决大型语言模型训练中的存储带宽瓶颈问题。现有的存储卸载训练方法虽然扩展了内存，但由于存储设备的带宽远低于GPU，导致性能受限。

核心思路：Smart-Infinity的核心思路是将参数更新移至存储侧，利用近存储处理设备来执行这些更新，从而显著减少存储流量。通过这种方式，训练过程中的存储带宽需求得以降低。

技术框架：Smart-Infinity的整体架构包括SmartUpdate模块和高效的数据传输处理结构。SmartUpdate在近存储加速器上执行参数更新，而数据传输处理结构则通过重用设备缓冲区实现数据传输的重叠。

关键创新：Smart-Infinity的主要创新在于将参数更新从GPU转移到近存储加速器，显著降低了存储流量。此外，采用加速器辅助的梯度压缩/解压缩技术，进一步提升了系统的可扩展性。

关键设计：在设计中，Smart-Infinity采用了固定内存消耗的重用缓冲区策略，并通过压缩梯度来减少共享通道的写入流量。这些设计使得系统在多设备扩展时仍能保持高效性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Smart-Infinity在与基线相比时实现了显著的加速，具体提升幅度未知。通过减少存储流量和优化数据传输，该方法在真实系统中展现了良好的性能，证明了其在大型语言模型训练中的有效性。

🎯 应用场景

Smart-Infinity的研究成果在大型语言模型的训练中具有广泛的应用潜力，尤其是在需要处理大规模数据集的自然语言处理任务中。通过提高训练效率，该方法可以加速模型的迭代与优化，推动AI技术的快速发展。未来，Smart-Infinity的开源将使更多研究者和开发者能够利用这一技术，促进相关领域的创新。

📄 摘要（原文）

The recent huge advance of Large Language Models (LLMs) is mainly driven by the increase in the number of parameters. This has led to substantial memory capacity requirements, necessitating the use of dozens of GPUs just to meet the capacity. One popular solution to this is storage-offloaded training, which uses host memory and storage as an extended memory hierarchy. However, this obviously comes at the cost of storage bandwidth bottleneck because storage devices have orders of magnitude lower bandwidth compared to that of GPU device memories. Our work, Smart-Infinity, addresses the storage bandwidth bottleneck of storage-offloaded LLM training using near-storage processing devices on a real system. The main component of Smart-Infinity is SmartUpdate, which performs parameter updates on custom near-storage accelerators. We identify that moving parameter updates to the storage side removes most of the storage traffic. In addition, we propose an efficient data transfer handler structure to address the system integration issues for Smart-Infinity. The handler allows overlapping data transfers with fixed memory consumption by reusing the device buffer. Lastly, we propose accelerator-assisted gradient compression/decompression to enhance the scalability of Smart-Infinity. When scaling to multiple near-storage processing devices, the write traffic on the shared channel becomes the bottleneck. To alleviate this, we compress the gradients on the GPU and decompress them on the accelerators. It provides further acceleration from reduced traffic. As a result, Smart-Infinity achieves a significant speedup compared to the baseline. Notably, Smart-Infinity is a ready-to-use approach that is fully integrated into PyTorch on a real system. We will open-source Smart-Infinity to facilitate its use.