SkipDecode: Авторегрессивное пропускное декодирование с пакетной обработкой и кэшированием для эффективного вывода в больших языковых моделях

SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference

July 5, 2023

Авторы: Luciano Del Corro, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, Subhabrata Mukherjee

cs.AI

Аннотация

Авторегрессионные большие языковые модели (LLM) достигли значительных успехов в различных задачах генерации естественного языка. Однако они требуют высоких вычислительных затрат и вызывают задержки из-за пошаговой генерации токенов. Для решения этой проблемы было предложено несколько подходов, направленных на снижение вычислительных затрат с использованием стратегий раннего выхода. Эти стратегии позволяют ускорить генерацию текста за счет сокращения вычислений, не применяя полный вычислительный граф к каждому токену. Хотя существующие методы раннего выхода на уровне токенов демонстрируют многообещающие результаты для онлайн-вывода, они не могут быть легко применены для пакетного вывода и кэширования ключей и значений (KV). Это связано с тем, что они должны ждать, пока последний токен в пакете не завершит обработку, прежде чем прекратить вычисления. Это серьезно ограничивает практическое применение таких методов. В данной статье мы предлагаем простой и эффективный метод раннего выхода на уровне токенов, SkipDecode, который разработан для бесперебойной работы с пакетным выводом и KV-кэшированием. Он преодолевает предыдущие ограничения, устанавливая единую точку выхода для каждого токена в пакете на каждой позиции последовательности. Кроме того, он гарантирует монотонное уменьшение точек выхода, тем самым устраняя необходимость пересчета KV-кэшей для предыдущих токенов. В отличие от преждевременного завершения вычислений, как в предыдущих работах, наш подход обходит нижние и средние слои, сосредотачивая большую часть вычислительных ресурсов на верхних слоях, что позволяет последующим токенам использовать вычислительные затраты, сделанные для предыдущих токенов. Наши экспериментальные результаты показывают, что SkipDecode может обеспечить ускорение вывода в 2–5 раз с незначительной потерей качества на различных задачах. Это достигается с использованием моделей OPT с 1,3 миллиарда и 6,7 миллиарда параметров, при этом метод полностью совместим с техниками оптимизации пакетной обработки и KV-кэширования.

English

Autoregressive large language models (LLMs) have made remarkable progress in various natural language generation tasks. However, they incur high computation cost and latency resulting from the autoregressive token-by-token generation. To address this issue, several approaches have been proposed to reduce computational cost using early-exit strategies. These strategies enable faster text generation using reduced computation without applying the full computation graph to each token. While existing token-level early exit methods show promising results for online inference, they cannot be readily applied for batch inferencing and Key-Value caching. This is because they have to wait until the last token in a batch exits before they can stop computing. This severely limits the practical application of such techniques. In this paper, we propose a simple and effective token-level early exit method, SkipDecode, designed to work seamlessly with batch inferencing and KV caching. It overcomes prior constraints by setting up a singular exit point for every token in a batch at each sequence position. It also guarantees a monotonic decrease in exit points, thereby eliminating the need to recompute KV Caches for preceding tokens. Rather than terminating computation prematurely as in prior works, our approach bypasses lower to middle layers, devoting most of the computational resources to upper layers, allowing later tokens to benefit from the compute expenditure by earlier tokens. Our experimental results show that SkipDecode can obtain 2x to 5x inference speedups with negligible regression across a variety of tasks. This is achieved using OPT models of 1.3 billion and 6.7 billion parameters, all the while being directly compatible with batching and KV caching optimization techniques.