Efficient LLM Inference using Dynamic Input Pruning and Cache-Aware Masking

📄 arXiv: 2412.01380v2 📥 PDF

作者: Marco Federici, Davide Belli, Mart van Baalen, Amir Jalalirad, Andrii Skliar, Bence Major, Markus Nagel, Paul Whatmough

分类: cs.LG, cs.CL

发布日期: 2024-12-02 (更新: 2025-04-03)

备注: Main Text: 10 pages, 11 figures. Appendix: 6 pages, 3 figures

🔗 代码/项目: GITHUB

💡 一句话要点

提出动态输入剪枝(DIP)和缓存感知掩码，高效推理移动端LLM。

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

关键词: 大语言模型 移动端推理 动态稀疏化 输入剪枝 缓存优化 LoRA微调 DRAM带宽 边缘计算

📋 核心要点

  1. 现有LLM推理受限于移动设备DRAM带宽瓶颈，传统基于ReLU激活稀疏性的方法在新兴SwiGLU模型上失效。
  2. 提出动态输入剪枝(DIP)，无需预测器即可实现动态稀疏化，并通过LoRA适配器恢复性能，保持精度。
  3. 实验表明，DIP在Phi-3-Medium上实现了显著的内存减少和吞吐量提升，同时保持了可接受的精度损失。

📝 摘要（中文）

移动设备的计算能力不断提升，但DRAM带宽的提升速度却慢得多。这对大型语言模型（LLM）的token生成来说是不利的，因为token生成严重依赖内存带宽。先前的工作提出利用ReLU激活的LLM中固有的动态激活稀疏性来降低每个token的有效DRAM带宽。然而，更新的LLM使用SwiGLU代替ReLU，导致固有的稀疏性很小。虽然可以基于幅度对SwiGLU激活进行剪枝，但由此产生的稀疏模式难以预测，使得先前的方法失效。为了解决这个问题，本文提出了动态输入剪枝（DIP）：一种无需预测器的动态稀疏化方法，它以最小的微调来保持精度。DIP可以进一步使用轻量级的LoRA适配器来恢复在稀疏化过程中损失的一些性能。最后，本文描述了一种新颖的缓存感知掩码策略，该策略考虑缓存状态和激活幅度，以进一步提高缓存命中率，从而提高移动设备上的LLM token速率。在模拟硬件设置中，DIP在精度、内存和吞吐量权衡方面优于其他方法。在Phi-3-Medium上，与从Flash流式传输密集模型相比，DIP实现了46%的内存减少和40%的吞吐量增加，且困惑度损失小于0.1。本文中硬件模拟器、方法和实验的开源代码可在https://github.com/Qualcomm-AI-research/dynamic-sparsity 获得。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在移动设备上部署时，由于DRAM带宽限制导致的推理效率低下问题。现有方法，如基于ReLU激活稀疏性的剪枝，在采用SwiGLU激活函数的现代LLM上效果不佳，因为SwiGLU激活的固有稀疏性很低，且剪枝后的稀疏模式难以预测。

核心思路：论文的核心思路是通过动态输入剪枝（DIP）来实现LLM的稀疏化，从而降低内存带宽需求并提高推理速度。DIP的关键在于它是一种无需预测器的动态稀疏化方法，这意味着它不需要预先训练的预测器来预测哪些激活可以被安全地剪枝。这种设计避免了预测器带来的额外计算开销和复杂性。

技术框架：DIP方法主要包含以下几个阶段：1) 动态输入剪枝：根据激活幅度动态地选择要保留的激活子集，从而实现稀疏化。2) LoRA适配器微调：使用轻量级的LoRA适配器对剪枝后的模型进行微调，以恢复在稀疏化过程中损失的精度。3) 缓存感知掩码：利用缓存状态和激活幅度信息，设计一种掩码策略，以提高缓存命中率。

关键创新：论文的关键创新在于提出了无需预测器的动态输入剪枝（DIP）方法。与现有方法相比，DIP不需要训练额外的预测器，从而降低了计算复杂度和内存开销。此外，缓存感知掩码策略也是一个创新点，它能够充分利用缓存信息来提高缓存命中率，从而进一步提高推理速度。

关键设计：DIP的关键设计包括：1) 动态剪枝阈值：根据激活幅度的分布动态地调整剪枝阈值，以实现所需的稀疏度。2) LoRA适配器配置：选择合适的LoRA适配器大小和训练参数，以在精度和计算开销之间取得平衡。3) 缓存感知掩码策略：根据缓存状态和激活幅度信息，设计一种掩码策略，以优先保留可能被缓存命中的激活。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DIP在Phi-3-Medium模型上实现了显著的性能提升。与从Flash流式传输密集模型相比，DIP实现了46%的内存减少和40%的吞吐量增加，同时困惑度损失小于0.1。这些结果表明DIP是一种有效的LLM推理加速方法，尤其适用于资源受限的移动设备。

🎯 应用场景

该研究成果可应用于移动设备上的LLM部署，例如智能手机、平板电脑等。通过降低内存带宽需求和提高推理速度，可以实现更流畅、更高效的本地LLM应用体验，例如离线对话、文本生成、图像处理等。该技术还有潜力扩展到其他资源受限的边缘计算设备。

📄 摘要（原文）

While mobile devices provide ever more compute power, improvements in DRAM bandwidth are much slower. This is unfortunate for large language model (LLM) token generation, which is heavily memory-bound. Previous work has proposed to leverage natural dynamic activation sparsity in ReLU-activated LLMs to reduce effective DRAM bandwidth per token. However, more recent LLMs use SwiGLU instead of ReLU, which results in little inherent sparsity. While SwiGLU activations can be pruned based on magnitude, the resulting sparsity patterns are difficult to predict, rendering previous approaches ineffective. To circumvent this issue, our work introduces Dynamic Input Pruning (DIP): a predictor-free dynamic sparsification approach, which preserves accuracy with minimal fine-tuning. DIP can further use lightweight LoRA adapters to regain some performance lost during sparsification. Lastly, we describe a novel cache-aware masking strategy, which considers the cache state and activation magnitude to further increase cache hit rate, improving LLM token rate on mobile devices. DIP outperforms other methods in terms of accuracy, memory and throughput trade-offs across simulated hardware settings. On Phi-3-Medium, DIP achieves a 46\% reduction in memory and 40\% increase in throughput with $<$ 0.1 loss in perplexity when compared to streaming the dense model from Flash. The open source code for HW simulator, methods, and experiments in this paper is available at https://github.com/Qualcomm-AI-research/dynamic-sparsity .