Pipeline Parallèle à Bulle Zéro

Zero Bubble Pipeline Parallelism

November 30, 2023

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

cs.AI

Résumé

Le parallélisme par pipeline est l'un des composants clés pour l'entraînement distribué à grande échelle, mais son efficacité est entravée par les bulles de pipeline, considérées jusqu'ici comme inévitables. Dans ce travail, nous introduisons une stratégie d'ordonnancement qui, à notre connaissance, est la première à réussir à éliminer complètement les bulles de pipeline dans le cadre d'une sémantique d'entraînement synchrone. L'idée centrale de cette amélioration consiste à diviser le calcul rétrograde en deux parties : l'une pour calculer le gradient des entrées, et l'autre pour celui des paramètres. Sur cette base, nous avons conçu manuellement des ordonnancements de pipeline novateurs qui surpassent significativement les méthodes de référence. Nous avons également développé un algorithme capable de trouver automatiquement un ordonnancement optimal en fonction de la configuration spécifique du modèle et de la limite de mémoire. Par ailleurs, pour véritablement atteindre zéro bulle, nous introduisons une technique innovante permettant de contourner les synchronisations lors de l'étape d'optimisation. Les évaluations expérimentales montrent que notre méthode surpasse l'ordonnancement 1F1B jusqu'à 23 % en termes de débit sous une limite de mémoire similaire. Ce chiffre peut être porté à 31 % lorsque la contrainte de mémoire est assouplie. Nous pensons que nos résultats représentent une avancée majeure dans l'exploitation du véritable potentiel du parallélisme par pipeline. Nous avons rendu notre implémentation open source, basée sur le dépôt populaire Megatron-LM, disponible sur 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.