EE-LLM: Treinamento e Inferência em Grande Escala de Modelos de Linguagem com Saída Antecipada 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

Resumo

Apresentamos o EE-LLM, um framework para treinamento e inferência em larga escala de modelos de linguagem grandes (LLMs) com saída antecipada. Embora trabalhos recentes tenham mostrado evidências preliminares da eficácia da saída antecipada na aceleração da inferência de LLMs, o EE-LLM dá um passo fundamental para escalar LLMs com saída antecipada, suportando seu treinamento e inferência com paralelismo 3D massivo. Construído sobre o Megatron-LM, o EE-LLM implementa uma variedade de inovações algorítmicas e otimizações de desempenho específicas para a saída antecipada, incluindo um método leve que facilita a retropropagação para o objetivo de treinamento de saída antecipada com paralelismo de pipeline, técnicas que aproveitam recursos ociosos no cronograma original de pipeline para computação relacionada às camadas de saída antecipada, e duas abordagens de inferência de saída antecipada que são compatíveis com o armazenamento em cache de KV para geração autoregressiva. Nosso estudo analítico e empírico mostra que o EE-LLM alcança grande eficiência de treinamento com sobrecarga computacional insignificante em comparação com o treinamento padrão de LLMs, além de um excelente aumento de velocidade na inferência sem comprometer a qualidade da saída. Para facilitar pesquisas e adoção futuras, disponibilizamos o EE-LLM em 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.