Het in evenwicht brengen van pijplijnparallelisme met woordenschatparallelisme

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

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

cs.AI

Samenvatting

Pipeline parallelisme wordt veel gebruikt om de training van op transformers gebaseerde grote taalmodellen te schalen, er zijn verschillende werken uitgevoerd om de doorvoer en geheugenfootprint te verbeteren. In dit artikel behandelen we een vaak over het hoofd gezien probleem: de woordenschatlagen kunnen een onevenwichtige berekening en geheugenverbruik veroorzaken over de verschillende fasen van de pipeline, wat de pipeline-onderbrekingen en het geheugenknelpunt verergert. Om dit aan te pakken, verdelen we de woordenschatlagen gelijkmatig over de pipeline-apparaten en groeperen we de berekening in pipeline-passages. Om de activatiegeheugenoverhead te verminderen, stellen we verschillende algoritmen voor om communicatiebarrières binnen de woordenschatlagen te verminderen. Daarnaast maken we gebruik van een generaliseerbare methode om Woordenschatparallelisme te integreren met bestaande pipeline-schema's. Door deze technieken te combineren, balanceren onze methoden effectief de berekening en parametergeheugen, met slechts een kleine constante activatiegeheugenoverhead. Opmerkelijk is dat, wanneer gecombineerd met activatiegeheugen-gebalanceerde schema's zoals V-Half, onze aanpak een perfecte balans bereikt in zowel geheugen als berekening. Uitgebreide evaluaties tonen aan dat onze methode berekening en geheugenbalans bereikt ongeacht de omvang van de woordenschat, wat resulteert in een verbetering van de doorvoer van 5% tot 51% in vergelijking met naïeve benaderingen, terwijl het piekgebruik van geheugen aanzienlijk wordt verminderd, vooral voor scenario's met een grote woordenschat. Onze implementatie is open source beschikbaar op 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 .