MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models
作者: Elias Frantar, Roberto L. Castro, Jiale Chen, Torsten Hoefler, Dan Alistarh
分类: cs.LG
发布日期: 2024-08-21
💡 一句话要点
MARLIN：在大语言模型上实现混合精度自回归并行推理，提升批量推理效率。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 批量推理 权重量化 混合精度 GPU加速
📋 核心要点
- 现有大语言模型批量推理中，权重量化虽能加速单用户推理，但多用户并行场景下的加速效果不明确，GPU内核设计面临挑战。
- MARLIN通过混合精度量化和优化的线性内核设计，结合异步内存访问、任务调度和流水线等技术，提升批量推理效率。
- 实验表明，MARLIN能有效加速LLM推理，与vLLM集成后可实现高达2.8倍的端到端加速，并可扩展到其他压缩技术。
📝 摘要（中文）
随着大语言模型（LLM）推理成为机器学习应用中的重要工作负载，权重量化已成为高效GPU部署的标准技术。量化不仅减小了模型大小，而且由于减少了内存移动，已被证明可以显著加速单用户推理，且对精度影响较小。然而，在具有多个并行客户端的批量设置中是否也能实现加速仍然未知，这对于实际服务至关重要。目前尚不清楚GPU内核是否可以设计为在支持批量工作负载大幅增加的计算需求的同时，保持实际的内存受限状态。本文通过描述混合精度自回归线性内核（称为MARLIN）的设计，积极地解决了这个问题。具体而言，对于权重通过量化压缩到例如每个元素4位的模型，MARLIN表明，可以支持高达16-32的批量大小，并接近最大（4倍）的量化加速，以及高达64-128的更大批量大小，且加速逐渐降低但仍然显著。MARLIN通过异步内存访问、复杂的任务调度和流水线以及定制的量化支持等技术的组合来实现这一点。实验表明，MARLIN在不同场景下对单个LLM层的近乎最优的性能，在与流行的vLLM服务引擎集成时，也可以带来端到端LLM推理加速（高达2.8倍）。最后，MARLIN可以扩展到进一步的压缩技术，如NVIDIA 2:4稀疏性，从而带来额外的加速。
🔬 方法详解
问题定义：论文旨在解决大语言模型（LLM）批量推理效率低下的问题。现有方法在单用户推理中通过权重量化实现了加速，但在实际部署中，需要支持多个并行用户的批量推理场景。现有的GPU内核难以在满足批量推理增加的计算需求的同时，保持内存受限状态，导致量化带来的加速效果不明显。
核心思路：论文的核心思路是设计一种混合精度的自回归线性内核（MARLIN），通过优化内存访问、任务调度和流水线，充分利用量化带来的计算优势，从而在批量推理场景下实现显著的加速。这种设计旨在克服GPU在处理大量并行请求时遇到的内存瓶颈，并最大化硬件利用率。
技术框架：MARLIN的技术框架主要包括以下几个关键模块：1) 混合精度量化：对模型权重进行量化，降低内存占用和计算复杂度。2) 异步内存访问：通过异步方式加载数据，减少内存访问延迟。3) 复杂任务调度：优化任务调度策略，提高GPU利用率。4) 流水线：将计算任务分解为多个阶段，并行执行，提高吞吐量。这些模块协同工作，共同实现高效的批量推理。
关键创新：MARLIN最重要的技术创新点在于其针对批量推理场景的优化设计。与现有方法相比，MARLIN不仅关注单用户推理的加速，更注重在多用户并行场景下如何充分利用量化的优势。通过异步内存访问、复杂的任务调度和流水线等技术，MARLIN能够有效地克服内存瓶颈，提高GPU利用率，从而实现显著的批量推理加速。
关键设计：MARLIN的关键设计包括：1) 混合精度量化方案的选择，需要在精度和性能之间进行权衡。2) 异步内存访问的实现细节，需要考虑数据依赖关系和同步机制。3) 任务调度策略的设计，需要根据GPU的硬件特性进行优化。4) 流水线的划分和调度，需要平衡各个阶段的负载，避免出现瓶颈。
🖼️ 关键图片
📊 实验亮点
实验结果表明，MARLIN在批量推理场景下能够实现显著的加速。对于权重量化到4位的模型，MARLIN可以支持高达16-32的批量大小，并接近4倍的量化加速。与流行的vLLM服务引擎集成后，MARLIN能够实现高达2.8倍的端到端LLM推理加速。此外，MARLIN还能够与NVIDIA 2:4稀疏性等其他压缩技术结合，进一步提升性能。
🎯 应用场景
MARLIN技术可广泛应用于需要高并发、低延迟的大语言模型服务场景，例如在线问答、对话机器人、文本生成等。该研究成果有助于降低LLM部署成本，提升用户体验，并推动LLM在更多实际应用中的落地。未来，该技术有望扩展到其他类型的深度学习模型和硬件平台。
📄 摘要（原文）
As inference on Large Language Models (LLMs) emerges as an important workload in machine learning applications, weight quantization has become a standard technique for efficient GPU deployment. Quantization not only reduces model size, but has also been shown to yield substantial speedups for single-user inference, due to reduced memory movement, with low accuracy impact. Yet, it remains open whether speedups are achievable also in \emph{batched} settings with multiple parallel clients, which are highly relevant for practical serving. It is unclear whether GPU kernels can be designed to remain practically memory-bound, while supporting the substantially increased compute requirements of batched workloads. This paper resolves this question positively by describing the design of Mixed-precision Auto-Regressive LINear kernels, called MARLIN. Concretely, given a model whose weights are compressed via quantization to, e.g., 4 bits per element, MARLIN shows that batchsizes up to 16-32 can be supported with close to maximum ($4\times$) quantization speedup, and larger batchsizes up to 64-128 with gradually decreasing, but still significant, acceleration. MARLIN accomplishes this via a combination of techniques, such as asynchronous memory access, complex task scheduling and pipelining, and bespoke quantization support. Our experiments show that MARLIN's near-optimal performance on individual LLM layers across different scenarios can also lead to end-to-end LLM inference speedups (of up to $2.8\times$) when integrated with the popular vLLM serving engine. Finally, MARLIN is extensible to further compression techniques, like NVIDIA 2:4 sparsity, leading to additional speedups.