Pipeline Paralelo com Bolha Zero

Zero Bubble Pipeline Parallelism

November 30, 2023

Autores: Penghui Qi, Xinyi Wan, Guangxing Huang, Min Lin

cs.AI

Resumo

O paralelismo de pipeline é um dos componentes-chave para o treinamento distribuído em larga escala, mas sua eficiência é prejudicada pelas bolhas de pipeline, que eram consideradas inevitáveis. Neste trabalho, introduzimos uma estratégia de escalonamento que, até onde sabemos, é a primeira a alcançar com sucesso zero bolhas de pipeline sob semântica de treinamento síncrono. A ideia central por trás dessa melhoria é dividir o cálculo de retropropagação em duas partes: uma que calcula o gradiente para a entrada e outra que calcula para os parâmetros. Com base nessa ideia, criamos manualmente novos esquemas de pipeline que superam significativamente os métodos de referência. Além disso, desenvolvemos um algoritmo que encontra automaticamente um escalonamento ideal com base na configuração específica do modelo e no limite de memória. Adicionalmente, para realmente alcançar zero bolhas, introduzimos uma técnica inovadora para contornar sincronizações durante a etapa do otimizador. Avaliações experimentais mostram que nosso método supera o escalonamento 1F1B em até 23% em taxa de transferência sob um limite de memória similar. Esse número pode ser ainda maior, chegando a 31%, quando a restrição de memória é relaxada. Acreditamos que nossos resultados representam um grande avanço no aproveitamento do verdadeiro potencial do paralelismo de pipeline. Disponibilizamos nossa implementação com base no popular repositório Megatron-LM em 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.