Конвейерный параллелизм с нулевым пузырем

Zero Bubble Pipeline Parallelism

November 30, 2023

Авторы: Penghui Qi, Xinyi Wan, Guangxing Huang, Min Lin

cs.AI

Аннотация

Параллелизм конвейера является одним из ключевых компонентов для крупномасштабного распределенного обучения, однако его эффективность страдает от так называемых "пузырей конвейера", которые ранее считались неизбежными. В данной работе мы представляем стратегию планирования, которая, насколько нам известно, впервые успешно достигает нулевых пузырей конвейера при синхронной семантике обучения. Основная идея этого улучшения заключается в разделении обратного вычисления на две части: одну, которая вычисляет градиент для входных данных, и другую, которая вычисляет градиент для параметров. На основе этой идеи мы вручную разработали новые расписания конвейера, которые значительно превосходят базовые методы. Мы также разработали алгоритм, который автоматически находит оптимальное расписание в зависимости от конкретной конфигурации модели и ограничений памяти. Кроме того, для полного устранения пузырей мы вводим новую технику, позволяющую обходить синхронизации на этапе оптимизатора. Экспериментальные оценки показывают, что наш метод превосходит расписание 1F1B до 23% по пропускной способности при аналогичном ограничении памяти. Этот показатель может быть увеличен до 31% при ослаблении ограничений памяти. Мы считаем, что наши результаты знаменуют собой важный шаг вперед в раскрытии истинного потенциала параллелизма конвейера. Мы открыли исходный код нашей реализации на основе популярного репозитория Megatron-LM по адресу https://github.com/sail-sg/zero-bubble-pipeline-parallelism.

English

Pipeline parallelism is one of the key components for large-scale distributed training, yet its efficiency suffers from pipeline bubbles which were deemed inevitable. In this work, we introduce a scheduling strategy that, to our knowledge, is the first to successfully achieve zero pipeline bubbles under synchronous training semantics. The key idea behind this improvement is to split the backward computation into two parts, one that computes gradient for the input and another that computes for the parameters. Based on this idea, we handcraft novel pipeline schedules that significantly outperform the baseline methods. We further develop an algorithm that automatically finds an optimal schedule based on specific model configuration and memory limit. Additionally, to truly achieve zero bubble, we introduce a novel technique to bypass synchronizations during the optimizer step. Experimental evaluations show that our method outperforms the 1F1B schedule up to 23% in throughput under a similar memory limit. This number can be further pushed to 31% when the memory constraint is relaxed. We believe our results mark a major step forward in harnessing the true potential of pipeline parallelism. We open sourced our implementation based on the popular Megatron-LM repository on https://github.com/sail-sg/zero-bubble-pipeline-parallelism.