Équilibrer le parallélisme du pipeline avec le parallélisme du vocabulaire

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

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

cs.AI

Résumé

Le parallélisme de pipeline est largement utilisé pour mettre à l'échelle l'entraînement des grands modèles de langage basés sur des transformateurs, divers travaux ont été réalisés pour améliorer son débit et son empreinte mémoire. Dans cet article, nous abordons un problème fréquemment négligé : les couches de vocabulaire peuvent entraîner un déséquilibre de calcul et une utilisation inégale de la mémoire à travers les étapes du pipeline, aggravant les bulles du pipeline et l'engorgement de la mémoire. Pour résoudre cela, nous partitionnons de manière équitable les couches de vocabulaire entre les appareils du pipeline et regroupons le calcul en passes de pipeline. Pour réduire la surcharge de mémoire d'activation, nous proposons plusieurs algorithmes pour réduire les barrières de communication au sein des couches de vocabulaire. De plus, nous utilisons une méthode généralisable pour intégrer le Parallélisme de Vocabulaire avec les plannings de pipeline existants. En combinant ces techniques, nos méthodes équilibrent efficacement le calcul et la mémoire des paramètres, avec seulement une petite surcharge de mémoire d'activation constante. Notamment, lorsqu'associée à des plannings d'équilibrage de mémoire d'activation comme V-Half, notre approche atteint un équilibre parfait à la fois en termes de mémoire et de calcul. Des évaluations approfondies démontrent que notre méthode atteint un équilibre entre le calcul et la mémoire indépendamment de la taille du vocabulaire, entraînant une amélioration du débit de 5 % à 51 % par rapport aux approches naïves, tout en réduisant significativement l'utilisation maximale de la mémoire, en particulier pour les scénarios de vocabulaire volumineux. Notre implémentation est open source sur 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 .