LayerSkip: Habilitando Inferência com Saída Antecipada e Decodificação Autosspeculativa

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

April 25, 2024

Autores: Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed A Aly, Beidi Chen, Carole-Jean Wu

cs.AI

Resumo

Apresentamos o LayerSkip, uma solução de ponta a ponta para acelerar a inferência de modelos de linguagem de grande escala (LLMs). Primeiro, durante o treinamento, aplicamos o dropout de camadas, com taxas de dropout baixas para as camadas iniciais e taxas mais altas para as camadas posteriores, além de uma função de perda de saída antecipada onde todas as camadas do transformer compartilham a mesma saída. Segundo, durante a inferência, demonstramos que essa abordagem de treinamento aumenta a precisão da saída antecipada nas camadas iniciais, sem adicionar camadas ou módulos auxiliares ao modelo. Terceiro, apresentamos uma nova solução de decodificação auto-especulativa, onde saímos nas camadas iniciais e verificamos e corrigimos com as camadas restantes do modelo. Nossa abordagem proposta de decodificação auto-especulativa tem uma pegada de memória menor do que outras abordagens de decodificação especulativa e se beneficia do compartilhamento de computação e ativações entre as etapas de rascunho e verificação. Realizamos experimentos em diferentes tamanhos de modelos Llama em diferentes tipos de treinamento: pré-treinamento do zero, pré-treinamento contínuo, ajuste fino em domínios de dados específicos e ajuste fino em tarefas específicas. Implementamos nossa solução de inferência e mostramos acelerações de até 2,16x na sumarização de documentos CNN/DM, 1,82x em codificação e 2,0x na tarefa de análise semântica TOPv2. Disponibilizamos nosso código e checkpoints em https://github.com/facebookresearch/LayerSkip.

English

We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. Second, during inference, we show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model. Third, we present a novel self-speculative decoding solution where we exit at early layers and verify and correct with remaining layers of the model. Our proposed self-speculative decoding approach has less memory footprint than other speculative decoding approaches and benefits from shared compute and activations of the draft and verification stages. We run experiments on different Llama model sizes on different types of training: pretraining from scratch, continual pretraining, finetuning on specific data domain, and finetuning on specific task. We implement our inference solution and show speedups of up to 2.16x on summarization for CNN/DM documents, 1.82x on coding, and 2.0x on TOPv2 semantic parsing task. We open source our code and checkpoints at https://github.com/facebookresearch/LayerSkip.