Die Balance zwischen Pipeline-Parallelität und Vokabular-Parallelität

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

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

cs.AI

Zusammenfassung

Die Pipeline-Parallelität wird weitgehend verwendet, um das Training von auf Transformer basierenden großen Sprachmodellen zu skalieren. Es wurden verschiedene Arbeiten durchgeführt, um die Durchsatzleistung und den Speicherbedarf zu verbessern. In diesem Papier behandeln wir ein häufig übersehenes Problem: Die Vokabularschichten können eine unausgewogene Berechnung und Speichernutzung über die Pipeline-Stufen hinweg verursachen, was zu einer Verschlechterung der Pipeline-Blasen und des Speicherengpasses führt. Um dies zu bewältigen, partitionieren wir die Vokabularschichten gleichmäßig über die Pipeline-Geräte und gruppieren die Berechnung in Pipeline-Durchläufen. Zur Reduzierung des Aktivierungsspeicher-Overheads schlagen wir mehrere Algorithmen vor, um die Kommunikationsbarrieren innerhalb der Vokabularschichten zu reduzieren. Darüber hinaus verwenden wir eine verallgemeinerbare Methode, um die Vokabular-Parallelität mit bestehenden Pipeline-Zeitplänen zu integrieren. Durch die Kombination dieser Techniken balancieren unsere Methoden die Berechnung und den Parameter-Speicher effektiv aus, mit nur geringem konstanten Aktivierungsspeicher-Overhead. Insbesondere erreicht unser Ansatz in Kombination mit aktivierungsspeicher-ausgeglichenen Zeitplänen wie V-Half eine perfekte Balance sowohl im Speicher als auch in der Berechnung. Umfangreiche Evaluationen zeigen, dass unsere Methode eine Berechnungs- und Speicherbalance unabhängig von der Vokabulargröße erreicht, was zu einer Steigerung der Durchsatzleistung um 5 % bis 51 % im Vergleich zu naiven Ansätzen führt, wobei der Spitzen-Speicherverbrauch insbesondere bei Szenarien mit großem Vokabular signifikant reduziert wird. Unsere Implementierung ist unter https://github.com/sail-sg/VocabularyParallelism als Open Source verfügbar.

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 .