Enabling High-Sparsity Foundational Llama Models with Efficient Pretraining and Deployment

📄 arXiv: 2405.03594v1 📥 PDF

作者: Abhinav Agarwalla, Abhay Gupta, Alexandre Marques, Shubhra Pandit, Michael Goin, Eldar Kurtic, Kevin Leong, Tuan Nguyen, Mahmoud Salem, Dan Alistarh, Sean Lie, Mark Kurtz

分类: cs.CL, cs.AI

发布日期: 2024-05-06

💡 一句话要点

提出高效稀疏预训练与部署方法，实现高稀疏度LLaMA模型加速与精度保持

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

关键词: 大型语言模型 稀疏性 剪枝 预训练 推理加速 LLaMA 模型优化

📋 核心要点

  1. 大型语言模型计算成本高昂，限制了其应用。本文旨在解决LLM的计算瓶颈问题。
  2. 结合SparseGPT剪枝和稀疏预训练，在保证精度的前提下，显著提高LLM的稀疏度。
  3. 实验表明，该方法在CPU和GPU上均能实现显著的推理加速，且训练过程也得到加速。

📝 摘要（中文）

本文提出了一种新颖的方法，旨在创建高性能LLM的精确、稀疏的基础版本，在高达70%的稀疏度下，实现微调任务的完全精度恢复。该方法针对LLaMA-2 7B模型，结合了SparseGPT单次剪枝方法，以及在SlimPajama数据集子集与The Stack数据集的Python子集上对稀疏模型进行预训练。在Cerebras CS-3芯片上，观察到由稀疏性带来的训练加速，与理论缩放相符。此外，通过Neural Magic的DeepSparse引擎在CPU上实现了高达3倍的推理加速，通过Neural Magic的nm-vllm引擎在GPU上实现了1.7倍的加速。上述增益仅通过稀疏性实现，因此可以通过量化进一步提升性能。具体而言，稀疏量化LLaMA模型在CPU上的总加速高达8.6倍。在聊天、指令跟随、代码生成、算术推理和摘要等多种具有挑战性的任务中验证了这些结果，证明了其通用性。这项工作为快速创建更小、更快的LLM铺平了道路，且不牺牲精度。

🔬 方法详解

问题定义：大型语言模型（LLM）虽然在自然语言处理领域取得了显著进展，但其庞大的规模带来了巨大的计算负担，包括训练和推理。现有的方法难以在保持模型精度的前提下，有效降低LLM的计算成本，尤其是在资源受限的场景下。因此，如何创建更小、更快且精度不下降的LLM是一个关键问题。

核心思路：本文的核心思路是通过引入稀疏性来降低LLM的计算复杂度。具体而言，首先使用SparseGPT进行一次性剪枝，移除模型中不重要的连接，从而实现模型的稀疏化。然后，对稀疏模型进行预训练，以恢复剪枝带来的精度损失。通过稀疏化和预训练的结合，可以在保持模型精度的前提下，显著降低模型的计算量。

技术框架：该方法主要包含两个阶段：稀疏化和预训练。在稀疏化阶段，使用SparseGPT对LLaMA-2 7B模型进行一次性剪枝，移除模型中不重要的连接。在预训练阶段，使用SlimPajama数据集的子集和The Stack数据集的Python子集对稀疏模型进行预训练，以恢复剪枝带来的精度损失。在推理阶段，使用Neural Magic的DeepSparse引擎和nm-vllm引擎对稀疏模型进行加速。

关键创新：本文的关键创新在于结合了SparseGPT剪枝和稀疏预训练，从而在保持模型精度的前提下，显著提高了LLM的稀疏度。与传统的剪枝方法相比，SparseGPT具有高效性，可以快速地对模型进行稀疏化。与传统的预训练方法相比，稀疏预训练可以更好地适应稀疏模型的特点，从而更有效地恢复精度。

关键设计：在稀疏化阶段，使用SparseGPT对LLaMA-2 7B模型进行剪枝，目标稀疏度高达70%。在预训练阶段，使用SlimPajama数据集的子集和The Stack数据集的Python子集对稀疏模型进行预训练。数据集的选择和混合比例对最终的性能有重要影响。在推理阶段，使用Neural Magic的DeepSparse引擎和nm-vllm引擎对稀疏模型进行加速，这些引擎针对稀疏模型进行了优化。

📊 实验亮点

实验结果表明，该方法可以在高达70%的稀疏度下，实现微调任务的完全精度恢复。在CPU上，通过Neural Magic的DeepSparse引擎实现了高达3倍的推理加速，通过Neural Magic的nm-vllm引擎在GPU上实现了1.7倍的加速。稀疏量化LLaMA模型在CPU上的总加速高达8.6倍。这些结果表明，该方法可以有效地降低LLM的计算成本，同时保持模型的精度。

🎯 应用场景

该研究成果可广泛应用于资源受限的场景，例如移动设备、边缘计算和嵌入式系统。通过降低LLM的计算成本，可以使其更容易部署在这些平台上，从而实现更广泛的应用。此外，该方法还可以用于加速LLM的训练和推理，从而提高开发效率和用户体验。未来，该方法有望成为构建高效LLM的重要技术手段。

📄 摘要（原文）

Large language models (LLMs) have revolutionized Natural Language Processing (NLP), but their size creates computational bottlenecks. We introduce a novel approach to create accurate, sparse foundational versions of performant LLMs that achieve full accuracy recovery for fine-tuning tasks at up to 70% sparsity. We achieve this for the LLaMA-2 7B model by combining the SparseGPT one-shot pruning method and sparse pretraining of those models on a subset of the SlimPajama dataset mixed with a Python subset of The Stack dataset. We exhibit training acceleration due to sparsity on Cerebras CS-3 chips that closely matches theoretical scaling. In addition, we establish inference acceleration of up to 3x on CPUs by utilizing Neural Magic's DeepSparse engine and 1.7x on GPUs through Neural Magic's nm-vllm engine. The above gains are realized via sparsity alone, thus enabling further gains through additional use of quantization. Specifically, we show a total speedup on CPUs for sparse-quantized LLaMA models of up to 8.6x. We demonstrate these results across diverse, challenging tasks, including chat, instruction following, code generation, arithmetic reasoning, and summarization to prove their generality. This work paves the way for rapidly creating smaller and faster LLMs without sacrificing accuracy.