Paralelismo de Tuberías con Cero Burbujas

Zero Bubble Pipeline Parallelism

November 30, 2023

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

cs.AI

Resumen

El paralelismo de pipeline es uno de los componentes clave para el entrenamiento distribuido a gran escala, aunque su eficiencia se ve afectada por las burbujas de pipeline, que se consideraban inevitables. En este trabajo, presentamos una estrategia de planificación que, hasta donde sabemos, es la primera en lograr con éxito cero burbujas de pipeline bajo semántica de entrenamiento sincrónico. La idea clave detrás de esta mejora es dividir el cálculo de retropropagación en dos partes: una que calcula el gradiente para la entrada y otra que lo calcula para los parámetros. Basándonos en esta idea, diseñamos manualmente nuevos planes de pipeline que superan significativamente a los métodos de referencia. Además, desarrollamos un algoritmo que encuentra automáticamente un plan óptimo basado en la configuración específica del modelo y el límite de memoria. Adicionalmente, para lograr verdaderamente cero burbujas, introducimos una técnica novedosa para evitar sincronizaciones durante el paso del optimizador. Las evaluaciones experimentales muestran que nuestro método supera al plan 1F1B hasta en un 23% en rendimiento bajo un límite de memoria similar. Este número puede aumentar hasta un 31% cuando se relaja la restricción de memoria. Creemos que nuestros resultados representan un avance importante en el aprovechamiento del verdadero potencial del paralelismo de pipeline. Hemos liberado nuestra implementación basada en el popular repositorio Megatron-LM en 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.