LayerSkip: Abilitazione dell'Inferenza con Uscita Anticipata e Decodifica Auto-Speculativa
LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
April 25, 2024
Autori: 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
Abstract
Presentiamo LayerSkip, una soluzione end-to-end per accelerare l'inferenza dei grandi modelli linguistici (LLM). In primo luogo, durante l'addestramento applichiamo il dropout a livello di layer, con tassi di dropout bassi per i layer iniziali e più alti per i layer successivi, insieme a una funzione di perdita per uscite anticipate in cui tutti i layer del transformer condividono la stessa uscita. In secondo luogo, durante l'inferenza, dimostriamo che questa procedura di addestramento aumenta l'accuratezza delle uscite anticipate nei layer iniziali, senza aggiungere alcun layer o modulo ausiliario al modello. In terzo luogo, presentiamo una nuova soluzione di decodifica auto-speculativa in cui usciamo dai layer iniziali e verifichiamo e correggiamo con i layer rimanenti del modello. Il nostro approccio di decodifica auto-speculativa proposto ha un'impronta di memoria inferiore rispetto ad altri approcci di decodifica speculativa e beneficia della condivisione del calcolo e delle attivazioni tra le fasi di bozza e verifica. Eseguiamo esperimenti su diverse dimensioni dei modelli Llama con diversi tipi di addestramento: pretraining da zero, pretraining continuo, fine-tuning su un dominio di dati specifico e fine-tuning su un compito specifico. Implementiamo la nostra soluzione di inferenza e dimostriamo accelerazioni fino a 2,16x nella summarizzazione di documenti CNN/DM, 1,82x nel coding e 2,0x nel task di parsing semantico TOPv2. Rendiamo disponibili il nostro codice e i checkpoint all'indirizzo 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.