Ускорьте обучение с помощью библиотеки Amazon SageMaker для параллельной обработки данных

Ускорение обучения с помощью библиотеки Amazon SageMaker для параллельной обработки данных

Тренировка моделей на крупномасштабном языковом моделировании (LLM) стала все более популярной за последний год с выпуском нескольких общедоступных моделей, таких как Llama2, Falcon и StarCoder. Теперь клиенты тренируют LLM’ы безпрецедентных размеров от 1 миллиарда до более 175 миллиардов параметров. Для обучения этих LLM’ов требуются значительные вычислительные ресурсы и время, поскольку для работы с сегодняшними огромными наборами данных для обучения и размерами модели необходимо использовать сотни и тысячи графических процессорных устройств (GPU). Одним из узких мест в распределенном обучении может быть коммуникация между GPU, которая обрабатывается библиотекой коллективной коммуникации NVIDIA Collective Communication Library (NCCL). В некоторых больших распределенных задачах обучения больше времени может уходить на меж-GPU коммуникацию, чем на само вычисление на GPU. Для устранения узкого места в коммуникации между GPU и обеспечения более быстрого обучения Amazon SageMaker с радостью объявляет об оптимизированной групповой операции AllGather в составе библиотеки распределенного параллелизма данных SageMaker (SMDDP). AllGather является наиболее используемой групповой операцией в популярных решениях памяти-эффективного параллелизма данных, таких как DeepSpeed Zero Redundancy Optimizer (ZeRO) и Fully Sharded Data Parallelism (FSDP), и она является основным компонентом для организации коммуникации на GPU. В этой статье мы покажем общий обзор работы SMDDP, как включить SMDDP в ваши скрипты обучения Amazon SageMaker и какие улучшения производительности вы можете ожидать.

Обзор решения

Традиционное параллельное обучение на данных включает репликацию всей модели на нескольких GPU, причем каждая модель обучается на различных фрагментах данных из набора данных. Во время прохода назад градиенты усредняются среди рабочих GPU, поэтому каждая реплика модели обновляется с одинаковыми значениями градиента, несмотря на то, что они обучены на различных фрагментах данных. Такая техника позволяет значительно ускорить обучение на огромных наборах данных за счет параллелизации потребления тренировочных данных. Однако некоторые современные большие модели (например, Llama2 70B) слишком велики, чтобы полностью поместиться в памяти GPU, что делает традиционное параллельное обучение невозможным. Чтобы продолжать получать преимущества параллельного обучения, преодолевая ограничения памяти GPU, становятся популярными решения параллелизма разбитых данных, такие как DeepSpeed ZeRO, PyTorch FSDP и библиотека SageMaker model parallelism library от Amazon.

В параллелизме разбитых данных, вместо репликации всей модели на рабочих GPU, параметры модели, градиенты и состояния оптимизатора разбиваются и распределяются (т.е. делятся на фрагменты) по GPU в задаче обучения. Для вычисления прямого и обратного прохода параметры собираются из фрагментов, находящихся на других рабочих GPU, чтобы сформировать один или несколько слоев модели. После выполнения вычислений эти слои освобождаются из памяти, чтобы можно было собрать следующий набор слоев. Обратите внимание, что существуют варианты параллелизма разбитых данных, где разбиваются только состояния оптимизатора и градиенты, а не параметры модели. В таком параллелизме разбитых данных также используется операция AllGather, но только перед прямым проходом для сбора модельных параметров, обновленных различными фрагментами градиента или состояния оптимизатора с других рабочих GPU. Подробнее смотрите разные этапы DeepSpeed ZeRO и стратегии разбиения FSDP с использованием SHARD_GRAD_OP для получения более подробной информации.

Групповая операция AllGather выполняется каждый раз, когда параметры разбираются – NCCL предоставляет стандартную реализацию этой операции с открытым исходным кодом. Как показано ниже, каждый рабочий GPU, участвующий в AllGather, начинает с входного буфера и заканчивает с объединением всех входных буферов от других рабочих устройств. Когда AllGather используется в параллелизме разбитых данных, входные буферы содержат фрагменты параметров модели, а большие выходные буферы содержат один или несколько слоев модели, сформированных из других фрагментов.

Before and after AllGather operation on 4 GPUs

Хотя NCCL обычно используется для AllGather в распределенном обучении, его низкоуровневая реализация не подходит для сетевой инфраструктуры инстансов Amazon Elastic Compute Cloud ( Amazon EC2 ) и, следовательно, его производительность может замедлить обучение от начала до конца. Библиотека SMDDP – это библиотека коллективных коммуникаций для графических процессоров NVIDIA, которая служит заменой NCCL и обеспечивает более высокую производительность для задач распределенного обучения с использованием PyTorch. Конкретно, SMDDP обеспечивает оптимизированную реализацию AllGather для типов экземпляров p4d / p4de.

Поскольку коллективные операции, такие как AllGather, блокируют переднюю и заднюю проходы вычислений, более быстрое выполнение этих операций напрямую приводит к более короткому общему времени обучения без побочных эффектов на сходимость. Другие коллективные операции, которые используются реже в обучении с распределенными параллельными данными, обрабатываются с использованием NCCL.

Путеводитель

Оптимизированная для AWS AllGather

Оптимизированный для AWS AllGather использует следующие техники для достижения лучшей производительности на инфраструктуре AWS по сравнению с NCCL:

  1. Мы перемещаем данные между экземплярами через сеть Elastic Fabric Adapter (EFA) с использованием образца связи всех ко всем. EFA – это низколатентная и высокопропускная сетевая решение AWS, а образец связи всех ко всем для межузлового сетевого взаимодействия более точно соответствует особенностям EFA и сетевой инфраструктуры AWS, требуя меньшего количества передач пакетов по сравнению с NCCL, который использует кольцевую или древовидную схему связи.
  2. GDRCopy для координации локального NVLink и сетевого трафика EFA. GDRCopy – это библиотека, обеспечивающая низколатентную связь между процессами ЦП и ядрами CUDA графического процессора. Благодаря этой технологии мы можем создавать конвейерную передачу данных внутри узла и между узлами.
  3. Снижение использования мультипроцессоров графического процессора для освобождения большей вычислительной мощности для ядер модели. Подключенные к AWS экземпляры P4d / P4de оснащены графическими процессорами NVIDIA A100, каждый из которых содержит 108 мультипроцессоров потоков. В то время как NCCL занимает до 24 мультипроцессоров потоков для выполнения коллективных операций, SMDDP Collectives использует только до девяти мультипроцессоров потоков. Освободившиеся мультипроцессоры потоков могут быть использованы ядрами вычисления модели для более быстрого выполнения.

Использование

SMDDP collectives нативно интегрируется с PyTorch через абстракцию группы процессов в модуле torch.distributed. Группа процессов определяет интерфейсы для распространенных коллективных операций, таких как AllGather, ReduceScatter, AllReduce, и т. д. Пользователи могут писать общий распределенный код, а затем выбирать применяемый backend, который предоставляет реализацию этих операций на основе использованного вычислительного устройства. Задачи обучения процессора часто используют бэкэнды gloo или mpi, а графические процессоры NVIDIA используют бэкэнд nccl.

Библиотека SMDDP вступает в картину, регистрируя себя как специальный бэкэнд в абстракции группы процессов. Это делается при помощи оператора импорта, который показан в следующих фрагментах кода. Затем, при выборе бэкэнда для вашей распределенной обучающей задачи на основе графического процессора просто замените nccl на smddp. Бэкэнд smddp соблюдает те же семантические правила, что и бэкэнд nccl и поддерживает те же сценарии обучения.

DeepSpeed

import smdistributed.dataparallel.torch.torch_smddpdeepspeed.init_distributed(dist_backend = "smddp") # замена "nccl" 

FSDP

import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend="smddp")  # замена "nccl"

Бенчмарки

Мы провели сравнительный анализ производительности автономной операции AllGather, в которой коллективная операция выполняется в изоляции без обучения модели. Ниже приведена образец результатов на 32-х экземплярах p4d, сравнивающих AllGather с использованием NCCL и SMDDP. Ось X представляет размер вывода AllGather, а ось Y представляет используемую сеть p4d с пропускной способностью 400 Гб/с. 4 подграфика представляют собой общие шаблоны группового взаимодействия, где участвуют 1, 2, 4 и 8 рангов на одном экземпляре p4d соответственно.

Использование сети SMDDP и NCCL AllGather на 32 узлах

Эти микробенчмарки показывают, что SMDDP превосходит NCCL по двум ключевым характеристикам:

  1. Пиковая производительность SMDDP (примерно 90% использования пропускной способности) выше, чем у NCCL (примерно 80% использования пропускной способности) во всех конфигурациях.
  2. SMDDP достигает пиковой производительности при гораздо меньших размерах буфера, чем NCCL. Это особенно улучшает скорость обучения для более маленьких моделей или когда пользователь задает маленький размер буфера AllGather в DeepSpeed (где размер AllGather не обязательно равен размеру слоя).

Бенчмарки обучения моделей

В задачах масштабного обучения, где коммуникация между GPU является значительным узким местом, SMDDP позволяет значительно увеличить скорость обучения, измеряемую в модельных TFLOPS/GPU.

Конфигурация Производительность
Модель/Обучение Кластер Решение с режимом параллелизма данных Модельные TFLOPS/GPU с NCCL Модельные TFLOPS/GPU с SMDDP % увеличения скорости
13B Llama2 Длина последовательности: 4096 Глобальный размер пакета: 4M токенов 64 экземпляра p4d.24xlarge (512 GPU NVIDIA A100) PyTorch FSDP 97.89 121.85 24.40%
65B GPT-NeoX Длина последовательности: 2048 Глобальный размер пакета: 4M токенов 64 экземпляра p4d.24xlarge (512 GPU NVIDIA A100) DeepSpeed ZeRO Stage 3* 99.23 108.66 9.50%

* Был использован репозиторий Megatron-DeepSpeed от EleutherAI. Был также включен тензорный параллелизм с тензорным параллельным уровнем восемь.

Примечание: Модельные TFLOPS/GPU основаны на расчете использования модельных FLOPS, определенных в статье, ссылка на которую приведена здесь, а в других источниках бенчмарков может быть указана производительность аппаратного обеспечения TFLOPS/GPU в качестве метрики производительности. Производительность аппаратного обеспечения TFLOPS/GPU может быть приближенно оценена как 4/3 x модельные TFLOPS/GPU.

Вывод

В этой статье мы показали, как значительно ускорить параллельное обучение с шардами данных в Amazon SageMaker всего с двумя строками изменений кода. Масштабное распределенное обучение становится все более популярным в связи с появлением LLM, но с этим масштабом приходят и высокие затраты. Уменьшая узкое место коммуникации между GPU, SMDDP помогает обучать модели быстрее и экономить вычислительные ресурсы. Вы можете найти больше примеров использования SMDDP с параллельным обучением данных в репозитории Amazon SageMaker Examples на GitHub.