Equilibrando el Paralelismo de Tuberías con el Paralelismo de Vocabulario

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

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

cs.AI

Resumen

El paralelismo de canalización se utiliza ampliamente para escalar el entrenamiento de modelos de lenguaje grandes basados en transformadores; se han realizado diversos trabajos para mejorar su rendimiento y huella de memoria. En este documento, abordamos un problema frecuentemente pasado por alto: las capas de vocabulario pueden causar desequilibrios en la computación y el uso de memoria en las etapas de la canalización, empeorando las burbujas de la canalización y el cuello de botella de memoria. Para abordar esto, dividimos equitativamente las capas de vocabulario entre los dispositivos de la canalización y agrupamos la computación en pases de canalización. Para reducir la sobrecarga de memoria de activación, proponemos varios algoritmos para disminuir las barreras de comunicación dentro de las capas de vocabulario. Además, utilizamos un método generalizable para integrar el Paralelismo de Vocabulario con los horarios de canalización existentes. Al combinar estas técnicas, nuestros métodos equilibran eficazmente la computación y la memoria de los parámetros, con solo una pequeña sobrecarga constante de memoria de activación. Especialmente, al combinarse con horarios de memoria de activación equilibrados como V-Half, nuestro enfoque logra un equilibrio perfecto tanto en la memoria como en la computación. Evaluaciones exhaustivas demuestran que nuestro método logra un equilibrio entre la computación y la memoria independientemente del tamaño del vocabulario, lo que resulta en una mejora del 5% al 51% en el rendimiento en comparación con enfoques ingenuos, al mismo tiempo que reduce significativamente el uso máximo de memoria, especialmente para escenarios de vocabulario grande. Nuestra implementación está disponible en código abierto en 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 .