LazyLLM: Poda Dinámica de Tokens para una Inferencia Eficiente de LLM de Contexto Largo

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

Autores: Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, Mahyar Najibi

cs.AI

Resumen

La inferencia de los modelos de lenguaje grandes basados en transformadores consta de dos etapas secuenciales: 1) una etapa de precarga para calcular la caché KV de las indicaciones y generar el primer token, y 2) una etapa de decodificación para generar tokens subsecuentes. Para indicaciones largas, la caché KV debe calcularse para todos los tokens durante la etapa de precarga, lo que puede aumentar significativamente el tiempo necesario para generar el primer token. En consecuencia, la etapa de precarga puede convertirse en un cuello de botella en el proceso de generación. Queda abierta la pregunta de si todos los tokens de la indicación son esenciales para generar el primer token. Para responder a esto, presentamos un método novedoso, LazyLLM, que calcula selectivamente el KV para los tokens importantes para la predicción del siguiente token en ambas etapas, tanto la de precarga como la de decodificación. A diferencia de los enfoques estáticos de poda que podan la indicación de una vez, LazyLLM permite a los modelos de lenguaje seleccionar dinámicamente diferentes subconjuntos de tokens del contexto en diferentes pasos de generación, aunque puedan haber sido podados en pasos anteriores. Experimentos exhaustivos en conjuntos de datos estándar en diversas tareas demuestran que LazyLLM es un método genérico que puede integrarse fácilmente con modelos de lenguaje existentes para acelerar significativamente la generación sin necesidad de ajustes finos. Por ejemplo, en la tarea de pregunta-respuesta multi-documento, LazyLLM acelera la etapa de precarga del modelo LLama 2 7B en un 2.34x manteniendo la precisión.

English

The inference of transformer-based large language models consists of two sequential stages: 1) a prefilling stage to compute the KV cache of prompts and generate the first token, and 2) a decoding stage to generate subsequent tokens. For long prompts, the KV cache must be computed for all tokens during the prefilling stage, which can significantly increase the time needed to generate the first token. Consequently, the prefilling stage may become a bottleneck in the generation process. An open question remains whether all prompt tokens are essential for generating the first token. To answer this, we introduce a novel method, LazyLLM, that selectively computes the KV for tokens important for the next token prediction in both the prefilling and decoding stages. Contrary to static pruning approaches that prune the prompt at once, LazyLLM allows language models to dynamically select different subsets of tokens from the context in different generation steps, even though they might be pruned in previous steps. Extensive experiments on standard datasets across various tasks demonstrate that LazyLLM is a generic method that can be seamlessly integrated with existing language models to significantly accelerate the generation without fine-tuning. For instance, in the multi-document question-answering task, LazyLLM accelerates the prefilling stage of the LLama 2 7B model by 2.34x while maintaining accuracy.