Балансировка параллелизма конвейера с параллелизмом словаря
С ростом размеров трансформерных моделей, такие как LLM (Large Language Models), параллелизм модели стал ключевым элементом в обучении этих моделей. Различные подходы к параллелизму, такие как Zero Redundancy Optimizer (ZeRO), Tensor Parallelism (TP) и Pipeline Parallelism (PP), были предложены для решения проблем, связанных с обучением больших моделей. В этой статье мы рассмотрим проблему несбалансированного использования вычислительных ресурсов и памяти в параллелизме конвейера, вызванную слоями словаря, и предложим методы для её устранения.
Проблемы параллелизма конвейера
Параллелизм конвейера (PP) является привлекательным из-за низкой стоимости коммуникации и высокой арифметической интенсивности. Однако он сталкивается с двумя основными проблемами:
- Конвейерные зазоры (Pipeline Bubbles): Это временные промежутки, когда устройства в конвейере простаивают, что приводит к неэффективному использованию ресурсов.
- Высокое потребление памяти: Активации нескольких микропакетов должны храниться для уменьшения конвейерных зазоров, что создает узкое место в памяти.
Имбаланс из-за слоёв словаря
Типичная архитектура трансформера включает в себя слои словаря в начале и конце конвейера, которые могут значительно увеличивать вычислительную нагрузку и потребление памяти:
- Входной слой (Input Layer) добавляет значительную память, особенно при большом словаре.
- Выходной слой (Output Layer) может быть в 5 раз более вычислительно нагруженным, чем стандартные трансформерные слои, что усугубляет проблему несбалансированности.
Предлагаемое решение: Параллелизм словаря
Чтобы решить проблему несбалансированности, мы предлагаем метод, называемый Параллелизмом словаря (Vocabulary Parallelism), который включает в себя:
- Разделение слоёв словаря по измерению словаря и распределение их равномерно по всем устройствам конвейера.
- Группировка вычислений в конвейерные проходы, аналогично форвардным и бэквардным проходам трансформера.
- Интеграция с существующими расписаниями конвейера, чтобы не влиять значительно на память и эффективность.
Алгоритмы для уменьшения коммуникационных барьеров
Мы предложили несколько алгоритмов для уменьшения числа коммуникационных барьеров в слоях словаря:
- Алгоритм 1: Использует два коммуникационных барьера, что позволяет снизить нагрузку на вычислительные ресурсы.
- Алгоритм 2: Уменьшает количество барьеров до одного, что минимизирует потребление памяти активаций.
Интеграция в существующие расписания
Мы разработали методы для интеграции этих проходов словаря в существующие расписания конвейера, такие как 1F1B и V-Half:
- 1F1B: Вставляем проходы словаря между форвардными и бэквардными проходами последнего трансформерного слоя.
- V-Half: Используем V-образное размещение устройств для балансировки памяти активаций, добавляя проходы словаря.
Эксперименты и результаты
Эксперименты показали, что наш подход:
- Улучшает пропускную способность на 5% до 51% по сравнению с наивными методами.
- Значительно снижает пиковое потребление памяти, особенно при больших размерах словаря.
- Обеспечивает баланс памяти и вычислений, когда используется вместе с расписаниями, сбалансированными по памяти, такими как V-Half.
Реализация и масштабируемость
Наша реализация основана на проекте Megatron-LM, с оптимизациями для уменьшения накладных расходов на коммуникацию и вычисления. Мы также исследовали масштабируемость вычислений слоёв словаря, показав, что наш метод масштабируется эффективно с увеличением числа устройств.
Заключение и будущее направление
Предложенный метод параллелизма словаря эффективно решает проблему несбалансированности в параллелизме конвейера, улучшая производительность и снижая потребление памяти. Будущие работы могут включать оптимизацию ядер CUDA для слияния форвардных и бэквардных проходов, а также исследование применимости метода к мультимодальным моделям LLM.
Наш код открыт для использования и дальнейшего развития на GitHub: sail-sg/VocabularyParallelism.