EE-LLM: Масштабное обучение и вывод раннего выхода в больших языковых моделях с использованием 3D-параллелизма

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

December 8, 2023

Авторы: Yanxi Chen, Xuchen Pan, Yaliang Li, Bolin Ding, Jingren Zhou

cs.AI

Аннотация

Мы представляем EE-LLM — фреймворк для масштабного обучения и вывода ранневыходных больших языковых моделей (LLM). Хотя недавние работы показали предварительные доказательства эффективности раннего выхода для ускорения вывода LLM, EE-LLM делает фундаментальный шаг к масштабированию таких моделей, поддерживая их обучение и вывод с использованием массивного 3D-параллелизма. Построенный на основе Megatron-LM, EE-LLM реализует ряд алгоритмических инноваций и оптимизаций производительности, адаптированных для раннего выхода, включая легковесный метод, облегчающий обратное распространение для цели обучения с ранним выходом при использовании конвейерного параллелизма, техники задействования простаивающих ресурсов в исходном расписании конвейера для вычислений, связанных с ранневыходными слоями, и два подхода к выводу с ранним выходом, совместимых с кэшированием ключей и значений (KV) для авторегрессионной генерации. Наши аналитические и эмпирические исследования показывают, что EE-LLM достигает высокой эффективности обучения с незначительными вычислительными накладными расходами по сравнению со стандартным обучением LLM, а также выдающегося ускорения вывода без ущерба для качества выходных данных. Для содействия дальнейшим исследованиям и внедрению мы публикуем EE-LLM по адресу 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.