Сбалансирование параллелизма конвейера с параллелизмом словаря

Balancing Pipeline Parallelism with Vocabulary Parallelism

November 8, 2024

Авторы: Man Tsung Yeung, Penghui Qi, Min Lin, Xinyi Wan

cs.AI

Аннотация

Параллелизм конвейера широко используется для масштабирования обучения больших языковых моделей на основе трансформера, было предпринято множество работ по улучшению его производительности и использования памяти. В данной статье мы обращаем внимание на часто пренебрегаемую проблему: слои словаря могут вызывать дисбаланс в вычислениях и использовании памяти на различных этапах конвейера, ухудшая проблемы с производительностью и памятью. Для решения этой проблемы мы равномерно разбиваем слои словаря по устройствам конвейера и группируем вычисления в проходы конвейера. Для снижения накладных расходов памяти активации мы предлагаем несколько алгоритмов для уменьшения барьеров взаимодействия в слоях словаря. Кроме того, мы используем обобщенный метод для интеграции Параллелизма Словаря с существующими расписаниями конвейера. Совмещая эти техники, наши методы эффективно балансируют вычисления и использование памяти параметров, при этом имеется лишь небольшой постоянный накладной расход памяти активации. Особенно стоит отметить, что при совмещении с расписаниями баланса памяти активации, такими как V-Half, наш подход достигает идеального баланса как в памяти, так и в вычислениях. Обширные оценки показывают, что наш метод достигает баланса между вычислениями и памятью независимо от размера словаря, что приводит к улучшению производительности на 5% до 51% по сравнению с наивными подходами, при этом значительно снижается пиковое использование памяти, особенно для сценариев с большим словарем. Наша реализация доступна в открытом доступе на 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 .