대규모 언어 모델의 효율적 추론을 위한 계층별 압축 KV 캐시

Layer-Condensed KV Cache for Efficient Inference of Large Language Models

May 17, 2024

저자: Haoyi Wu, Kewei Tu

cs.AI

초록

대용량 메모리 소비는 실제 애플리케이션에서 고처리량 대규모 언어 모델을 배포하는 데 있어 주요한 병목 현상으로 작용해 왔습니다. 많은 수의 파라미터 외에도, 트랜스포머 아키텍처의 어텐션 메커니즘을 위한 키-값(Key-Value, KV) 캐시는 특히 깊은 언어 모델에서 레이어 수가 많을 때 상당한 양의 메모리를 소비합니다. 본 논문에서는 소수의 레이어에 대해서만 KV를 계산하고 캐싱함으로써 메모리 소비를 크게 절약하고 추론 처리량을 향상시키는 새로운 방법을 제안합니다. 대규모 언어 모델에 대한 실험 결과, 우리의 방법은 표준 트랜스포머 대비 최대 26배 높은 처리량을 달성하며, 언어 모델링 및 다운스트림 작업에서 경쟁력 있는 성능을 보였습니다. 또한, 우리의 방법은 기존의 트랜스포머 메모리 절약 기술과 직교적(orthogonal)이므로, 이를 우리 모델과 쉽게 통합하여 추론 효율성을 더욱 개선할 수 있습니다. 우리의 코드는 https://github.com/whyNLP/LCKV에서 확인할 수 있습니다.

English

Huge memory consumption has been a major bottleneck for deploying high-throughput large language models in real-world applications. In addition to the large number of parameters, the key-value (KV) cache for the attention mechanism in the transformer architecture consumes a significant amount of memory, especially when the number of layers is large for deep language models. In this paper, we propose a novel method that only computes and caches the KVs of a small number of layers, thus significantly saving memory consumption and improving inference throughput. Our experiments on large language models show that our method achieves up to 26times higher throughput than standard transformers and competitive performance in language modeling and downstream tasks. In addition, our method is orthogonal to existing transformer memory-saving techniques, so it is straightforward to integrate them with our model, achieving further improvement in inference efficiency. Our code is available at https://github.com/whyNLP/LCKV.