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.
翻译:我们提出LayerSkip,一种用于加速大型语言模型(LLM)推理的端到端解决方案。首先,在训练阶段,我们应用层丢弃法,对较早层采用低丢弃率,对较晚层采用高丢弃率,并引入早期退出损失,使所有Transformer层共享相同的退出点。其次,在推理阶段,我们发现这种训练方案能提高模型较早层的早期退出准确性,且无需添加任何辅助层或模块。第三,我们提出一种新颖的自推测解码方案:在早期层退出,并通过模型剩余层进行验证与修正。与其它推测解码方法相比,我们提出的自推测解码方法内存占用更小,并能利用草稿与验证阶段的共享计算与激活值。我们在不同Llama模型规模上开展了实验,涵盖四种训练类型:从头预训练、持续预训练、特定数据域微调,以及特定任务微调。我们实现了推理方案,在CNN/DM文档摘要任务上获得高达2.16倍加速,代码生成任务上1.82倍加速,TOPv2语义解析任务上2.0倍加速。