EE-LLM: Großskaliges Training und Inferenz von Early-Exit Large Language Models mit 3D-Parallelismus

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

December 8, 2023

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

cs.AI

Zusammenfassung

Wir präsentieren EE-LLM, ein Framework für das großskalige Training und die Inferenz von Early-Exit Large Language Models (LLMs). Während aktuelle Arbeiten erste Hinweise auf die Wirksamkeit von Early Exiting zur Beschleunigung der LLM-Inferenz gezeigt haben, macht EE-LLM einen grundlegenden Schritt zur Skalierung von Early-Exit-LLMs, indem es deren Training und Inferenz mit massiver 3D-Parallelität unterstützt. Basierend auf Megatron-LM implementiert EE-LLM eine Vielzahl von algorithmischen Innovationen und Leistungsoptimierungen, die speziell auf Early Exiting zugeschnitten sind, darunter eine leichtgewichtige Methode, die die Backpropagation für das Early-Exit-Trainingsziel mit Pipeline-Parallelität erleichtert, Techniken zur Nutzung von Leerlaufressourcen im ursprünglichen Pipeline-Zeitplan für Berechnungen im Zusammenhang mit Early-Exit-Schichten sowie zwei Ansätze für die Early-Exit-Inferenz, die mit KV-Caching für autoregressive Generierung kompatibel sind. Unsere analytische und empirische Studie zeigt, dass EE-LLM eine hohe Trainingseffizienz mit vernachlässigbarem Rechenaufwand im Vergleich zum Standard-LLM-Training erreicht sowie eine herausragende Beschleunigung der Inferenz ohne Einbußen bei der Ausgabequalität. Um weitere Forschung und die Verbreitung zu fördern, veröffentlichen wir EE-LLM unter 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.