Back

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

Yanxi ChenXuchen PanYaliang LiBolin DingJingren Zhou

article2024ICML96 citations

Presents EE-LLM, a scalable framework built on Megatron-LM that enables efficient 3D parallel training of early-exit language models across pipeline stages and resolves the incompatibility between early exiting and KV caching during autoregressive inference.

Listen

Deploying large language models incurs significant computational costs and high latency during inference, as these massive systems traditionally pass every input through their entire network regardless of task difficulty. Early exiting offers an intuitive solution by allowing simpler tokens or inputs to finish processing at intermediate layers. However, early-exit models have remained confined to relatively small scales because standard distributed training systems—which rely on multidimensional parallel computing—lack native support for calculating multiple loss functions across separate hardware stages, while autoregressive text generation conflicts with standard key-value memory caching.

The article introduces and evaluates EE-LLM, an open-source framework designed to enable the large-scale training and inference of early-exit language models using comprehensive three-dimensional parallel computing. The main objective is to demonstrate that massive early-exit models can be trained with negligible computational overhead compared to standard models, while achieving substantial inference speedups without degrading output quality.

The evaluation relies on empirical pre-training experiments and standard downstream benchmarks. The authors implemented algorithmic innovations atop the Megatron-LM framework, including a lightweight mathematical reformulation for backpropagation across pipeline stages and memory optimizations that shift intermediate computations into existing hardware idle periods. Training efficiency was tested on models ranging from 1.3 billion to 30 billion parameters across an 8-node cluster of 64 Nvidia A100 GPUs, using pre-training datasets from Data-Juicer. Downstream performance and inference latency were measured across six question-answering and summarization benchmarks, including BoolQ, TruthfulQA, NaturalQuestions, XSUM, and CNN/DailyMail.

The investigation produced four primary findings. First, training early-exit models with EE-LLM introduces virtually zero peak GPU memory overhead and negligible runtime penalties compared to standard models by repurposing idle pipeline bubbles and deferring early-exit forward calculations to backward steps. Second, early-exit models converge at the same pace as standard architectures, with final-layer loss curves matching or slightly exceeding baseline full-model performance. Third, during inference, the proposed pipeline-parallel execution method achieves a 2x or higher speedup over full-model inference across multiple tasks while maintaining comparable benchmark evaluation scores. Fourth, the framework successfully resolves the conflict between early exiting and key-value caching during autoregressive generation via both key-value recomputation and a newly introduced pipeline-parallel inference mechanism.

These results demonstrate that organizations can deploy early-exit mechanisms to cut operational inference costs and reduce latency without needing additional hardware budgets for training. Because training overhead is minimal, engineering teams do not face the typical cost penalties associated with adding multi-exit capacity. Furthermore, the findings show that early exiting does not compromise output quality, proving that large language models naturally possess intermediate features capable of handling simpler text generation steps efficiently.

Organizations aiming to reduce inference latency should consider adopting EE-LLM for large-scale training pipelines. When configuring models, practitioners should place early exits primarily in middle pipeline stages rather than initial or terminal stages to balance hardware loads, while avoiding an excessive number of exits per network to minimize vocabulary projection costs. Before full deployment, teams should conduct task-specific threshold calibration to optimize the trade-off between execution speed and task accuracy for their specific domain.

Readers should note certain limitations: the empirical training evaluations were capped at 30 billion parameters due to hardware constraints, though theoretical scaling extends further. Additionally, the highest-efficiency inference method requires multi-device pipeline parallelism, which may not suit single-GPU edge deployments. Nonetheless, given the rigorous mathematical proofs and consistent multi-benchmark results, decision-makers can have high confidence in the framework's training efficiency and inference acceleration capabilities.

Alibaba Groupbackpropagation algorithmcachinginferenceearly-exit inferenceearly-exit LLM training3D parallelism
pan-x-c/EE-LLMChen et al (2024).pdf
  • Paper: Confident Adaptive Language Modeling, Tal Schuster et al. (2022). Confident Adaptive Language Modeling establishes calibrated token-level early exiting for LLM generation, providing the early-exit inference setting that EE-LLM scales to large, parallel training and serving.
Cover for EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism

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.

Topics

backpropagation algorithmcachinginferencelanguage modellarge language modelsoptimizationparallelismprobabilistic inferencedata parallelprompt engineering

Keywords

early-exit inferenceearly-exit LLM training3D parallelismdeep supervisionearly-exit loss weightsunbiased gradient estimate

Table of Contents

  • Abstract
  • 1. Introduction
  • 2. An overview of EE-LLM
  • 3. Training
  • 3.1. Backpropagation through pipeline stages
  • 3.1.1. METHODOLOGY
  • 3.1.2. RATIONALE
  • 3.1.3. INTEGRATION WITH THE 1F1B SCHEDULE
  • 3.2. Training efficiency
  • 3.3. Advanced features
  • 4. Inference
  • 5. Experiments
  • 5.1. Training
  • 5.2. Inference
  • 6. Conclusions
  • Acknowledgements
  • Impact Statement
  • References
  • Structure of the appendix
  • A. Training efficiency
  • A.1. Utilization of idle resources
  • A.2. Further performance optimizations
  • A.3. Theoretical analysis for general cases
  • A.3.1. TRAINING TIME PER ITERATION
  • A.3.2. PEAK GPU MEMORY
  • B. Additional experiments
  • B.1. A comparison between two inference methods
  • B.2. Examples of generated texts
  • B.3. Training with different early-exit configurations
  • C. Advanced features
  • C.1. Non-constant early-exit loss weights
  • C.2. Filling explicit bubbles
  • C.2.1. METHODOLOGY
  • C.2.2. ANALYSIS
  • D. Implementations
  • D.1. Model architectures
  • D.2. Pipeline scheduling
  • D.3. Inference service
  • E. Preliminaries
  • F. Related works

Knowls

  1. Knowl 1 — Auxiliary losses enable exact early-exit backpropagation across pipeline stages

    model/method

    For a model split across KK pipeline stages, let LiL_i be the weighted sum of early-exit and/or final-exit losses computed on Stage ii, so the full training objective is L=∑i=1KLiL=\sum_{i=1}^{K}L_i. Let xix_i be the hidden-state tensor passed from Stage ii to Stage i+1i+1. Define local auxiliary losses in reverse stage order by LKaux=LKL_K^{\mathrm{aux}}=L_K and, for i=K−1,…,1i=K-1,\ldots,1, gi=∂Li+1aux/∂xig_i=\partial L_{i+1}^{\mathrm{aux}}/\partial x_i and Liaux=Li+⟨gi,xi⟩L_i^{\mathrm{aux}}=L_i+\langle g_i,x_i\rangle. Stage ii treats the received tensor gig_i as constant while differentiating its auxiliary loss. This lets each stage compute the gradient contributions from its own and all later stages’ losses using the usual adjacent-stage pipeline communication; it requires no additional inter-stage communication. The resulting parameter and activation gradients equal those of the full objective. Tied parameters across stages are supported by first computing gradients as if the parameters were untied, then summing and synchronizing gradients for the shared parameters.

  2. Knowl 2 — Configurable early-exit Transformer architecture

    model/method

    EE-LLM extends a GPT-style Transformer with early-exit layers that map intermediate hidden states to outputs, such as next-token logits. Users can choose the locations and number of exits, and can select a minimal exit consisting of an output embedding matrix with an optional normalization layer, or add modules such as an MLP or a full Transformer layer. Input and output embedding matrices across early and final exits can be tied or untied. This allows exit expressivity and parameter sharing to be configured for a target training or inference use case.

  3. Knowl 3 — Pipeline scheduling bounds the training overhead of early exits

    theoretical result

    For 1F1B training with PP pipeline stages and MM microbatches, suppose every stage has the same backbone structure, all early-exit layers have the same structure, and pipeline communication latency is ignored. Let NiN_i be the number of exits on Stage ii. For any module type o∈{IN,EE,FE,BB}o\in\{\mathrm{IN},\mathrm{EE},\mathrm{FE},\mathrm{BB}\}—input processing, early exit, final exit, and per-stage Transformer backbone—let fof_o and bob_o be its forward and backward time per microbatch. Let 1(E)\mathbf{1}(E) equal 1 when condition EE holds and 0 otherwise. An upper bound on iteration time is

    fIN+bIN+(P−1)(fBB+bBB)+∑i=1P−1Ni(fEE+bEE)+Mmax⁡i∈[P]{fBB+bBB+1(i=1)(fIN+bIN)+1(i=P)(fFE+bFE)+Ni(fEE+bEE)}. f_{\mathrm{IN}}+b_{\mathrm{IN}}+(P-1)(f_{\mathrm{BB}}+b_{\mathrm{BB}})+\sum_{i=1}^{P-1}N_i(f_{\mathrm{EE}}+b_{\mathrm{EE}})+M\max_{i\in[P]}\left\{f_{\mathrm{BB}}+b_{\mathrm{BB}}+\mathbf{1}(i=1)(f_{\mathrm{IN}}+b_{\mathrm{IN}})+\mathbf{1}(i=P)(f_{\mathrm{FE}}+b_{\mathrm{FE}})+N_i(f_{\mathrm{EE}}+b_{\mathrm{EE}})\right\}.

    Here [P]={1,…,P}[P]=\{1,\ldots,P\}. This bound captures how exits on earlier stages add work to the critical path and how exit work can be absorbed when it does not change the maximum stage time in the steady 1F1B phase. In particular, adding one minimalistic exit to each of kk middle stages adds only k(fEE+bEE)k(f_{\mathrm{EE}}+b_{\mathrm{EE}}) to iteration time under the paper’s implicit-bubble scheduling analysis.

    The method also defers each early-exit forward computation from the forward pass to its corresponding backward step. If sequence length is ss, microbatch size is bb, vocabulary size is VV, and an exit is on Stage ii, this reduces the peak activation storage for its logits from sbV(P−i+1)sbV(P-i+1) to sbVsbV. With no exit on the first stage, peak GPU memory remains unchanged when sbVsbV is below the activation memory of one stage’s backbone for one microbatch. The general analysis assumes untied input and output embedding matrices; it ignores point-to-point communication and models gradient and optimizer-state memory as a constant multiple of parameter memory.

  4. Knowl 4 — Filling 1F1B pipeline bubbles with partial microbatch computation

    algorithm

    EE-LLM uses otherwise-idle explicit bubbles in a 1F1B schedule to perform useful partial forward and backward computation. Let PP be the number of pipeline stages, and let ff and bb be estimated forward and backward times per microbatch. Each bubble region can accommodate at most ⌊(P−1)/(f/b+1)⌋\left\lfloor(P-1)/(f/b+1)\right\rfloor inserted microbatches without increasing iteration time. In the bubble between warm-up and steady state, inserted microbatch ii runs forward through the first K+1−iK+1-i stages, where KK is the number inserted, and backpropagates the early-exit losses encountered on those stages. It need only forward as far as the last exit-bearing stage in that range. In cool-down bubbles, an inserted microbatch runs the full forward pass, then backpropagates only through the final and early-exit losses on the last few stages; the number of backward stages for inserted microbatch ii can be set to ⌊P−i(f/b+1)⌋\left\lfloor P-i(f/b+1)\right\rfloor. EE-LLM’s implementation keeps the microbatch count per iteration unchanged by replacing some full passes with partial passes and truncating backward computation for the final microbatches. The method is intended for settings without shared parameters across pipeline stages; partial work in early stages can also increase activation memory, and some data receives only partial computation.

  5. Knowl 5 — Partial bubble computation can preserve unbiased gradients while reducing variance

    theoretical result

    Assume microbatches are sampled independently from a data distribution and model parameters are not tied across pipeline stages. The partial gradients produced by bubble-filling computation can contribute useful updates to the original early-exit training objective without increasing iteration time. For partial backward work confined to later stages, entrywise rescaling of the affected parameter gradients yields an unbiased estimate of the corresponding population-risk gradient, with reduced variance. For partial forward and backward work confined to earlier stages, rescaling the gradients associated with the losses computed there likewise restores an unbiased estimate; the paper’s variance claim holds unless gradients from early and later losses have sufficiently strong negative correlation. These are conditional results: partial computations do not generally provide a full-objective gradient for every parameter on every added microbatch.

  6. Knowl 6 — Pipeline inference overlaps next-token generation with KV-cache completion

    model/method

    In EE-LLM’s pipeline-based autoregressive inference, when a token meets the exit condition at a pipeline stage, that token is returned to the first stage and generation of the next token begins immediately. In parallel, the current token’s forward computation continues from its exit layer through later stages, filling the missing key-value (KV) caches needed by future tokens. Thus, computation after the chosen exit overlaps with computation for subsequent tokens, resolving the incompatibility between early exiting and ordinary KV caching. The latency to produce a token matches the time to compute up to its selected exit, except when the exit lies within the first stage: the next token must then wait for that stage’s current-token forward pass to finish. This method requires multiple devices for pipeline parallelism and does not depend on batching effects for its theoretical acceleration.

  7. Knowl 7 — KV recomputation repairs caches for recently early-exited tokens

    algorithm

    For early-exit inference without pipeline parallelism, EE-LLM maintains a list of recent tokens whose KV caches are missing in deeper layers because those tokens exited early. On each subsequent forward pass, it batches those tokens with the current token and recomputes their missing caches. This permits later tokens to use ordinary cached autoregressive attention while relying on the batching effect of GPU computation for speed. To prevent the list from growing indefinitely, the method forces a full-model forward pass when the number of early-exit tokens reaches a user-specified limit. Its acceleration may not transfer to hardware that does not benefit from batched computation.

  8. Knowl 8 — Training-time and memory measurements validate exit-placement optimizations

    empirical result

    A controlled 1F1B experiment compared standard GPT training with early-exit training and two optimizations: Optimization 1 defers early-exit forward computation to backward steps; Optimization 2 moves an exit from the end of one stage to the beginning of the next. The setup used pipeline parallelism degree 4, tensor and data parallelism degree 1, microbatch size 2, global batch size 128, and sequence length 2048. The base early-exit setup placed exits at one-quarter depth at the end of Stage 1 and one-half depth at the end of Stage 2. Time per iteration and peak memory were, respectively: for 1.3B parameters, standard 5.23 s and 19.85 GB; early exit without these optimizations 5.31 s and 24.05 GB; Optimization 1 only 5.29 s and 22.56 GB; Optimization 2 only 5.28 s and 20.23 GB; both optimizations 5.24 s and 19.85 GB. For 7B parameters, the corresponding results were standard 17.75 s and 62.27 GB; early exit 17.93 s and 67.42 GB; Optimization 1 only 17.91 s and 65.79 GB; Optimization 2 only 17.81 s and 62.27 GB; both optimizations 17.79 s and 62.27 GB. Stage 2 was the peak-memory bottleneck for the Optimization 2-only results; Stage 1 was the bottleneck otherwise. In separate scaling experiments on a 64-GPU cluster, EE-LLM trained models up to 30B parameters. With pipeline parallelism, peak memory stayed unchanged as exits were added to middle stages and increased when an exit was added to the first stage; without pipeline parallelism, memory increased with the number of exits.

  9. Knowl 9 — Early-exit training retains comparable final-exit loss convergence

    empirical result

    EE-LLM trained a 1.3B GPT with 24 Transformer layers and minimal exits at layers 6 and 12, using exit-loss weights 1/41/4 and 1/21/2, final-loss weight 1, and tied embedding matrices. It also trained a 7B GPT with 32 layers and exits at layers 8 and 16, using weights 0.1 and 0.2 and untied embedding matrices. In both experiments, batch size and sequence length were 2048, and standard models with the same architecture, data, and training hyperparameters served as controls. The early- and final-exit negative log-likelihood losses declined at similar rates; early-exit losses were somewhat higher than final-exit losses. The final-exit loss of each early-exit model remained close to, and sometimes slightly below, the corresponding standard model’s loss in these experiments. This demonstrates compatibility of the weighted multi-exit training objective with comparable final-exit convergence in the tested configurations, not a universal guarantee.

  10. Knowl 10 — Pipeline early-exit decoding yields speedups with largely preserved benchmark scores

    empirical result

    The inference evaluation used 1.3B and 7B early-exit models pretrained from scratch on 300B and 150B tokens, respectively, with four pipeline stages on a server with four Nvidia A100-40GB GPUs. Models used greedy decoding and exited when the maximum next-token probability at an exit exceeded a threshold; lowering the threshold increased speedup, while threshold 1 disabled early exits and supplied the full-model latency baseline. Evaluation covered BoolQ and TruthfulQA (exact match), NaturalQuestions open-book and closed-book (F1), and XSUM and CNN/DailyMail (ROUGE-L). Across many task and threshold settings, early exiting achieved at least 2× speedup over full-model inference with comparable or sometimes higher evaluation scores; the results do not imply that all thresholds preserve quality. Example generations also showed identical output at a moderate threshold with lower latency, while more aggressive thresholds could change generated text.

Coverage note — No substantial contributed method or result was deliberately omitted. Low-level implementation class details and individual generated-text examples were excluded because they add little independent knowledge beyond the architecture and empirical findings captured here.

References

  1. 1.Ba, J., Kiros, J. R., and Hinton, G. E. Layer normalization. ArXiv, abs/1607.06450, 2016.
  2. 2.Bae, S., Ko, J., Song, H., and Yun, S.-Y. Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding. In EMNLP, 2023.
  3. 3.Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. In ICLR, 2015.
  4. 4.Bengio, E., Bacon, P.-L., Pineau, J., and Precup, D. Conditional computation in neural networks for faster models. ArXiv, abs/1511.06297, 2015.
  5. 5.Bengio, Y., Leonard, N., and Courville, A. C. Estimating or propagating gradients through stochastic neurons for conditional computation. ArXiv, abs/1308.3432, 2013.
  6. 6.Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosselut, A., Brunskill, E., Brynjolfsson, E., Buch, S., Card, D., Castellon, R., Chatterji, N. S., Chen, A. S., Creel, K. A., Davis, J., Demszky, D., Donahue, C., Doumbouya, M., Durmus, E., Ermon, S., Etchemendy, E., Ethayarajh, K., Fei-Fei, L., Finn, C., Gale, T., Gillespie, L., Goel, K., Goodman, N. D., Grossman, S., Guha, N., Hashimoto, T., Henderson, P., Hewitt, J., Ho, D. E., Hong, J., Hsu, K., Huang, J., Icard, T. F., Jain, S., Jurafsky, D., Kalluri, P., Karamcheti, V., Keeling, G., Khani, F., Khattab, O., Koh, P. W., Krass, M. S., Krishna, R., Kuditipudi, R., Kumar, A., Ladhak, F., Lee, M., Lee, T., Leskovec, J., Levent, I., Li, X. L., Li, X., Ma, T., Malik, A., Manning, C. D., Mirchandani, S., Mitchell, E., Munyikwa, Z., Nair, S., Narayan, A., Narayanan, D., Newman, B., Nie, A., Niebles, J. C., Nilforoshan, H., Nyarko, J. F., Ogut, G., Orr, L. J., Papadimitriou, I., Park, J. S., Piech, C., Portelance, E., Potts, C., Raghunathan, A., Reich, R., Ren, H., Rong, F., Roohani, Y. H., Ruiz, C., Ryan, J., R’e, C., Sadigh, D., Sagawa, S., Santhanam, K., Shih, A., Srinivasan, K. P., Tamkin, A., Taori, R., Thomas, A. W., Tramer, F., Wang, R. E., Wang, W., Wu, B., Wu, J., Wu, Y., Xie, S. M., Yasunaga, M., You, J., Zaharia, M. A., Zhang, M., Zhang, T., Zhang, X., Zhang, Y., Zheng, L., Zhou, K., and Liang, P. On the opportunities and risks of foundation models. ArXiv, abs/2108.07258, 2021.
  7. 7.Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In NeurIPS, 2020.
  8. 8.Chen, D., Huang, Y., Ma, Z., Chen, H., Pan, X., Ge, C., Gao, D., Xie, Y., Liu, Z., Gao, J., Li, Y., Ding, B., and Zhou, J. Data-juicer: A one-stop data processing system for large language models. ArXiv, abs/2309.02033, 2023.
  9. 9.Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. ArXiv, abs/1604.06174, 2016.
  10. 10.Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A. M., Pillai, T. S., Pellat, M., Lewkowycz, A., Moreira, E., Child, R., Polozov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier-Hellstern, K., Eck, D., Dean, J., Petrov, S., and Fiedel, N. Palm: Scaling language modeling with pathways. J. Mach. Learn. Res., 24:240:1–240:113, 2023.
  11. 11.Clark, C., Lee, K., Chang, M., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. In NAACL, 2019.
  12. 12.Corro, L. D., Giorno, A. D., Agarwal, S., Yu, T., Awadallah, A. H., and Mukherjee, S. Skipdecode: Autoregressive skip decoding with batching and caching for efficient llm inference. ArXiv, abs/2307.02628, 2023.
  13. 13.Dai, Y., Pan, R., Iyer, A., Li, K., and Netravali, R. Apparate: Rethinking early exits to tame latency-throughput tensions in ml serving. ArXiv, abs/2312.05385, 2023.
  14. 14.Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In NeurIPS, 2022.
  15. 15.Data-Juicer. Refined open source dataset by data-juicer. https://github.com/alibaba/data-juicer/blob/main/configs/data_juicer_recipes/README.md, 2023.
  16. 16.Devlin, J., Chang, M., Lee, K., and Toutanova, K. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, pp. 4171–4186, 2019.
  17. 17.Din, A. Y., Karidi, T., Choshen, L., and Geva, M. Jump to conclusions: Short-cutting transformers with linear transformations. ArXiv, abs/2303.09435, 2023.
  18. 18.Duggal, R., Freitas, S., Dhamnani, S., Chau, D. H., and Sun, J. Elf: An early-exiting framework for long-tailed classification. ArXiv, abs/2006.11979, 2020.
  19. 19.Elbayad, M., Gu, J., Grave, E., and Auli, M. Depth-adaptive transformer. In ICLR, 2020.
  20. 20.Fan, S., Rong, Y., Meng, C., Cao, Z., Wang, S., Zheng, Z., Wu, C., Long, G., Yang, J., Xia, L., Diao, L., Liu, X., and Lin, W. DAPPLE: a pipelined data parallel approach for training large models. In PPoPP, pp. 431–445, 2021.
  21. 21.Fedus, W., Zoph, B., and Shazeer, N. M. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1–120:39, 2021.
  22. 22.Gera, A., Friedman, R., Arviv, O., Gunasekara, C., Sznajder, B., Slonim, N., and Shnarch, E. The benefits of bad advice: Autocontrastive decoding across model layers. In ACL, 2023.
  23. 23.Graves, A. Adaptive computation time for recurrent neural networks. ArXiv, abs/1603.08983, 2016.
  24. 24.Han, Y., Huang, G., Song, S., Yang, L., Wang, H., and Wang, Y. Dynamic neural networks: A survey. IEEE Trans. Pattern Anal. Mach. Intell., 44(11):7436–7456, 2022.
  25. 25.Hou, L., Huang, Z., Shang, L., Jiang, X., Chen, X., and Liu, Q. Dynabert: Dynamic BERT with adaptive width and depth. In NeurIPS, 2020.
  26. 26.Hu, B., Zhu, Y., Li, J., and Tang, S. Smartbert: A promotion of dynamic early exiting mechanism for accelerating bert inference. In IJCAI, 2023.
  27. 27.Huang, G., Chen, D., Li, T., Wu, F., van der Maaten, L., and Weinberger, K. Q. Multi-scale dense networks for resource efficient image classification. In ICLR, 2018.
  28. 28.Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M. X., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, Z. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In NeurIPS, pp. 103–112, 2019.
  29. 29.Ilhan, F., Su, G., and Liu, L. Scalefl: Resource-adaptive federated learning with heterogeneous clients. In CVPR, pp. 24532–24541, 2023.
  30. 30.Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. Adaptive mixtures of local experts. Neural Computation, 3:79–87, 1991.
  31. 31.Kaya, Y., Hong, S., and Dumitras, T. Shallow-deep networks: Understanding and mitigating network overthinking. In ICML, volume 97, pp. 3301–3310, 2019.
  32. 32.Kim, S., Hooper, C., Wattanawong, T., Kang, M., Yan, R., Genc¸, H., Dinh, G., Huang, Q., Keutzer, K., Mahoney, M. W., Shao, Y. S., and Gholami, A. Full stack optimization of transformer inference: a survey. ArXiv, abs/2302.14017, 2023a.
  33. 33.Kim, S., Mangalam, K., Moon, S., Malik, J., Mahoney, M. W., Gholami, A., and Keutzer, K. Speculative decoding with big little decoder. In NeurIPS, 2023b.
  34. 34.Kim, Y., Denton, C., Hoang, L., and Rush, A. M. Structured attention networks. In ICLR, 2017.
  35. 35.Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In ICLR, 2014.
  36. 36.Korthikanti, V. A., Casper, J., Lym, S., McAfee, L. C., Andersch, M., Shoeybi, M., and Catanzaro, B. Reducing activation recomputation in large transformer models. ArXiv, abs/2205.05198, 2022.
  37. 37.Kwiatkowski, T., Palomaki, J., Redfield, O., Collins, M., Parikh, A. P., Alberti, C., Epstein, D., Polosukhin, I., Devlin, J., Lee, K., Toutanova, K., Jones, L., Kelcey, M., Chang, M., Dai, A. M., Uszkoreit, J., Le, Q., and Petrov, S. Natural questions: a benchmark for question answering research. Trans. Assoc. Comput. Linguistics, 7:452–466, 2019.
  38. 38.Langedijk, A., Mohebbi, H., Sarti, G., Zuidema, W. H., and Jumelet, J. Decoderlens: Layerwise interpretation of encoder-decoder transformers. ArXiv, abs/2310.03686, 2023.
  39. 39.Laskaridis, S., Kouris, A., and Lane, N. D. Adaptive inference through early-exit networks: Design, challenges and directions. In International Workshop on Embedded and Mobile Deep Learning, 2021.
  40. 40.Lee, C.-Y., Xie, S., Gallagher, P. W., Zhang, Z., and Tu, Z. Deeply-supervised nets. In AISTATS, 2014.
  41. 41.Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In ICML, 2022.
  42. 42.Li, S. and Hoefler, T. Chimera: efficiently training large-scale neural networks with bidirectional pipelines. In SC, pp. 27, 2021.
  43. 43.Li, X., Shao, Y., Sun, T., Yan, H., Qiu, X., and Huang, X. Accelerating bert inference for sequence labeling via early-exit. In ACL, 2021.
  44. 44.Liang, P., Bommasani, R., Lee, T., Tsipras, D., Soylu, D., Yasunaga, M., Zhang, Y., Narayanan, D., Wu, Y., Kumar, A., Newman, B., Yuan, B., Yan, B., Zhang, C., Cosgrove, C., Manning, C. D., R’e, C., Acosta-Navas, D., Hudson, D. A., Zelikman, E., Durmus, E., Ladhak, F., Rong, F., Ren, H., Yao, H., Wang, J., Santhanam, K., Orr, L. J., Zheng, L., Yuksekgonul, M., Suzgun, M., Kim, N. S., Guha, N., Chatterji, N. S., Khattab, O., Henderson, P., Huang, Q., Chi, R., Xie, S. M., Santurkar, S., Ganguli, S., Hashimoto, T., Icard, T. F., Zhang, T., Chaudhary, V., Wang, W., Li, X., Mai, Y., Zhang, Y., and Koreeda, Y. Holistic evaluation of language models. Annals of the New York Academy of Sciences, 1525:140 – 146, 2023.
  45. 45.Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human falsehoods. In ACL, pp. 3214–3252, 2022.
  46. 46.Liu, D., Kan, M., Shan, S., and Chen, X. A simple romance between multi-exit vision transformer and token reduction. In ICLR, 2024.
  47. 47.Liu, W., Zhou, P., Wang, Z., Zhao, Z., Deng, H., and Ju, Q. Fastbert: a self-distilling BERT with adaptive inference time. In ACL, pp. 6035–6044, 2020.
  48. 48.Narayan, S., Cohen, S. B., and Lapata, M. Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In EMNLP, pp. 1797–1807, 2018.
  49. 49.Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. Pipedream: generalized pipeline parallelism for DNN training. In SOSP, pp. 1–15, 2019.
  50. 50.Narayanan, D., Phanishayee, A., Shi, K., Chen, X., and Zaharia, M. Memory-efficient pipeline-parallel DNN training. In ICML, volume 139, pp. 7937–7947, 2021a.
  51. 51.Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M. Efficient large-scale language model training on GPU clusters using megatron-lm. In SC, pp. 58, 2021b.
  52. 52.OpenAI. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
  53. 53.Osawa, K., Li, S., and Hoefler, T. Pipefisher: Efficient training of large language models using pipelining and fisher information matrices. In MLSys, 2023.
  54. 54.Parikh, A., Täckström, O., Das, D., and Uszkoreit, J. A decomposable attention model for natural language inference. In EMNLP, pp. 2249–2255, 2016.
  55. 55.Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. ArXiv, abs/2211.05102, 2022.
  56. 56.Press, O. and Wolf, L. Using the output embedding to improve language models. In EACL, pp. 157–163, 2017.
  57. 57.Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre-training, 2018.
  58. 58.Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners, 2019.
  59. 59.Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: Memory optimizations toward training trillion parameter models. In SC, 2019.
  60. 60.Rasley, J., Rajbhandari, S., Ruwase, O., and He, Y. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In KDD, pp. 3505–3506, 2020.
  61. 61.Santilli, A., Severino, S., Postolache, E., Maiorca, V., Mancusi, M., Marin, R., and Rodola, E. Accelerating transformer inference for translation via parallel decoding. In ACL, 2023.
  62. 62.Scardapane, S., Scarpiniti, M., Baccarelli, E., and Uncini, A. Why should we add early exits to neural networks? Cognitive Computation, 12:954 – 966, 2020.
  63. 63.Schuster, T., Fisch, A., Jaakkola, T. S., and Barzilay, R. Consistent accelerated inference via confident adaptive transformers. In EMNLP, pp. 4962–4979, 2021.
  64. 64.Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D., Tran, V., Tay, Y., and Metzler, D. Confident adaptive language modeling. In NeurIPS, 2022.
  65. 65.Schwartz, R., Stanovsky, G., Swayamdipta, S., Dodge, J., and Smith, N. A. The right tool for the job: Matching model and instance complexities. In ACL, pp. 6640–6651, 2020.
  66. 66.Shazeer, N., Cheng, Y., Parmar, N., Tran, D., Vaswani, A., Koanantakool, P., Hawkins, P., Lee, H., Hong, M., Young, C., Sepassi, R., and Hechtman, B. Mesh-tensorflow: Deep learning for supercomputers. In NIPS, pp. 10435–10444, 2018.
  67. 67.Shazeer, N. M., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. V., Hinton, G. E., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In ICLR, 2017.
  68. 68.Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-lm: Training multi-billion parameter language models using model parallelism. ArXiv, abs/1909.08053, 2019.
  69. 69.Smith, S., Patwary, M., Norick, B., LeGresley, P., Rajbhandari, S., Casper, J., Liu, Z., Prabhumoye, S., Zerveas, G., Korthikanti, V. A., Zhang, E., Child, R., Aminabadi, R. Y., Bernauer, J., Song, X., Shoeybi, M., He, Y., Houston, M., Tiwary, S., and Catanzaro, B. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. ArXiv, abs/2201.11990, 2022.
  70. 70.Szegedy, C., Liu, W., Jia, Y., Sermanet, P., Reed, S. E., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A. Going deeper with convolutions. In CVPR, 2015.
  71. 71.Tang, P., Zhu, P., Li, T., Appalaraju, S., Mahadevan, V., and Manmatha, R. Deed: Dynamic early exit on decoder for accelerating encoder-decoder transformer models. ArXiv, abs/2311.08623, 2023.
  72. 72.Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. ACM Comput. Surv., 55(6), 2022. doi: 10.1145/3530811.
  73. 73.Team, I. Internlm: A multilingual language model with progressively enhanced capabilities. https://github.com/InternLM/InternLM, 2023.
  74. 74.Teerapittayanon, S., McDanel, B., and Kung, H. T. Branchynet: Fast inference via early exiting from deep neural networks. In ICPR, pp. 2464–2469, 2016.
  75. 75.Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Roziere, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models. ArXiv, abs/2302.13971, 2023a.
  76. 76.Touvron, H., Martin, L., Stone, K. R., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D. M., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, E., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A. S., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I. M., Korenev, A. V., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I., Nie, Y., Poulton, A., Reizenstein, J., Rungta, S., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, X., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models. ArXiv, abs/2307.09288, 2023b.
  77. 77.Varshney, N., Chatterjee, A., Parmar, M., and Baral, C. Accelerating llama inference by enabling intermediate layer decoding via instruction tuning with lite. ArXiv, abs/2310.18581, 2023.
  78. 78.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In NeurIPS, pp. 5998–6008, 2017.
  79. 79.Wang, H., Wang, Y., Liu, T., Zhao, T., and Gao, J. Hadskip: Homotopic and adaptive layer skipping of pre-trained language models for efficient inference. In EMNLP, pp. 4283–4294, 2023.
  80. 80.Wang, J., Chen, K., Chen, G., Shou, L., and McAuley, J. Skipbert: Efficient inference with shallow layer skipping. In ACL, 2022.
  81. 81.Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. SmoothQuant: Accurate and efficient post-training quantization for large language models. In ICML, volume 202, pp. 38087–38099, 2023.
  82. 82.Xin, J., Tang, R., Lee, J., Yu, Y., and Lin, J. Deebert: Dynamic early exiting for accelerating BERT inference. In ACL, pp. 2246–2251, 2020.
  83. 83.Xin, J., Tang, R., Yu, Y., and Lin, J. Berxit: Early exiting for BERT with better fine-tuning and extension to regression. In EACL, pp. 91–104, 2021.
  84. 84.Xu, C. and McAuley, J. J. A survey on dynamic neural networks for natural language processing. In EACL, pp. 2325–2336, 2023.
  85. 85.Yang, F., Peng, S., Sun, N., Wang, F., Tan, K., Wu, F., Qiu, J., and Pan, A. Holmes: Towards distributed training across clusters with heterogeneous nic environment. ArXiv, abs/2312.03549, 2023.
  86. 86.Yao, Z., Aminabadi, R. Y., Zhang, M., Wu, X., Li, C., and He, Y. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. In NeurIPS, 2022.
  87. 87.Yuan, B., He, Y., Davis, J., Zhang, T., Dao, T., Chen, B., Liang, P., Re, C., and Zhang, C. Decentralized training of foundation models in heterogeneous environments. In NeurIPS, 2022.
  88. 88.Zeng, D., Du, N., Wang, T., Xu, Y., Lei, T., Chen, Z., and Cui, C. Learning to skip for language modeling. ArXiv, abs/2311.15436, 2023.
  89. 89.Zhang, B. and Sennrich, R. Root mean square layer normalization. In NeurIPS, pp. 12360–12371, 2019.
  90. 90.Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M. T., Li, X., Lin, X. V., Mihaylov, T., Ott, M., Shleifer, S., Shuster, K., Simig, D., Koura, P. S., Sridhar, A., Wang, T., and Zettlemoyer, L. Opt: Open pre-trained transformer language models. ArXiv, abs/2205.01068, 2022.
  91. 91.Zhao, W. X., Zhou, K., Li, J., Tang, T., Wang, X., Hou, Y., Min, Y., Zhang, B., Zhang, J., Dong, Z., Du, Y., Yang, C., Chen, Y., Chen, Z., Jiang, J., Ren, R., Li, Y., Tang, X., Liu, Z., Liu, P., Nie, J., and rong Wen, J. A survey of large language models. ArXiv, abs/2303.18223, 2023a.
  92. 92.Zhao, Y., Xie, Z., Zhuang, C., and Gu, J. Lookahead: An inference acceleration framework for large language model with lossless generation accuracy. ArXiv, abs/2312.12728, 2023b.
  93. 93.Zheng, L., Li, Z., Zhang, H., Zhuang, Y., Chen, Z., Huang, Y., Wang, Y., Xu, Y., Zhuo, D., Xing, E. P., Gonzalez, J. E., and Stoica, I. Alpa: Automating inter- and Intra-Operator parallelism for distributed deep learning. In OSDI, pp. 559–578, 2022.
  94. 94.Zhong, Z., Wang, J., Bao, W., Zhou, J., Zhu, X., and Zhang, X. Semi-hfl: semi-supervised federated learning for heterogeneous devices. Complex & Intelligent Systems, 9:1995–2017, 2022.
  95. 95.Zhou, W., Xu, C., Ge, T., McAuley, J. J., Xu, K., and Wei, F. BERT loses patience: Fast and robust inference with early exit. In NeurIPS, 2020.

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF
License: https://creativecommons.org/licenses/by/4.0/