Balancing Pipeline Parallelism with Vocabulary Parallelism

TL;DR

This paper proposes a novel method to balance computation and memory in pipeline parallelism for large language models by partitioning vocabulary layers and optimizing communication, leading to significant throughput and memory improvements.

Contribution

It introduces a new approach to evenly partition vocabulary layers and integrate vocabulary parallelism with pipeline schedules, addressing imbalance issues in large language model training.

Findings

Achieves 5% to 51% throughput improvement

Reduces peak memory usage significantly

Balances computation and memory across pipeline stages

Abstract

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 .