EE-LLM: Addestramento su larga scala e inferenza di modelli linguistici di grandi dimensioni con uscita anticipata mediante parallelismo 3D

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

December 8, 2023

Autori: Yanxi Chen, Xuchen Pan, Yaliang Li, Bolin Ding, Jingren Zhou

cs.AI

Abstract

Presentiamo EE-LLM, un framework per l'addestramento e l'inferenza su larga scala di modelli linguistici di grandi dimensioni (LLM) con uscita anticipata (early-exit). Mentre lavori recenti hanno mostrato prove preliminari dell'efficacia dell'uscita anticipata nell'accelerare l'inferenza degli LLM, EE-LLM compie un passo fondamentale verso la scalabilità degli LLM con uscita anticipata, supportandone l'addestramento e l'inferenza con un massiccio parallelismo 3D. Basato su Megatron-LM, EE-LLM implementa una varietà di innovazioni algoritmiche e ottimizzazioni delle prestazioni specifiche per l'uscita anticipata, tra cui un metodo leggero che facilita la retropropagazione per l'obiettivo di addestramento con uscita anticipata in presenza di parallelismo pipeline, tecniche per sfruttare le risorse inattive nel programma pipeline originale per il calcolo relativo ai livelli di uscita anticipata, e due approcci per l'inferenza con uscita anticipata compatibili con la memorizzazione della cache KV per la generazione autoregressiva. Il nostro studio analitico ed empirico dimostra che EE-LLM raggiunge un'elevata efficienza nell'addestramento con un overhead computazionale trascurabile rispetto all'addestramento standard degli LLM, oltre a un eccezionale aumento della velocità di inferenza senza compromettere la qualità dell'output. Per favorire ulteriori ricerche e l'adozione, rilasciamo EE-LLM all'indirizzo https://github.com/pan-x-c/EE-LLM.

English

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.