Equilibrando o Paralelismo de Pipeline com o Paralelismo de Vocabulário

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

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

cs.AI

Resumo

A paralelismo de pipeline é amplamente utilizada para escalar o treinamento de grandes modelos de linguagem baseados em transformadores, diversos trabalhos têm sido realizados para melhorar sua taxa de processamento e pegada de memória. Neste artigo, abordamos um problema frequentemente negligenciado: as camadas de vocabulário podem causar desequilíbrio na computação e uso de memória entre os estágios do pipeline, piorando as bolhas do pipeline e o gargalo de memória. Para lidar com isso, particionamos uniformemente as camadas de vocabulário entre os dispositivos do pipeline e agrupamos a computação em passagens do pipeline. Para reduzir a sobrecarga de memória de ativação, propomos vários algoritmos para reduzir as barreiras de comunicação dentro das camadas de vocabulário. Além disso, utilizamos um método generalizável para integrar o Paralelismo de Vocabulário com os cronogramas de pipeline existentes. Ao combinar essas técnicas, nossos métodos equilibram efetivamente a computação e a memória de parâmetros, com apenas uma pequena sobrecarga de memória de ativação constante. Notavelmente, ao ser combinada com cronogramas de memória de ativação equilibrados como o V-Half, nossa abordagem alcança um equilíbrio perfeito tanto na memória quanto na computação. Avaliações extensivas demonstram que nosso método alcança equilíbrio de computação e memória, independentemente do tamanho do vocabulário, resultando em uma melhoria de 5% a 51% na taxa de processamento em comparação com abordagens ingênuas, ao mesmo tempo em que reduz significativamente o uso máximo de memória, especialmente para cenários de vocabulário grande. Nossa implementação está disponível em código aberto em 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 .