Research map / Large language models / Reasoning and chain of thought
EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Presents EE-LLM, a Megatron-LM-based framework for training and inference of early-exit LLMs with 3D parallelism, including an auxiliary-loss method for backpropagating multi-exit losses across pipeline stages and two KV-cache-compatible early-exit inference methods.
Problem
Early-exit LLMs have only been trained at small scale, and standard pipeline parallelism cannot backpropagate losses located at multiple exits on different stages; early exiting also conflicts with KV caching during autoregressive generation.
Approach
Rewrites the multi-exit objective as a sum of per-stage losses and defines per-stage auxiliary losses with a linear term <g_i, x_i> carrying downstream gradients, so standard 1F1B P2P communication suffices; adds a pipeline-based early-exit inference that continues the KV computation of exited tokens in parallel on later stages, plus KV recomputation.
Key results
- Training time overhead from added early exits is small under pipeline parallelism; e.g., 1.3B model time per iteration rises from 5.23s to 5.31s with early exits (vs 5.24s with both optimizations) and peak memory is unchanged with all optimizations applied.
- Early-exit inference with confidence thresholds achieves 2x or higher speedup over full-model inference in many cases with comparable or better HELM scores.
- Pipeline-based inference outperforms KV recomputation with TP=1 for thresholds below 1, though KV recomputation with TP=4 is faster.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add early exits at 1/4 and 1/2 depth with optimizations 1 and 2 | GPT 1.3B training · Training iteration | time per iteration (s): 5.23 -> 5.24 | Hurt |
| Add early exits at 1/4 and 1/2 depth with optimizations 1 and 2 | GPT 1.3B training · Training iteration | peak memory (GB): 19.85 -> 19.85 | No clear effect |
| Add early exits at 1/4 and 1/2 depth with optimizations 1 and 2 | GPT 7B training · Training iteration | time per iteration (s): 17.75 -> 17.79 | Hurt |
| Add early exits at 1/4 and 1/2 depth with optimizations 1 and 2 | GPT 7B training · Training iteration | peak memory (GB): 62.27 -> 62.27 | No clear effect |
| Defer early-exit forward computation to backward steps (ablation) | Early-exit GPT training · Training iteration, 1.3B GPT | peak memory (GB): 24.05 -> 22.56 | Helped |
| Defer early-exit forward computation to backward steps (ablation) | Early-exit GPT training · Training iteration, 7B GPT | peak memory (GB): 67.42 -> 65.79 | Helped |
| Move each exit from the end of a pipeline stage to the beginning of the next stage (ablation) | Early-exit GPT training · Training iteration, 1.3B GPT | time per iteration (s): 5.31 -> 5.28 | Helped |
| Move each exit from the end of a pipeline stage to the beginning of the next stage (ablation) | Early-exit GPT training · Training iteration, 7B GPT | peak memory (GB): 67.42 -> 62.27 | Helped |
| Apply both training optimizations to the early exits (ablation) | Early-exit GPT training · Training iteration, 1.3B GPT | peak memory (GB): 24.05 -> 19.85 | Helped |
| Apply both training optimizations to the early exits (ablation) | Early-exit GPT training · Training iteration, 7B GPT | peak memory (GB): 67.42 -> 62.27 | Helped |
| Use pipeline-based early-exit inference instead of full-model inference | EE-LLM 1.3B and 7B · HELM tasks: BoolQ, TruthfulQA, NaturalQuestions, XSUM, CNN/DailyMail | relative speedup | Helped |
Surprising findings
- Adding early exits to the last pipeline stage's bottleneck can be avoided: early exits in middle stages cost nearly nothing in training time and peak memory because they fill implicit pipeline bubbles.
- Final-exit loss of early-exit models is close to or slightly below that of the standard model, suggesting early-exit losses do not hurt the full-model output.
Limitations
- Pipeline-based inference requires multiple devices for pipeline parallelism
- KV recomputation has high theoretical complexity and benefits depend on GPU batching effects
- Early-exit models trained in experiments are small-ish and not fine-tuned or aligned; 7B models show more loss spikes when embeddings are untied and no layer norm is used
Experimental setup
Datasets: Data-Juicer pre-training data subset BoolQ TruthfulQA NaturalQuestions XSUM CNN/DailyMail
Benchmarks: HELM BoolQ HELM TruthfulQA HELM NaturalQuestions open-book HELM NaturalQuestions closed-book HELM XSUM HELM CNN/DailyMail
Metrics: exact match F1 ROUGE-L training time per iteration peak GPU memory inference speedup
Baselines: standard GPT LLM of same architecture KV recomputation
Builds on: Megatron-LM PipeDream-Flush (1F1B) HELM
Closest papers by meaning
- Puzzle: Distillation-Based NAS for Inference-Optimized LLMs (ICML 2025 · similarity 0.89)
- xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference (ICML 2025 · similarity 0.89)
- Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (ICML 2024 · similarity 0.89)
- ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.89)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.89)
- Scaling Laws Meet Model Architecture: Toward Inference-Efficient LLMs (ICLR 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.