LayerSkip: Habilitando Inferencia con Salida Temprana y Decodificación Auto-Especulativa

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

Resumen

Presentamos LayerSkip, una solución integral para acelerar la inferencia de modelos de lenguaje grandes (LLMs). Primero, durante el entrenamiento aplicamos dropout por capas, con tasas de dropout bajas para las capas iniciales y tasas más altas para las capas posteriores, junto con una pérdida de salida temprana donde todas las capas del transformador comparten la misma salida. Segundo, durante la inferencia, demostramos que esta estrategia de entrenamiento aumenta la precisión de la salida temprana en las capas iniciales, sin añadir capas o módulos auxiliares al modelo. Tercero, presentamos una novedosa solución de decodificación auto-especulativa donde salimos en las capas iniciales y verificamos y corregimos con las capas restantes del modelo. Nuestro enfoque de decodificación auto-especulativa propuesto tiene una huella de memoria menor que otros enfoques de decodificación especulativa y se beneficia del cómputo y activaciones compartidos entre las etapas de borrador y verificación. Realizamos experimentos en diferentes tamaños de modelos Llama con distintos tipos de entrenamiento: preentrenamiento desde cero, preentrenamiento continuo, ajuste fino en un dominio de datos específico y ajuste fino en una tarea específica. Implementamos nuestra solución de inferencia y mostramos aceleraciones de hasta 2.16x en resúmenes de documentos CNN/DM, 1.82x en codificación y 2.0x en la tarea de análisis semántico TOPv2. Publicamos nuestro código y checkpoints en 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.