Balancing Pipeline Parallelism with Vocabulary Parallelism
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 .
67316a1ad0049e65ffa6a74e
Man Tsung Yeung
67316a1ad0049e65ffa6a74f
Penghui Qi
claimed_verified
2024-11-15T09:28:31.663Z
63885f1d0bebb233d8ad6e5b
https://cdn-avatars.huggingface.co/v1/production/uploads/1669881620925-noauth.jpeg
Penghui Qi
user
QPHutu
67316a1ad0049e65ffa6a750
Min Lin
67316a1ad0049e65ffa6a751
Xinyi Wan
extracted_confirmed
2024-11-11T02:39:29.434Z
63510eea0b94548566dad923
/avatars/629eaaf810718259bf7588dc2e6cc0d5.svg
Xinyi Wan
user
ufotalent
Balancing Pipeline Parallelism with Vocabulary Parallelism
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 .
Proposed techniques achieve balanced computation and memory usage in pipeline parallelism for transformer-based models by partitioning vocabulary layers and optimizing communication barriers, resulting in improved throughput and reduced peak memory usage.
pipeline parallelism transformer-based models vocabulary layers pipeline stages pipeline bubbles memory bottleneck vocabulary parallelism activation memory overhead communication barriers parameter memory activation memory-balanced schedules V-Half throughput memory usage