LayerSkip : Permettant l'inférence à sortie précoce et le décodage auto-spéculatif

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

Résumé

Nous présentons LayerSkip, une solution de bout en bout pour accélérer l'inférence des grands modèles de langage (LLM). Premièrement, lors de l'entraînement, nous appliquons un dropout par couche, avec des taux de dropout faibles pour les premières couches et plus élevés pour les dernières couches, ainsi qu'une perte de sortie précoce où toutes les couches du transformateur partagent la même sortie. Deuxièmement, lors de l'inférence, nous montrons que cette méthode d'entraînement augmente la précision des sorties précoces dans les premières couches, sans ajouter de couches ou modules auxiliaires au modèle. Troisièmement, nous proposons une nouvelle solution de décodage auto-spéculatif où nous sortons aux premières couches et vérifions et corrigons avec les couches restantes du modèle. Notre approche de décodage auto-spéculatif a une empreinte mémoire plus faible que les autres méthodes de décodage spéculatif et bénéficie du partage des calculs et des activations entre les étapes de brouillon et de vérification. Nous menons des expériences sur différentes tailles de modèles Llama avec différents types d'entraînement : pré-entraînement à partir de zéro, pré-entraînement continu, fine-tuning sur un domaine de données spécifique, et fine-tuning sur une tâche spécifique. Nous implémentons notre solution d'inférence et montrons des accélérations allant jusqu'à 2,16x sur la tâche de résumé de documents CNN/DM, 1,82x sur le codage, et 2,0x sur la tâche d'analyse sémantique TOPv2. Nous mettons notre code et nos checkpoints en open source à l'adresse 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.