title: Balancing Pipeline Parallelism with Vocabulary Parallelism
summary: 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 .
paper_authors._id: 67316a1ad0049e65ffa6a74e
paper_authors.name: Man Tsung Yeung
paper_authors._id: 67316a1ad0049e65ffa6a74f
paper_authors.name: Penghui Qi
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-11-15T09:28:31.663Z
paper_authors._id: 63885f1d0bebb233d8ad6e5b
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/1669881620925-noauth.jpeg
paper_authors.fullname: Penghui Qi
paper_authors.type: user
paper_authors.user: QPHutu
paper_authors._id: 67316a1ad0049e65ffa6a750
paper_authors.name: Min Lin
paper_authors._id: 67316a1ad0049e65ffa6a751
paper_authors.name: Xinyi Wan
paper_authors.status: extracted_confirmed
paper_authors.statusLastChangedAt: 2024-11-11T02:39:29.434Z
paper_authors._id: 63510eea0b94548566dad923
paper_authors.avatarUrl: /avatars/629eaaf810718259bf7588dc2e6cc0d5.svg
paper_authors.fullname: Xinyi Wan
paper_authors.type: user
paper_authors.user: ufotalent
paper_title: Balancing Pipeline Parallelism with Vocabulary Parallelism
paper_summary: 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 .
paper_ai_keywords: pipeline parallelism transformer-based large language models memory footprint vocabulary layers pipeline bubbles memory bottleneck token imbalanced computation pipeline stages pipeline devices activation memory overhead communication barriers Vocabulary Parallelism parameter-efficient fine-tuning activation memory-balanced schedules V-Half computations and memory balance memory usage peak memory usage