Optimizing Distributed Training on Frontier for Large Language...
Optimizing Distributed Training on Frontier for Large Language Models
Sajal Dash, Isaac Lyngaas, Junqi Yin, Xiao Wang, Romain Egele, Guojing Cong, Feiyi Wang, Prasanna Balaprakash
Статья: https://arxiv.org/abs/2312.12705
Интересная статья про то, как обучали большие LLM на суперкомпьютере Frontier (https://www.olcf.ornl.gov/frontier/), номер 1 в рейтинге TOP500 (https://www.top500.org/lists/top500/list/2023/11/).
Frontier интересен тем, что он построен на процессорах и GPU от AMD, EPYC и MI250X (https://www.amd.com/en/products/accelerators/instinct/mi200/mi250x.html) соответственно. Второй суперкомпьютер в TOP500, Aurora, кстати, построен целиком на Интеле (тоже процы и GPU). И только третий, Eagle, на NVIDIA H100 и Xeon.
Обучали модели 22B, 175B (размер GPT-3) и 1T. Две последние на 1024 и 3072 MI250X (из 37,888).
Самое интересное для меня в этой истории -- это как именно обучали на этих карточках, потому что долгое время пользоваться ничем кроме Нвидии и CUDA было невозможно, а тут пошли валом работы про обучение больших сеток на AMD. Можно вспомнить финский суперкомпьютер LUMI (https://t.me/gonzo_ML/2084), поддержку в HF Transformers (https://t.me/gonzo_ML/2167), обучение RetNet в MS (https://t.me/gonzo_ML/1754). Наконец какая-то реальная конкуренция, кажется. Ну и большой масштаб и возникающие там проблемы тоже интересны, конечно.
В mixed precision обучении нужно 6 байт на каждый параметр модели (4 для модели в fp32, 2 для вычислений в fp16), 4 байта на параметр оптимизатора (для сохранения момента в Adam в fp32), и 4 байта на градиент каждого параметра (fp32). Итого, для 1T модели надо 14T параметров. Каждый узел Frontier содержит 8 MI250X каждый с 64 GB HBM, так что без model parallelism не обойтись.
Возможны варианты. В Tensor parallelism большие матрицы весов разбивается по строкам или столбцам. В Pipeline Parallelism модель разбивается по слоям на несколько живущих на одном GPU этапов, в каждом из которых несколько слоёв. Sharded Data Parallelism похож на классический Data Parallelism, но вместо того, чтобы хостить копии целой модели на каждом устройстве (анрил), на каждом устройстве размещается только вычисляемый в данный момент слой. Разные методы можно смешивать, получая гибрид (минутка рекламы, в моей книге про JAX (https://www.manning.com/books/deep-learning-with-jax) есть пример смешивания data + tensor parallelism при обучении, ну и вообще про параллелизацию получились самые большие главы, и их целых три). В текущей работе для лучшей утилизации ресурсов используют 3D parallelism включающий tensor, pipeline, и data (обычный и sharded) parallelism.
Взяли фреймворк Megatron-DeepSpeed с поддержкой разных видов параллелизма (tensor, pipeline, data, sharded data parallelism). В оригинале он заточен на NVIDIA (и является форком их Megatron-LM с добавленными фичами), в работе его адаптировали под AMD ROCm.
Адаптация включала:
1) Конвертацию CUDA кода в HIP код через утилиту hipify, компилирование so через hipcc и привязку к питону через pybind.
2) Подготовку DeepSpeed Ops (расширений CUDA, полученных в оригинале JIT компиляцией) под ROCm и отключение JIT на всякий пожарный.
3) Инициализацию PyTorch Distributed Environment хостом из SLURM.
4) Работу с AMD над получением ROCm версий нужных библиотек, типа APEX для mixed precision, FlashAttention и FlashAttention2 (для этих использовали Composable Kernel library, https://github.com/ROCm/composable_kernel).
Экспериментировали с разными стратегиями распараллеливания и их параметрами, а также с гиперпараметрами обучения, подобрали оптимум. За деталями и наблюдениями вэлкам в статью.
В итоге получили рабочую конфигурацию, достигающую 32-38% от пиковых FLOPS. Через Roofline Analysis (писал про это когда-то тут https://moocaholic.medium.com/hardware-for-deep-learning-part-3-gpu-8906c1644664#8dd5) показали, что обучение не memory-bound. Потестировали скейлинг, получили 100% weak scaling и 87.05% strong scaling efficiency на 1T модели.
Sajal Dash, Isaac Lyngaas, Junqi Yin, Xiao Wang, Romain Egele, Guojing Cong, Feiyi Wang, Prasanna Balaprakash
Статья: https://arxiv.org/abs/2312.12705
Интересная статья про то, как обучали большие LLM на суперкомпьютере Frontier (https://www.olcf.ornl.gov/frontier/), номер 1 в рейтинге TOP500 (https://www.top500.org/lists/top500/list/2023/11/).
Frontier интересен тем, что он построен на процессорах и GPU от AMD, EPYC и MI250X (https://www.amd.com/en/products/accelerators/instinct/mi200/mi250x.html) соответственно. Второй суперкомпьютер в TOP500, Aurora, кстати, построен целиком на Интеле (тоже процы и GPU). И только третий, Eagle, на NVIDIA H100 и Xeon.
Обучали модели 22B, 175B (размер GPT-3) и 1T. Две последние на 1024 и 3072 MI250X (из 37,888).
Самое интересное для меня в этой истории -- это как именно обучали на этих карточках, потому что долгое время пользоваться ничем кроме Нвидии и CUDA было невозможно, а тут пошли валом работы про обучение больших сеток на AMD. Можно вспомнить финский суперкомпьютер LUMI (https://t.me/gonzo_ML/2084), поддержку в HF Transformers (https://t.me/gonzo_ML/2167), обучение RetNet в MS (https://t.me/gonzo_ML/1754). Наконец какая-то реальная конкуренция, кажется. Ну и большой масштаб и возникающие там проблемы тоже интересны, конечно.
В mixed precision обучении нужно 6 байт на каждый параметр модели (4 для модели в fp32, 2 для вычислений в fp16), 4 байта на параметр оптимизатора (для сохранения момента в Adam в fp32), и 4 байта на градиент каждого параметра (fp32). Итого, для 1T модели надо 14T параметров. Каждый узел Frontier содержит 8 MI250X каждый с 64 GB HBM, так что без model parallelism не обойтись.
Возможны варианты. В Tensor parallelism большие матрицы весов разбивается по строкам или столбцам. В Pipeline Parallelism модель разбивается по слоям на несколько живущих на одном GPU этапов, в каждом из которых несколько слоёв. Sharded Data Parallelism похож на классический Data Parallelism, но вместо того, чтобы хостить копии целой модели на каждом устройстве (анрил), на каждом устройстве размещается только вычисляемый в данный момент слой. Разные методы можно смешивать, получая гибрид (минутка рекламы, в моей книге про JAX (https://www.manning.com/books/deep-learning-with-jax) есть пример смешивания data + tensor parallelism при обучении, ну и вообще про параллелизацию получились самые большие главы, и их целых три). В текущей работе для лучшей утилизации ресурсов используют 3D parallelism включающий tensor, pipeline, и data (обычный и sharded) parallelism.
Взяли фреймворк Megatron-DeepSpeed с поддержкой разных видов параллелизма (tensor, pipeline, data, sharded data parallelism). В оригинале он заточен на NVIDIA (и является форком их Megatron-LM с добавленными фичами), в работе его адаптировали под AMD ROCm.
Адаптация включала:
1) Конвертацию CUDA кода в HIP код через утилиту hipify, компилирование so через hipcc и привязку к питону через pybind.
2) Подготовку DeepSpeed Ops (расширений CUDA, полученных в оригинале JIT компиляцией) под ROCm и отключение JIT на всякий пожарный.
3) Инициализацию PyTorch Distributed Environment хостом из SLURM.
4) Работу с AMD над получением ROCm версий нужных библиотек, типа APEX для mixed precision, FlashAttention и FlashAttention2 (для этих использовали Composable Kernel library, https://github.com/ROCm/composable_kernel).
Экспериментировали с разными стратегиями распараллеливания и их параметрами, а также с гиперпараметрами обучения, подобрали оптимум. За деталями и наблюдениями вэлкам в статью.
В итоге получили рабочую конфигурацию, достигающую 32-38% от пиковых FLOPS. Через Roofline Analysis (писал про это когда-то тут https://moocaholic.medium.com/hardware-for-deep-learning-part-3-gpu-8906c1644664#8dd5) показали, что обучение не memory-bound. Потестировали скейлинг, получили 100% weak scaling и 87.05% strong scaling efficiency на 1T модели.
Источник: gonzo-обзоры ML статей
2024-01-16 00:00:42