Bilanciare il parallelismo del pipeline con il parallelismo del vocabolario

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

Autori: Man Tsung Yeung, Penghui Qi, Min Lin, Xinyi Wan

cs.AI

Abstract

Il parallelismo a pipeline è ampiamente utilizzato per scalare l'addestramento di modelli linguistici di grandi dimensioni basati su trasformatori, sono stati compiuti vari lavori per migliorarne la capacità di elaborazione e l'occupazione di memoria. In questo articolo, affrontiamo un problema spesso trascurato: i livelli di vocabolario possono causare un disavanzo nell'elaborazione e nell'utilizzo di memoria tra le fasi della pipeline, peggiorando le interruzioni della pipeline e il collo di bottiglia della memoria. Per affrontare questo problema, suddividiamo uniformemente i livelli del vocabolario tra i dispositivi della pipeline e raggruppiamo l'elaborazione in passaggi della pipeline. Per ridurre il sovraccarico di memoria di attivazione, proponiamo diversi algoritmi per ridurre le barriere di comunicazione all'interno dei livelli del vocabolario. Inoltre, utilizziamo un metodo generalizzabile per integrare il Parallelismo del Vocabolario con gli attuali programmi di pipeline. Combinando queste tecniche, i nostri metodi bilanciano efficacemente l'elaborazione e la memoria dei parametri, con un sovraccarico di memoria di attivazione costante e ridotto. In particolare, quando combinato con programmi di memoria di attivazione bilanciati come V-Half, il nostro approccio raggiunge un perfetto equilibrio sia nella memoria che nell'elaborazione. Valutazioni approfondite dimostrano che il nostro metodo raggiunge un equilibrio tra elaborazione e memoria indipendentemente dalle dimensioni del vocabolario, portando a un miglioramento del 5% al 51% nella capacità di elaborazione rispetto agli approcci ingenui, riducendo significativamente l'utilizzo massimo di memoria soprattutto per scenari con un ampio vocabolario. La nostra implementazione è open source su https://github.com/sail-sg/VocabularyParallelism.

English

Pipeline parallelism is widely used to scale the training of transformer-based large language models, various works have been done to improve its throughput and memory footprint. In this paper, we address a frequently overlooked issue: the vocabulary layers can cause imbalanced computation and memory usage across pipeline stages, worsening pipeline bubbles and the memory bottleneck. To tackle this, we partition the vocabulary layers evenly across pipeline devices and group the computation into pipeline passes. To reduce the activation memory overhead, we propose several algorithms to reduce communication barriers within vocabulary layers. Additionally, we utilize a generalizable method to integrate Vocabulary Parallelism with existing pipeline schedules. By combining these techniques, our methods effectively balance the computation and parameter memory, with only a small constant activation memory overhead. Notably, when combined with activation memory-balanced schedules like V-Half, our approach achieves perfect balance in both memory and computation. Extensive evaluations demonstrate that our method achieves computation and memory balance regardless of the vocabulary size, resulting in a 5% to 51% improvement in throughput compared to naive approaches, meanwhile significantly reducing peak memory usage especially for large vocabulary scenarios. Our implementation is open-sourced at https://github.com/sail-sg/VocabularyParallelism .