Practical offloading for fine-tuning LLM on commodity GPU via learned sparse projectors

📄 arXiv: 2406.10181v2 📥 PDF

作者: Siyuan Chen, Zhuofeng Wang, Zelong Guan, Yudong Liu, Phillip B. Gibbons

分类: cs.DC, cs.AI

发布日期: 2024-06-14 (更新: 2025-02-09)

🔗 代码/项目: GITHUB

💡 一句话要点

LSP-Offload：通过学习稀疏投影，在消费级GPU上实现LLM微调的实用卸载框架

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

关键词: 大语言模型 模型微调 GPU卸载 稀疏投影 通信优化

📋 核心要点

  1. 现有LLM微调方法受限于单GPU的内存容量，通常需要将计算和数据从GPU卸载到CPU，但CPU与GPU之间的带宽限制阻碍了性能。
  2. LSP-Offload框架的核心思想是学习高效的稀疏压缩器，减少CPU和GPU之间的通信量，并设计逐层通信调度以提升并行性。
  3. 实验结果表明，LSP-Offload框架在消费级GPU上实现了高效的LLM微调，相比现有方法，显著降低了端到端微调时间。

📝 摘要（中文）

本文提出了一种名为LSP-Offload的卸载框架，旨在通过学习到的稀疏投影器，在消费级硬件上实现接近原生速度的LLM微调。该框架采用数据驱动方法，学习高效的稀疏压缩器，以最小的精度损失最小化通信量。此外，还引入了一种新颖的逐层通信调度，以最大限度地提高通信和计算之间的并行性。实验结果表明，该框架可以在4GB笔记本电脑GPU上微调一个13亿参数的模型，在24GB NVIDIA RTX 4090 GPU上微调一个67亿参数的模型。与最先进的卸载框架相比，在收敛到相同精度时，该方法可将端到端微调时间减少33.1%-62.5%。该框架已开源。

🔬 方法详解

问题定义：论文旨在解决在消费级GPU上微调大型语言模型（LLM）时，由于GPU内存限制和CPU-GPU通信瓶颈导致的性能问题。现有方法如模型并行和卸载计算虽然可以解决内存问题，但要么需要多GPU环境，要么受到CPU-GPU带宽的限制，导致微调速度显著下降。

核心思路：论文的核心思路是利用学习到的稀疏投影器来压缩需要从GPU卸载到CPU的数据，从而减少CPU-GPU之间的通信量。通过数据驱动的方式学习这些稀疏投影器，可以在保证精度损失最小化的前提下，显著降低通信开销。此外，论文还设计了一种逐层通信调度策略，以最大化通信和计算之间的并行性。

技术框架：LSP-Offload框架主要包含以下几个阶段：1) 稀疏投影器学习：使用数据驱动的方法学习稀疏投影矩阵，用于压缩从GPU卸载到CPU的数据。2) 数据卸载与压缩：将部分模型参数或激活值从GPU卸载到CPU，并使用学习到的稀疏投影器进行压缩。3) CPU计算：在CPU上进行部分计算，例如梯度计算。4) 数据上传与解压缩：将计算结果从CPU上传回GPU，并进行解压缩。5) 逐层通信调度：设计一种逐层通信调度策略，以最大化通信和计算之间的并行性。

关键创新：论文的关键创新在于：1) 学习到的稀疏投影器：通过数据驱动的方式学习稀疏投影矩阵，可以有效地压缩数据，减少CPU-GPU通信量，同时保证精度损失最小化。2) 逐层通信调度：通过精心设计的通信调度策略，可以最大化通信和计算之间的并行性，进一步提升微调速度。

关键设计：稀疏投影器的学习过程采用了一种基于梯度下降的优化方法，目标是最小化压缩和解压缩过程中的精度损失。逐层通信调度策略根据每一层的计算量和通信量，动态地调整通信和计算的顺序，以最大化并行性。具体的参数设置包括稀疏投影器的稀疏度、学习率、以及通信调度策略中的时间片大小等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LSP-Offload框架可以在4GB笔记本电脑GPU上微调一个13亿参数的模型，在24GB NVIDIA RTX 4090 GPU上微调一个67亿参数的模型。与现有最先进的卸载框架相比，在收敛到相同精度时，LSP-Offload可以将端到端微调时间减少33.1%-62.5%。

🎯 应用场景

该研究成果可广泛应用于资源受限的场景下的大型语言模型微调，例如在个人电脑、笔记本电脑等消费级设备上进行LLM的个性化定制和领域适配。这降低了LLM微调的硬件门槛，使得更多开发者和研究人员能够参与到LLM的应用开发中，加速LLM在各个领域的普及和应用。

📄 摘要（原文）

Fine-tuning large language models (LLMs) requires significant memory, often exceeding the capacity of a single GPU. A common solution to this memory challenge is offloading compute and data from the GPU to the CPU. However, this approach is hampered by the limited bandwidth of commodity hardware, which constrains communication between the CPU and GPU, and by slower matrix multiplications on the CPU. In this paper, we present an offloading framework, LSP-Offload, that enables near-native speed LLM fine-tuning on commodity hardware through learned sparse projectors. Our data-driven approach involves learning efficient sparse compressors that minimize communication with minimal precision loss. Additionally, we introduce a novel layer-wise communication schedule to maximize parallelism between communication and computation. As a result, our framework can fine-tune a 1.3 billion parameter model on a 4GB laptop GPU and a 6.7 billion parameter model on a 24GB NVIDIA RTX 4090 GPU. Compared to state-of-the-art offloading frameworks, our approach reduces end-to-end fine-tuning time by 33.1%-62.5% when converging to the same accuracy. We open source our framework at https://github.com/gulang2019/LSP-Offload.