We present EE-LLM, a framework for large-scale training and inference of early-exit large language models (LLMs). While recent works have shown preliminary evidence for the efficacy of early exiting in accelerating LLM inference, EE-LLM makes a foundational step towards scaling up early-exit LLMs by supporting their training and inference with massive 3D parallelism. Built upon Megatron-LM, EE-LLM implements a variety of algorithmic innovations and performance optimizations tailored to early exiting, including a lightweight method that facilitates backpropagation for the early-exit training objective with pipeline parallelism, techniques of leveraging idle resources in the original pipeline schedule for computation related to early-exit layers, and two approaches of early-exit inference that are compatible with KV caching for autoregressive generation. Our analytical and empirical study shows that EE-LLM achieves great training efficiency with negligible computational overhead compared to standard LLM training, as well as outstanding inference speedup without compromising output quality. To facilitate further research and adoption, we release EE-LLM at https://github.com/pan-x-c/EE-LLM.

翻译：我们提出EE-LLM框架，用于支持早期退出大语言模型（LLM）的大规模训练与推理。尽管近期研究已初步证明早期退出在加速LLM推理中的有效性，但EE-LLM通过支持基于大规模3D并行的训练与推理，向早期退出LLM的规模化应用迈出了奠基性一步。基于Megatron-LM框架，EE-LLM实现了针对早期退出场景的一系列算法创新与性能优化，包括：一种轻量化方法，支持在流水线并行下完成早期退出训练目标的反向传播；利用原始流水线调度中的闲置资源执行早期退出层相关计算的技术；以及两种兼容自回归生成中KV缓存的早期退出推理方案。我们的分析及实证研究表明，EE-LLM在保持标准LLM训练效率的同时，仅引入可忽略的计算开销，并能在不降低输出质量的情况下实现显著的推理加速。为促进后续研究与推广应用，我们已在https://github.com/pan-x-c/EE-LLM 开源EE-LLM。