LazyLLM: Potatura Dinamica dei Token per un'Inferenza Efficiente nei Modelli Linguistici con Contesti Lunghi

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

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

cs.AI

Abstract

L'inferenza dei modelli linguistici di grandi dimensioni basati su trasformatori consiste in due fasi sequenziali: 1) una fase di prefilling per calcolare la cache KV dei prompt e generare il primo token, e 2) una fase di decodifica per generare i token successivi. Per prompt lunghi, la cache KV deve essere calcolata per tutti i token durante la fase di prefilling, il che può aumentare significativamente il tempo necessario per generare il primo token. Di conseguenza, la fase di prefilling può diventare un collo di bottiglia nel processo di generazione. Rimane una questione aperta se tutti i token del prompt siano essenziali per generare il primo token. Per rispondere a questo, introduciamo un metodo innovativo, LazyLLM, che calcola selettivamente la KV per i token importanti per la previsione del token successivo sia nella fase di prefilling che in quella di decodifica. Contrariamente agli approcci di pruning statico che potano il prompt in una sola volta, LazyLLM consente ai modelli linguistici di selezionare dinamicamente diversi sottoinsiemi di token dal contesto in diversi passaggi di generazione, anche se potrebbero essere stati potati in passaggi precedenti. Esperimenti estesi su dataset standard in vari compiti dimostrano che LazyLLM è un metodo generico che può essere integrato senza soluzione di continuità con i modelli linguistici esistenti per accelerare significativamente la generazione senza bisogno di fine-tuning. Ad esempio, nel compito di risposta a domande su più documenti, LazyLLM accelera la fase di prefilling del modello LLama 2 7B di 2,34x mantenendo l'accuratezza.

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.