Обучение на больших языковых моделях (LLM) становится все более популярным за последний год с выпуском нескольких общедоступных моделей, таких как Llama2, Falcon и StarCoder. Сейчас клиенты обучают LLM беспрецедентных размеров — от 1 до более чем 175 миллиардов параметров. Обучение этих LLM требует значительных вычислительных ресурсов и времени, поскольку для обработки современных огромных наборов обучающих данных и размеров моделей необходимо использовать сотни и тысячи графических процессоров (GPU). Одним из узких мест в распределенном обучении может быть связь с графическим процессором, обслуживаемая библиотекой коллективных коммуникаций NVIDIA (NCCL). В некоторых крупнораспределенных учебных заданиях больше времени может быть потрачено на взаимодействие между графическими процессорами, чем на фактические вычисления на графическом процессоре. Чтобы устранить узкие места связи с графическим процессором и обеспечить более быстрое обучение, Amazon SageMaker рад объявить об оптимизированной коллективной операции AllGather как части параллельной библиотеки распределенных данных SageMaker (SMDDP). AllGather — наиболее часто используемая коллективная операция в популярных решениях для эффективного использования памяти, таких как Оптимизатор нулевой избыточности DeepSpeed ​​(ZeRO) и Полностью сегментированный параллелизм данных (FSDP), и это основной источник накладных расходов на связь с графическим процессором. В этом посте мы покажем общий обзор того, как работает SMDDP, как включить SMDDP в сценарии обучения Amazon SageMaker, а также ожидаемое повышение производительности.

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

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

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

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

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

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

Прохождение

AllGather, оптимизированный для AWS

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

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

Применение

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

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

ДипСпид

import smdistributed.dataparallel.torch.torch_smddp
deepspeed.init_distributed(dist_backend="smddp")  # replacing "nccl"

ФСДП

import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp")  # replacing "nccl"

Тесты

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

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

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

Тесты обучения модели

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

*Megatron-DeepSpeed ​​от EleutherAI использовался репозиторий. Тензорный параллелизм также был включен со степенью тензорной параллели, равной восьми.

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

Заключение

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

Об авторах

Апурв Гупта — инженер-разработчик программного обеспечения в AWS, специализирующийся на создании оптимальных систем глубокого обучения для инфраструктуры и оборудования AWS. Он интересуется распределенными вычислениями, системами глубокого обучения и ускорителями машинного обучения. Помимо работы Апурв любит путешествовать, ходить в походы и видеоигры.

Каран Диман — инженер-разработчик программного обеспечения в AWS, расположенном в Торонто, Канада. Он очень увлечен областью машинного обучения и созданием решений для ускорения распределенных вычислительных задач.

Рухан Прасад — инженер-разработчик программного обеспечения в AWS, который работает над тем, чтобы сделать распределенное обучение глубокому обучению более быстрым, дешевым и простым в использовании в SageMaker. Помимо работы Рухан любит играть в теннис, путешествовать и готовить.

Чжаоци Чжу — старший инженер по разработке программного обеспечения в AWS, увлечен распределенными системами и оптимизацией низкого уровня. Ему нравится смотреть футбольные матчи и пить (недиетическую) газировку.