LayerSkip: Ermöglichen von frühzeitigem Beenden der Inferenz und selbstspekulatives Decodieren

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

April 25, 2024

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

Zusammenfassung

Wir präsentieren LayerSkip, eine End-to-End-Lösung zur Beschleunigung der Inferenz großer Sprachmodelle (LLMs). Während des Trainings wenden wir zunächst Layer-Dropout an, mit niedrigen Dropout-Raten für frühere Schichten und höheren Dropout-Raten für spätere Schichten, sowie einen Early-Exit-Verlust, bei dem alle Transformer-Schichten denselben Exit teilen. Zweitens zeigen wir während der Inferenz, dass dieses Trainingsrezept die Genauigkeit des Early-Exits in früheren Schichten erhöht, ohne dem Modell zusätzliche Schichten oder Module hinzuzufügen. Drittens präsentieren wir eine neuartige selbstspekulative Decodierungslösung, bei der wir in frühen Schichten aussteigen und mit den verbleibenden Schichten des Modells überprüfen und korrigieren. Unser vorgeschlagener selbstspekulativer Decodierungsansatz hat einen geringeren Speicherbedarf als andere spekulative Decodierungsansätze und profitiert von gemeinsam genutzten Berechnungen und Aktivierungen der Entwurfs- und Verifikationsphasen. Wir führen Experimente mit verschiedenen Llama-Modellgrößen und verschiedenen Arten des Trainings durch: Pretraining von Grund auf, kontinuierliches Pretraining, Feinabstimmung auf spezifische Datendomänen und Feinabstimmung auf spezifische Aufgaben. Wir implementieren unsere Inferenzlösung und zeigen Beschleunigungen von bis zu 2,16-fach bei der Zusammenfassung von CNN/DM-Dokumenten, 1,82-fach beim Codieren und 2,0-fach bei der TOPv2-semantischen Parsing-Aufgabe. Wir stellen unseren Code und Checkpoints unter https://github.com/facebookresearch/LayerSkip als Open Source zur Verfügung.

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.