EE-LLM : Entraînement à grande échelle et inférence de modèles de langage massifs à sortie précoce avec parallélisme 3D

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

Résumé

Nous présentons EE-LLM, un cadre pour l'entraînement à grande échelle et l'inférence de modèles de langage massifs (LLMs) à sortie précoce. Alors que des travaux récents ont montré des preuves préliminaires de l'efficacité de la sortie précoce pour accélérer l'inférence des LLMs, EE-LLM représente une étape fondamentale vers la mise à l'échelle des LLMs à sortie précoce en supportant leur entraînement et leur inférence avec un parallélisme 3D massif. Basé sur Megatron-LM, EE-LLM met en œuvre diverses innovations algorithmiques et optimisations de performance spécifiquement adaptées à la sortie précoce, incluant une méthode légère qui facilite la rétropropagation pour l'objectif d'entraînement à sortie précoce avec parallélisme en pipeline, des techniques pour exploiter les ressources inactives dans le plan de pipeline original pour les calculs liés aux couches de sortie précoce, et deux approches d'inférence à sortie précoce compatibles avec la mise en cache KV pour la génération autoregressive. Notre étude analytique et empirique montre qu'EE-LLM atteint une grande efficacité d'entraînement avec un surcoût computationnel négligeable par rapport à l'entraînement standard des LLMs, ainsi qu'une accélération remarquable de l'inférence sans compromettre la qualité des sorties. Pour faciliter la recherche et l'adoption ultérieures, nous publions EE-LLM à l'adresse 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.