[Paper Review] EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism

Yanxi Chen, Xuchen Pan+3 more|arXiv (Cornell University)|Dec 8, 2023

Topic Modeling4 citations

TL;DR

EE-LLM introduces a scalable framework for training and inferring early-exit large language models using 3D parallelism (tensor, pipeline, and data parallelism). It enables efficient backpropagation through pipeline-parallel stages, leverages idle compute for early-exit computation, and supports KV-caching-compatible inference, achieving near-zero training overhead and significant inference speedups without quality loss on large-scale models like 13B and 7B Llama 2.

ABSTRACT

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.

Motivation & Objective

• To enable large-scale training and inference of early-exit LLMs at scale, addressing the lack of support in existing frameworks.
• To overcome the challenge of training early-exit LLMs with pipeline parallelism, where multiple loss terms are distributed across stages.
• To minimize computational overhead during training by optimizing gradient accumulation and resource utilization.
• To support efficient, autoregressive generation with KV caching in early-exit inference.
• To release a production-ready, open-source framework for community adoption and further research.

Proposed method

• Introduces a lightweight gradient accumulation scheme that enables unbiased backpropagation through pipeline-parallel early-exit models by adjusting loss scaling and microbatch handling.
• Leverages idle compute resources in the pipeline schedule to perform early-exit layer computations during forward passes, improving hardware utilization.
• Proposes two inference strategies compatible with KV caching: one using a shared KV cache across all exits and another using per-exit cache management.
• Employs 3D parallelism (tensor, pipeline, and data) to scale training and inference of large early-exit models beyond the capacity of single devices.
• Optimizes training efficiency by minimizing additional computational cost for early-exit heads, especially the large output embedding matrices.
• Uses a variance-reduction technique in gradient estimation that combines gradients from early and final exits across microbatches, reducing gradient noise.

Experimental results

Research questions

• RQ1 Can early-exit LLMs be trained efficiently at scale using pipeline parallelism, which traditionally conflicts with multi-exit training objectives?
• RQ2 What is the computational overhead of training early-exit LLMs compared to standard LLM training, and can it be minimized?
• RQ3 How can early-exit inference be made compatible with KV caching for autoregressive generation?
• RQ4 What is the impact of early exiting on inference speedup and output quality in large-scale LLMs?
• RQ5 Can the proposed framework achieve high training efficiency and inference speedup without compromising model accuracy?

Key findings

• EE-LLM achieves near-zero training overhead compared to standard LLM training, with negligible additional computational cost for early-exit heads.
• The framework supports training of early-exit LLMs at scale, including 13B and 7B Llama 2 models, using full 3D parallelism.
• Inference with EE-LLM achieves significant speedups—up to 2.5x faster on average—while maintaining output quality comparable to full-model generation.
• The proposed gradient accumulation method reduces gradient variance by up to 30% in experiments, improving training stability.
• The framework enables efficient KV-caching-compatible inference, making early exiting viable for autoregressive generation tasks.
• Empirical results show that early-exit LLMs trained with EE-LLM maintain strong performance on downstream tasks, even when exiting early on simple inputs.