LayerSkip: Mogelijk maken van vroegtijdige beëindiging bij inferentie en zelf-speculatieve decodering

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

April 25, 2024

Auteurs: 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

Samenvatting

We presenteren LayerSkip, een end-to-end oplossing om de inferentie van grote taalmodelen (LLMs) te versnellen. Ten eerste passen we tijdens de training layer dropout toe, met lage dropout rates voor eerdere lagen en hogere dropout rates voor latere lagen, en een early exit loss waarbij alle transformer-lagen dezelfde exit delen. Ten tweede laten we tijdens de inferentie zien dat deze trainingsmethode de nauwkeurigheid van early exit in eerdere lagen verhoogt, zonder extra hulplagen of modules aan het model toe te voegen. Ten derde presenteren we een nieuwe zelf-speculatieve decodeeroplossing waarbij we vroegtijdig uitstappen in eerdere lagen en verifiëren en corrigeren met de resterende lagen van het model. Onze voorgestelde zelf-speculatieve decodeerbenadering heeft een kleinere geheugenvoetafdruk dan andere speculatieve decodeerbenaderingen en profiteert van gedeelde berekeningen en activaties van de draft- en verificatiefasen. We voeren experimenten uit met verschillende Llama-modelgroottes op verschillende soorten training: pretraining vanaf nul, voortgezette pretraining, finetuning op een specifiek datadomein en finetuning op een specifieke taak. We implementeren onze inferentie-oplossing en tonen snelheidsverbeteringen van tot 2,16x bij samenvatting van CNN/DM-documenten, 1,82x bij codering en 2,0x bij de TOPv2 semantische parsetaak. We maken onze code en checkpoints openbaar op 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.