# EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism

arXiv: https://arxiv.org/abs/2312.04916 · HTML: https://mlsanity.com/paper/2312.04916/

## Key facts

- Cited by 69 papers in the map (data through 14 Oct 2026).
- Placed in Early exit and edge inference, one of 12 directions in Efficient LLM inference and serving.
- Builds on 73 of its 101 references that are in the map.
- First posted on arXiv 8 Dec 2023.

## Placed in

- Efficient LLM inference and serving / [Early exit and edge inference](https://mlsanity.com/inference/early-exit-and-edge-inference/)

## Builds on (most cited first, 30 of 73)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)

## Built on by (most cited first, 30 of 69)

- [LLM Inference Unveiled: Survey and Roofline Model Insights](https://mlsanity.com/paper/2402.16363/) (2024)
- [Mobile Edge Intelligence for Large Language Models: A Contemporary Survey](https://mlsanity.com/paper/2407.18921/) (2024)
- [MiniCache: KV Cache Compression in Depth Dimension for Large Language Models](https://mlsanity.com/paper/2405.14366/) (2024)
- [DeeR-VLA: Dynamic Inference of Multimodal Large Language Models for Efficient Robot Execution](https://mlsanity.com/paper/2411.02359/) (2024)
- [SLEB: Streamlining LLMs through Redundancy Verification and Elimination of Transformer Blocks](https://mlsanity.com/paper/2402.09025/) (2024)
- [Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA](https://mlsanity.com/paper/2410.20672/) (2024)
- [Inner Thinking Transformer: Leveraging Dynamic Depth Scaling to Foster Adaptive Internal Thinking](https://mlsanity.com/paper/2502.13842/) (2025)
- [CE-CoLLM: Efficient and Adaptive Large Language Models Through Cloud-Edge Collaboration](https://mlsanity.com/paper/2411.02829/) (2024)
- [FFN-SkipLLM: A Hidden Gem for Autoregressive Decoding with Adaptive Feed Forward Skipping](https://mlsanity.com/paper/2404.03865/) (2024)
- [Think-at-Hard: Dynamic Looped Transformers for Improved Reasoning](https://mlsanity.com/paper/2511.08577/) (2025)
- [Accelerating Multimodal Large Language Models via Dynamic Visual-Token Exit and the Empirical Findings](https://mlsanity.com/paper/2411.19628/) (2024)
- [The Energy Cost of Reasoning: Analyzing Energy Usage in LLMs with Test-time Compute](https://mlsanity.com/paper/2505.14733/) (2025)
- [EE-Tuning: An Economical yet Scalable Solution for Tuning Early-Exit Large Language Models](https://mlsanity.com/paper/2402.00518/) (2024)
- [Demystifying AI Platform Design for Distributed Inference of Next-Generation LLM models](https://mlsanity.com/paper/2406.01698/) (2024)
- [AdaSkip: Adaptive Sublayer Skipping for Accelerating Long-Context LLM Inference](https://mlsanity.com/paper/2501.02336/) (2025)
- [Certified Self-Consistency: Statistical Guarantees and Test-Time Training for Reliable Reasoning in LLMs](https://mlsanity.com/paper/2510.17472/) (2025)
- [From Prompts to Power: Measuring the Energy Footprint of LLM Inference](https://mlsanity.com/paper/2511.05597/) (2025)
- [Fast and Cost-effective Speculative Edge-Cloud Decoding with Early Exits](https://mlsanity.com/paper/2505.21594/) (2025)
- [SkipGPT: Dynamic Layer Pruning Reinvented with Token Awareness and Module Decoupling](https://mlsanity.com/paper/2506.04179/) (2025)
- [DLO: Dynamic Layer Operation for Efficient Vertical Scaling of LLMs](https://mlsanity.com/paper/2407.11030/) (2024)
- [The Efficiency vs. Accuracy Trade-off: Optimizing RAG-Enhanced LLM Recommender Systems Using Multi-Head Early Exit](https://mlsanity.com/paper/2501.02173/) (2025)
- [From LLMs to LRMs: Rethinking Pruning for Reasoning-Centric Models](https://mlsanity.com/paper/2601.18091/) (2026)
- [AdaPonderLM: Gated Pondering Language Models with Token-Wise Adaptive Depth](https://mlsanity.com/paper/2603.01914/) (2026)
- [Systematic Characterization of LLM Quantization: A Performance, Energy, and Quality Perspective](https://mlsanity.com/paper/2508.16712/) (2025)
- [HELIOS: Adaptive Model And Early-Exit Selection for Efficient LLM Inference Serving](https://mlsanity.com/paper/2504.10724/) (2025)
- [OrthoRank: Token Selection via Sink Token Orthogonality for Efficient LLM inference](https://mlsanity.com/paper/2507.03865/) (2025)
- [Skipping Computations in Multimodal LLMs](https://mlsanity.com/paper/2410.09454/) (2024)
- [ES-dLLM: Efficient Inference for Diffusion Large Language Models by Early-Skipping](https://mlsanity.com/paper/2603.10088/) (2026)
- [Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching](https://mlsanity.com/paper/2608.09444/) (2026)
- [A Comprehensive Survey of Accelerated Generation Techniques in Large Language Models](https://mlsanity.com/paper/2405.13019/) (2024)

## Abstract (arXiv)

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.