LoRA-Switch: Boosting the Efficiency of Dynamic LLM Adapters via System-Algorithm Co-design

📄 arXiv: 2405.17741v1 📥 PDF

作者: Rui Kong, Qiyang Li, Xinyu Fang, Qingtian Feng, Qingfeng He, Yazhu Dong, Weijun Wang, Yuanchun Li, Linghe Kong, Yunxin Liu

分类: cs.AI

发布日期: 2024-05-28

💡 一句话要点

LoRA-Switch：通过系统-算法协同设计提升动态LLM适配器的效率

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

关键词: 低秩适配器 动态适配器 大语言模型 系统-算法协同设计 CUDA内核优化 Token级别路由 推理加速

📋 核心要点

  1. 动态适配器能有效定制和提升LLM，但其碎片化的CUDA内核调用导致推理延迟显著增加。
  2. LoRA-Switch采用token级别的路由机制，通过优化的CUDA内核融合LoRA适配器，提升效率。
  3. 实验表明，LoRA-Switch在保持精度提升的同时，显著降低了解码延迟，提升幅度超过2.4倍。

📝 摘要（中文）

本文提出了一种名为LoRA-Switch的系统-算法协同设计的架构，旨在提升动态LLM适配器的效率。现有的动态适配器，如带有混合专家（MoE）结构的低秩适配器（LoRA），虽然计算复杂度不高，但会导致巨大的推理延迟，使解码速度降低2.5倍以上。通过对动态适配器的细粒度成本分析，发现碎片化的CUDA内核调用是根本原因。LoRA-Switch引入了一种token级别的路由机制，为每个token切换LoRA适配器和权重，并将它们合并到主干网络中进行推理。为了提高效率，这种切换通过优化的CUDA内核实现，该内核一次性融合所有LoRA适配器的合并操作。在常用基准测试中，基于流行的开源LLM的实验表明，该方法在实现与现有动态适配器相似的精度提升的同时，将解码延迟降低了2.4倍以上。

🔬 方法详解

问题定义：现有动态适配器（如LoRA+MoE）虽然计算复杂度不高，但由于其碎片化的CUDA内核调用，导致LLM推理过程中产生巨大的延迟开销，严重影响了解码速度。因此，需要解决的问题是如何在保持动态适配器带来的性能提升的同时，显著降低其推理延迟。

核心思路：LoRA-Switch的核心思路是采用token级别的动态路由机制，并结合优化的CUDA内核，实现高效的LoRA适配器切换和融合。通过token级别的路由，可以更精细地控制每个token使用的适配器，从而提高模型的表达能力。同时，通过优化的CUDA内核，可以一次性融合所有LoRA适配器的合并操作，减少内核调用次数，从而降低延迟。

技术框架：LoRA-Switch的整体架构包括以下几个主要模块：1) Token级别的路由模块：负责为每个token选择合适的LoRA适配器。2) LoRA适配器切换模块：根据路由结果，切换相应的LoRA适配器和权重。3) 融合模块：将选择的LoRA适配器和权重合并到主干网络中。4) 优化的CUDA内核：用于高效地执行LoRA适配器的切换和融合操作。

关键创新：LoRA-Switch的关键创新在于其token级别的路由机制和优化的CUDA内核。与现有的layer-wise或block-wise动态路由不同，LoRA-Switch可以为每个token选择不同的适配器，从而实现更精细的控制。此外，优化的CUDA内核通过融合多个操作，显著减少了内核调用次数，从而降低了延迟。

关键设计：LoRA-Switch的关键设计包括：1) Token路由策略：如何为每个token选择合适的LoRA适配器，例如可以使用一个小型神经网络来预测每个token应该使用的适配器。2) CUDA内核优化：如何高效地实现LoRA适配器的切换和融合操作，例如可以使用CUDA shared memory来减少数据传输，并使用CUDA streams来实现并发执行。3) 损失函数设计：如何训练token路由模块，使其能够选择合适的适配器，例如可以使用交叉熵损失函数来衡量路由结果的准确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LoRA-Switch在保持与现有动态适配器相似的精度提升的同时，将解码延迟降低了2.4倍以上。具体来说，在常用的LLM和基准测试集上，LoRA-Switch的性能优于现有的layer-wise或block-wise动态路由方法，并且能够显著降低推理延迟。

🎯 应用场景

LoRA-Switch可广泛应用于需要高效LLM推理的场景，如实时对话系统、智能客服、机器翻译等。通过降低动态适配器的推理延迟，可以提升用户体验，并降低部署成本。未来，该技术有望进一步扩展到其他类型的动态模型和硬件平台，推动LLM在更多领域的应用。

📄 摘要（原文）

Recent literature has found that an effective method to customize or further improve large language models (LLMs) is to add dynamic adapters, such as low-rank adapters (LoRA) with Mixture-of-Experts (MoE) structures. Though such dynamic adapters incur modest computational complexity, they surprisingly lead to huge inference latency overhead, slowing down the decoding speed by 2.5+ times. In this paper, we analyze the fine-grained costs of the dynamic adapters and find that the fragmented CUDA kernel calls are the root cause. Therefore, we propose LoRA-Switch, a system-algorithm co-designed architecture for efficient dynamic adapters. Unlike most existing dynamic structures that adopt layer-wise or block-wise dynamic routing, LoRA-Switch introduces a token-wise routing mechanism. It switches the LoRA adapters and weights for each token and merges them into the backbone for inference. For efficiency, this switching is implemented with an optimized CUDA kernel, which fuses the merging operations for all LoRA adapters at once. Based on experiments with popular open-source LLMs on common benchmarks, our approach has demonstrated similar accuracy improvement as existing dynamic adapters, while reducing the decoding latency by more than 2.4 times.