EE-LLM: Grootschalige training en inferentie van Early-Exit Large Language Modellen met 3D-parallelisme

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

December 8, 2023

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

cs.AI

Samenvatting

We presenteren EE-LLM, een raamwerk voor grootschalige training en inferentie van early-exit large language models (LLMs). Hoewel recente werken voorlopig bewijs hebben getoond voor de effectiviteit van early exiting bij het versnellen van LLM-inferentie, zet EE-LLM een fundamentele stap richthet opschalen van early-exit LLMs door hun training en inferentie te ondersteunen met massieve 3D-parallelisme. Gebouwd op Megatron-LM, implementeert EE-LLM een verscheidenheid aan algoritmische innovaties en prestatieoptimalisaties die zijn toegespitst op early exiting, waaronder een lichtgewicht methode die backpropagatie vergemakkelijkt voor het early-exit trainingsdoel met pipeline-parallelisme, technieken om inactieve resources in het oorspronkelijke pipelineschema te benutten voor berekeningen gerelateerd aan early-exit lagen, en twee benaderingen van early-exit inferentie die compatibel zijn met KV-caching voor autoregressieve generatie. Onze analytische en empirische studie toont aan dat EE-LLM een uitstekende trainingsefficiëntie bereikt met verwaarloosbare rekenkundige overhead in vergelijking met standaard LLM-training, evenals een opmerkelijke versnelling van de inferentie zonder in te leveren op de uitvoerkwaliteit. Om verder onderzoek en adoptie te vergemakkelijken, hebben we EE-LLM vrijgegeven op 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.