Large Parallelism Post: Part I

В этой части разобраны самые основные методы параллельного обучения сетей - Data Parallel, Distributed Data Parallel, Model Parallelism и Pipeline Parallelism.

Source: Model Parallelism HF, GPipe

Data Parallel (DP)

Классический параллелизм - параллелизм данных, реализованный в PyTorch и применяющийся одной строчкой:

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

В этом случае модель сначала копируется с "0" GPU на остальные, а данные делятся между карточками. Forward и Backward проходят на каждом GPU отдельно, а результаты полученных градиентов после каждого прохода суммируются в модель на "0" GPU. С нового цикла Forward и Backward все начинается заново.

Важные уточнения:

1. При каждом Forward проходе модель копируется с "0" GPU на остальные - то есть если существуют процесс модели на "1" карте, то он будет потерян.
2. Тензоры распределяются между GPU, но типы tuple, list и dict будут скопированы (именно скопированы, без deepcopy). Остальные типы данных будут использоваться сразу всеми карточкмаи, поэтому будут повреждены. Так что лучше все сразу переводить в тензора.
3. Очень важно подчеркнуть, что DataParallel использует многопоточность GPU, что снижает их эффективность.
4. DataParallel может параллелить GPU только на одной машине, что ограничивает размер модели и объем данных.

Прокомментирую пункт 3 подробнее. Многопоточность предполагает общую память между потоками GPU, а значит существует риск непредвиденного обновления ячеек памяти. Это может возникнуть из-за неправильного освобождения и ее выделения - такое встречалось в работе с NCCL (вот ссылка на тред).

Также необходимо добавить, что потоки Python зависят от GIL, который может блокироваться C bindings (это структурное связывание двух языков программирования C и Python - подробнее тут). Вычисления на GPU обычно используют C bindings, а значит могут иногда блокировать GIL, что остановит работу основного кода.

Distributed Data Parallel (DDP)

Для решения описанных проблем, сам PyTorch настоятельно рекомендует использовать метод DistributedDataParallel (DDP). Он вызывается также легко, как и DataParallel:

from torch.nn.parallel import DistributedDataParallel as DDP

ddp_model = DDP(model, device_ids=[rank])

Подход остался тем же самым - одна модель с "0" GPU копируется на остальные, а данные распределяются между процессами. Кстати делается это с помощью ссылки на состояние модели (state_dict), которое с помощью broadcast распространяется на все GPU. Это нужно для того, чтобы все копии модели запускались из того же самого состояния, что и оригинальная модель. При этом важно уточнить, что теперь вы не ограничены одной GPU для модели - вы работаете с машинами, где может находится сколько угодно GPU. Это означает, что для запуска больших моделей вы ограничены только бюджетом кластера.

Как я и сказал - подход остается действительно таким же, но все дело в мелочах. Например, DDP работает на мультипроцессинге, что и позволяет запускать его на нескольких машинах (и неважно сколько GPU на каждой машине). К тому же, так как каждый GPU имеет свой выделенный процесс, то это позволяет избежать перерасхода производительности, вызванного блокировками GIL. Пример работы DDP показан на рисунке ниже.

Продолжая разговор о мелочах, необходимо добавить, что каждый процесс имеет свой метод reduce, который нужен для запуска AllReduce усреднения градиентов при Forward и Backward проходах. Для повышения эффективности работы с градиентами, Reducer разбивает градиенты по buckets (размер можно настроить с помощью bucket_cap_mb), чтобы применять операции не к каждому градиенту, а сразу к их группе. Важно отметить, что параметры модели распределяются по бакетам примерно в порядке, обратном порядку Model.parameters() данной модели - это логично, ведь получение градиентов начинается с последних слоев к первым.

Optimizer Step происходит для модели на "0" процессе, а затем ее состояние снова рассылается по всем процессам и цикл начинается заново.

Model Parallelism (MP)

Очень хорошо, когда ваша модель помещается на одном GPU или даже на ноде кластера, но при работе с большими моделями такого ожидать не приходится - что делать если мы не обладаем достаточными мощностями для работы с большими моделями? Ответ очевиден - давайте параллелить модель.

Существует два варианта "резки" слоев модели - вертикальный и горизонтальный. Как обычно бывает в научной среде, в терминологиях путаются и иногда горизонтальный параллелизм называют вертикальным и наоборот (вот в этом видео Aleksa именно так и делает). Я буду придерживаться терминологии HuggingFace. Для начала расскажу про вертикальный параллелизм модели, а к горизонтальному мы вернемся в статье про Tensor Parallelism.

Представьте, что вы можете нарезать слои модели и каждую такую часть отправить на отдельный GPU.

Тогда вы будете пропускать данные сначала через первые слои модели на GPU[0], выход отправите на GPU[1], а после на GPU[2], где рассчитаете Loss, и начнете передавать градиенты с последних слоев к началу GPU[2->0].

В статье GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, к которой мы обратимся чуть позднее, изображена схема такого пайплайна, на которой сразу видны очевидные минусы данного подхода.

Пока проходит Forward или Backward pass на каждом GPU, остальные простаивают. Более того, общие данные придется копировать между каждым GPU, что займет время и ресурсы.

Чтобы решить данную проблему, разработчики Google решили придумать Pipeline Parallelism.

Pipeline Parallelism (PP)

Снова обратимся к работе GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. Pipeline Parallelism очень похож на MP - в нем также реализовано разделение модели на слои, которые хранятся на каждом GPU. Но в нем есть небольшое отличие в работе с входящими данными - каждый mini-batch разбивается на несколько micro-batches, то есть на еще более мелкие пакеты. Это позволяет обрабатывать каждый micro-batch параллельно и зон, где GPU простаивает, становится гораздо меньше. Оставшаяся зона простойки GPU называется Bubble.

Условно, вместо обработки одного большого mini-batch F0 из прошлой схемы, GPU нужно обработать четыре последовательные части F0. После расчета каждой части, результаты можно передавать на следующий GPU. Аналогичная ситуация происходит на Backward pass. Конечно, некоторый простой GPU сохраняется и в этом варианте (Bubble на рисунке), но он значительно меньше, чем при MP, что ускоряет обучение.

Спасибо, что дочитали до конца!
В следующей части углубимся в Tensor Parallelism на основе статьи Megatron-LM.