LayerSkip: Включение раннего завершения вывода и самоспекулятивного декодирования

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

April 25, 2024

Авторы: 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

Аннотация

Мы представляем LayerSkip, комплексное решение для ускорения вывода больших языковых моделей (LLM). Во-первых, во время обучения мы применяем отсев слоев с низкими значениями отсева для более ранних слоев и более высокими значениями отсева для более поздних слоев, а также потери на раннем выходе, где все трансформерные слои используют одинаковый выход. Во-вторых, во время вывода мы показываем, что данная методика обучения повышает точность раннего выхода на более ранних слоях, не добавляя дополнительных слоев или модулей к модели. В-третьих, мы представляем новое решение самоспекулятивного декодирования, где мы выходим на ранних слоях, а затем проверяем и корректируем с помощью оставшихся слоев модели. Наше предложенное самоспекулятивное декодирование требует меньше памяти, чем другие подобные подходы, и использует общие вычисления и активации этапов черновика и верификации. Мы проводим эксперименты на различных размерах моделей Llama при различных типах обучения: предварительное обучение с нуля, продолжительное предварительное обучение, донастройка на конкретной области данных и донастройка на конкретную задачу. Мы реализуем наше решение вывода и демонстрируем ускорение до 2.16x для суммаризации документов CNN/DM, 1.82x для кодирования и 2.0x для семантического разбора TOPv2. Мы открываем исходный код и контрольные точки по адресу 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.