РАСПРЕДЕЛЕННОЕ ОБУЧЕНИЕ

Распределенное параллельное обучение — параллельное обучение модели

Параллельное обучение распределенной модели для больших моделей в PyTorch

В последние годы наблюдается экспоненциальный рост масштабов моделей глубокого обучения и проблема распределенного параллельного обучения. Например, знаменитый GPT-3 имеет 175 миллиардов параметров и 96 уровней внимания с размером пакета 3,2 млн и 499 миллиардов слов. Учебная платформа Amazon SageMaker может обеспечить пропускную способность 32 выборки в секунду на инстансах 120 мл.p4d.24xlarge и 175 миллиардах параметров. Если мы увеличим это число до 240 экземпляров, обучение полной модели займет 25 дней.

Параллелизм обучения на графических процессорах становится необходимым для больших моделей. Существует три типичных типа распределенного параллельного обучения: параллельное распределенное обучение данным, параллельное моделирование и параллельное тензорное обучение. Мы часто объединяем последние два типа в одну категорию: модельный параллелизм, а затем делим ее на два подтипа: конвейерный параллелизм и тензорный параллелизм. Здесь мы сосредоточимся на параллельном обучении распределенной модели и продемонстрируем, как разрабатывать в PyTorch.

Понимание параллельного обучения распределенной модели

Параллелизм моделей разделяет модель на несколько графических процессоров, в отличие от параллелизма данных, который воспроизводит одну и ту же модель для всех обучающих графических процессоров. Есть три критических момента для параллельного обучения модели: 1. как эффективно разделить слои модели; 2. как обучать сегментированные слои параллельно, а не последовательно; 3. как спроектировать отличную пропускную способность между узлами. Если мы не сможем сбалансировать сегменты или запустить их параллельно, мы не сможем достичь цели параллелизма моделей.

Распределенное обучение — это своего рода параллелизм обучения с несколькими узлами в кластере или пуле ресурсов. Контейнеризация упрощает масштабирование узлов, а Kubernetes хорошо их организует. Каждый узел может иметь несколько графических процессоров и несколько контейнеров. Контейнер может управлять одним или несколькими графическими процессорами. Параллелизм моделей может распределять слои модели по кластеру распределенных узлов графического процессора. Такой подход может эффективно масштабировать обучение.

Мы можем проиллюстрировать параллельное обучение распределенной модели ниже.

В приведенном выше примере в кластере есть два узла. Каждый узел имеет один или два контейнера, и каждый контейнер также имеет один или два графических процессора. Семь уровней модели распределяются между этими графическими процессорами.

У каждого графического процессора или узла есть ограничения по ресурсам. Параллелизм модели позволяет распараллелить обучение, в то время как распределенное обучение может в конечном итоге масштабироваться.

Параллелизм моделей в PyTorch

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

  1. Используйте «to(device)», чтобы указать конкретное устройство (или графический процессор) для определенных слоев (или подсетей) модели.
  2. Добавьте соответствующий метод «forward», чтобы перемещать промежуточные выходные данные между устройствами.
  3. Укажите выходы меток на том же устройстве при вызове функции потерь. А «backward()» и «torch.optim» будут автоматически обрабатывать градиенты, как если бы они работали на одном графическом процессоре.

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

import torch
import torch.nn as nn
class DummyModel(nn.Module):
    def __init__(self):
        super(DummyModel, self).__init__()
        self.net0 = nn.Linear(20, 10).to('cuda:0')
        self.relu = nn.ReLU()
        self.net1 = nn.Linear(10, 10).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net0(x.to('cuda:0')))
        return self.net1(x.to('cuda:1'))

Теперь мы можем добавить обучающий код с функцией потерь ниже:

import torch
import torch.nn as nn
import torch.optim as optim
import DummyModel
model = DummyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()
outputs = model(torch.randn(30, 10))
labels = torch.randn(30, 10).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()

Та же идея может быть быстро распространена на более сложные модели. Мы можем сгруппировать несколько слоев в «Sequential» и распределить их по конкретному графическому процессору через «to(device).». Конечно, есть и другие улучшения для оптимизации параллельной эффективности, но мы не будем их здесь рассматривать.

TL;DR

Распределенное параллельное обучение модели имеет две основные концепции. Параллелизм моделей реализует обучение больших моделей, которые не могут работать на одном графическом процессоре или устройстве. Распределенное обучение можно эффективно масштабировать за счет сегментирования модели на распределенных устройствах. PyTorch и другие библиотеки (например, SakeMaker) облегчают жизнь с минимальными изменениями, хотя их сложно реализовать внутри.

Рекомендации