Pipeline Parallelo a Bolla Zero

Zero Bubble Pipeline Parallelism

November 30, 2023

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

cs.AI

Abstract

Il parallelismo pipeline è uno dei componenti chiave per l'addestramento distribuito su larga scala, ma la sua efficienza è compromessa dalle bolle pipeline, ritenute inevitabili. In questo lavoro, introduciamo una strategia di scheduling che, a nostra conoscenza, è la prima a raggiungere con successo zero bolle pipeline sotto semantica di addestramento sincrono. L'idea centrale di questo miglioramento è dividere il calcolo all'indietro in due parti: una che calcola il gradiente per l'input e un'altra che lo calcola per i parametri. Basandoci su questa idea, abbiamo progettato manualmente nuovi scheduli pipeline che superano significativamente i metodi di base. Abbiamo inoltre sviluppato un algoritmo che trova automaticamente uno schedulo ottimale in base alla configurazione specifica del modello e al limite di memoria. Inoltre, per raggiungere veramente zero bolle, introduciamo una nuova tecnica per bypassare le sincronizzazioni durante il passo dell'ottimizzatore. Le valutazioni sperimentali mostrano che il nostro metodo supera lo schedulo 1F1B fino al 23% in termini di throughput sotto un limite di memoria simile. Questo numero può essere ulteriormente spinto al 31% quando il vincolo di memoria è rilassato. Crediamo che i nostri risultati rappresentino un importante passo avanti nello sfruttare il vero potenziale del parallelismo pipeline. Abbiamo reso open source la nostra implementazione basata sul popolare repository Megatron-LM su 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.