Zero Bubble Pipeline Parallelisme

Samenvatting

Pipeline-parallelisme is een van de belangrijkste componenten voor grootschalige gedistribueerde training, maar de efficiëntie ervan wordt belemmerd door pipeline-bubbels die als onvermijdelijk werden beschouwd. In dit werk introduceren we een planningsstrategie die, voor zover wij weten, als eerste succesvol nul pipeline-bubbels weet te bereiken onder synchrone trainingssemantiek. De kern van deze verbetering ligt in het splitsen van de backward-berekening in twee delen: één die de gradiënt voor de invoer berekent en een ander die de gradiënt voor de parameters berekent. Op basis van dit idee hebben we nieuwe pipelineschema's ontworpen die de baseline-methoden aanzienlijk overtreffen. We ontwikkelen verder een algoritme dat automatisch een optimaal schema vindt op basis van een specifieke modelconfiguratie en geheugenlimiet. Daarnaast introduceren we een nieuwe techniek om synchronisaties tijdens de optimizer-stap te omzeilen, om zo echt nul bubbels te bereiken. Experimentele evaluaties tonen aan dat onze methode de 1F1B-planningsstrategie tot 23% overtreft in doorvoer onder een vergelijkbare geheugenlimiet. Dit getal kan verder worden opgedreven tot 31% wanneer de geheugenbeperking wordt versoepeld. Wij geloven dat onze resultaten een belangrijke stap voorwaarts betekenen in het benutten van het ware potentieel van pipeline-parallelisme. We hebben onze implementatie, gebaseerd op de populaire Megatron-LM-repository, open source gemaakt op 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.