EE-LLM: Entrenamiento e Inferencia a Gran Escala de Modelos de Lenguaje con Salidas Tempranas utilizando Paralelismo 3D

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

December 8, 2023

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

cs.AI

Resumen

Presentamos EE-LLM, un marco para el entrenamiento e inferencia a gran escala de modelos de lenguaje grandes (LLMs) con salida temprana. Si bien trabajos recientes han mostrado evidencia preliminar de la eficacia de la salida temprana para acelerar la inferencia de LLMs, EE-LLM da un paso fundamental hacia la escalabilidad de LLMs con salida temprana al soportar su entrenamiento e inferencia con paralelismo 3D masivo. Basado en Megatron-LM, EE-LLM implementa una variedad de innovaciones algorítmicas y optimizaciones de rendimiento específicas para la salida temprana, incluyendo un método ligero que facilita la retropropagación para el objetivo de entrenamiento con salida temprana en paralelismo de tuberías, técnicas que aprovechan los recursos inactivos en el plan original de la tubería para cálculos relacionados con las capas de salida temprana, y dos enfoques de inferencia con salida temprana compatibles con el almacenamiento en caché de KV para la generación autoregresiva. Nuestro estudio analítico y empírico muestra que EE-LLM logra una gran eficiencia en el entrenamiento con un sobrecosto computacional insignificante en comparación con el entrenamiento estándar de LLMs, así como una aceleración sobresaliente en la inferencia sin comprometer la calidad de la salida. Para facilitar futuras investigaciones y su adopción, publicamos EE-LLM en 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.