SkipDecode: Decodificación Autoregresiva con Omisión mediante Agrupación y Almacenamiento en Caché para Inferencia Eficiente en Modelos de Lenguaje Grande

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

July 5, 2023

Autores: Luciano Del Corro, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, Subhabrata Mukherjee

cs.AI

Resumen

Los modelos de lenguaje autoregresivos de gran escala (LLMs, por sus siglas en inglés) han logrado avances notables en diversas tareas de generación de lenguaje natural. Sin embargo, incurren en altos costos computacionales y latencia debido a la generación token por token de manera autoregresiva. Para abordar este problema, se han propuesto varios enfoques para reducir el costo computacional utilizando estrategias de salida temprana. Estas estrategias permiten una generación de texto más rápida al utilizar un cómputo reducido sin aplicar el grafo computacional completo a cada token. Si bien los métodos existentes de salida temprana a nivel de token muestran resultados prometedores para la inferencia en línea, no pueden aplicarse fácilmente para la inferencia por lotes y el almacenamiento en caché de claves-valores (KV). Esto se debe a que deben esperar hasta que el último token en un lote salga antes de poder detener el cómputo, lo que limita severamente la aplicación práctica de tales técnicas. En este artículo, proponemos un método simple y efectivo de salida temprana a nivel de token, llamado SkipDecode, diseñado para funcionar de manera fluida con la inferencia por lotes y el almacenamiento en caché KV. Este método supera las limitaciones anteriores al establecer un punto de salida único para cada token en un lote en cada posición de la secuencia. Además, garantiza una disminución monótona en los puntos de salida, eliminando así la necesidad de recalcular las cachés KV para los tokens anteriores. En lugar de terminar el cómputo prematuramente como en trabajos anteriores, nuestro enfoque omite las capas inferiores y medias, dedicando la mayor parte de los recursos computacionales a las capas superiores, permitiendo que los tokens posteriores se beneficien del cómputo invertido en los tokens anteriores. Nuestros resultados experimentales muestran que SkipDecode puede obtener aceleraciones en la inferencia de 2x a 5x con una regresión insignificante en una variedad de tareas. Esto se logra utilizando modelos OPT de 1.3 mil millones y 6.7 mil millones de parámetros, siendo además directamente compatible con técnicas de optimización de lotes y almacenamiento en caché 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.