Авторы Оливье Крючан, Эмили Уэббер и команда AWS

Облака Гавайев. Прекрасные напоминания о мире, безмятежности и потенциально впечатлениях ждут нас, когда мы выйдем из глобальной пандемии 2020 года и вернемся в мир, позволяющий путешествовать по всему миру. Но до тех пор почему бы не потратить некоторое время на изучение того, как новые библиотеки параллелизма данных и параллелизма моделей в Amazon SageMaker могут значительно повысить производительность нашей модели при распределенном обучении для глубокого обучения!

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

Следует ли вам просто принять статус-кво и ограничить свои приложения моделями, которые соответствуют существующим аппаратным возможностям или которые обучаются в приемлемое время? Конечно, нет! Используйте параллелизм моделей и параллелизм данных в Amazon SageMaker.

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

Сразу после этого мы наблюдаем впечатляющие результаты для PyTorch и TensorFlow. Используя параллелизм моделей SageMaker, мы смогли обучать модели, размер которых намного превышает объем памяти одного графического процессора. Используя параллелизм данных SageMaker, мы наблюдали улучшение пропускной способности по сравнению с PyTorch DistributedDataParallel от 15% до 41% при обучении BERT-Large (размер кластера 16–64 графических процессора) и от 4% до 13% при обучении MaskRCNN ( размер кластера 16–64 GPU).

Это означает две вещи. Во-первых, вы можете тренировать более крупные модели. Во-вторых, вы можете сделать это быстрее, чем когда-либо. Частично эти улучшения связаны с тем, что мы в полной мере используем преимущества сетевой инфраструктуры AWS и топологии инстансов EC2. Мы представляем их разработчикам менее чем в 10 строках кода.

В этом посте мы объясним, как использовать распределенные обучающие инструменты SageMaker с PyTorch, и объяснение аналогично с TensorFlow. Вы можете просмотреть наши примеры на GitHub прямо здесь. Если хотите, вы можете прямо здесь погрузиться в полное руководство по обучению TensorFlow и PyTorch.

Проблема с общими методами распределения данных для глубокого обучения

Если вы не знакомы с мини-пакетным стохастическим градиентным спуском (SGD), мы быстро напомним об этом здесь, но вам лучше не вдаваться в подробности в другом месте, таком как Оптимизация глава (https://d2l.ai/chapter_optimization/index.html) книги Dive Into Deep Learning (D2L.ai), соавторами которой являются ученые Amazon Чжан, Липтон, Ли и Смола.

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

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

Проблема с универсальными алгоритмами AllReduce заключается в том, что они имеют тенденцию замедлять процесс обучения по мере роста наборов данных. Это может значительно сократить время обучения, если не будет реализовано должным образом.

Параллелизм данных в SageMaker

Используя интеллектуальный сервер AllReduce на основе сервера параметров, SageMaker data parallelism (SDP) обучается быстрее, чем PyTorch DistributedDataParallel и Horovod, в ряде конкретных задач NLP и видения, при этом будучи совместимым с их популярными API-интерфейсами. SDP передает параметры через сбалансированные буферы слияния (BFB) одинакового размера, которые способствуют более эффективному использованию доступной полосы пропускания, чем чистая связь на основе параметров. SDP требует незначительных изменений для использования в существующем совместимом коде, и в результате вы получите более быстрое распределение заданий.

Использование параллелизма данных SageMaker в PyTorch
< br /> При запуске SDP поддерживает как PyTorch, так и TensorFlow, но здесь мы рассмотрим пример PyTorch. На самом деле это довольно доступно! Сначала вы импортируете SDK в свой скрипт. Затем вы используете SDK, чтобы сообщить загрузчику данных, на каком узле он находится, и о ранге этого узла. Вы загрузите саму модель в структуру SDP. Наконец, вы также можете указать в оценщике для своей работы, что вы собираетесь использовать этот метод распределения.

  1. Импортируйте и инициализируйте клиент SDP в своем скрипте.
  2. Оберните вашу модель DistributedDataParallel методом SDP
  3. Используйте методы SDP, чтобы управлять тем, какие узлы в вашем кластере обрабатывают определенные элементы:
  4. Какой блок данных используется каждым рабочим - это можно контролировать с помощью torch.utils.data.distributed.DistributedSampler
  5. Какой воркер сохраняет модель - в большинстве случаев мы просто сохраняем это на 0-м узле.
  6. При необходимости масштабируйте размер пакета и адаптируйте гиперпараметры, чтобы полностью занять обучающий кластер, сохраняя при этом сходимость.

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

# SDP: Import SDP PyTorch API, DDP and initialize
import smdistributed.dataparallel.torch.distributed as dist
from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP

dist.init_process_group()

def main():
    
    # Prepare dataset
    train_dataset = torchvision.datasets.MNIST(...)
 
    # SDP: Set num_replicas and rank in DistributedSampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(
            train_dataset,
            num_replicas=dist.get_world_size(),
            rank=dist.get_rank())
 
    train_loader = torch.utils.data.DataLoader(train_sampler=..., ...)
 
    # SDP: Wrap the PyTorch model with SDP’s DDP
    model = DDP(Net().to(device))
    
    # SDP: Pin each GPU to a single SDP process. 
    torch.cuda.set_device(local_rank)
    model.cuda(local_rank)
    
    # Train
    optimizer = optim.Adadelta(...)  # use your own optimizer and hyperparameters
    scheduler = StepLR(...)
    for epoch in range(1, args.epochs + 1):
        train(...)
        if rank == 0:
            test(...)
        scheduler.step()

    # SDP: Save model on master node.
    if dist.get_rank() == 0:
        torch.save(...)

Чтобы запустить распределенное обучение в SageMaker, вы можете выполнить задание с помощью SageMaker Python SDK:

from sagemaker.pytorch import PyTorch

job = PyTorch(entry_point="my_sdp_script.py",
              role=role,  # IAM role for the training cluster
              framework_version='1.6.0',
              py_version="py36",
              train_instance_count=2,
              train_instance_type="ml.p3.16xlarge",
              distributions={"smdistributed": {"dataparallel":           {"enabled": True}}} # SDP distribution method enabled here
)

job.fit({"mychannel":"s3://bucket/path/to/training/data"})
# In this case the s3 data will by copied locally at opt/ml/input/data/mychannel

Теперь, когда мы узнали, как использовать SDP, давайте посмотрим на параллелизм моделей SageMaker (SMP)!

Параллелизм моделей SageMaker

Современные инструменты параллелизма данных, такие как PyTorch DistributedDataParallel и Uber Horovod, по-прежнему требуют, чтобы разработчики подогнали хотя бы модель и одну запись для вычислительного работника. Эти текущие методы распределения блокируют инновации в рамках глубокого обучения, поскольку недавние результаты показали, что более крупные модели, как правило, связаны с более точными результатами.

При обучении очень больших моделей, таких как большие преобразователи в НЛП или модели сегментации по изображениям с высоким разрешением, это может оказаться невозможным. Альтернативный способ разделить рабочую нагрузку для размещения моделей - разделить модель на несколько вычислительных устройств. Это называется модельным параллельным распределенным обучением.

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

Исторически сложилось так, что реализация параллельного обучения моделей самостоятельно была чрезмерно сложной. Даже для экспертов в области глубокого обучения простые реализации содержат специальный код для управления распределенными вычислениями и взаимодействием между сотрудниками. В результате GPU используется неравномерно или даже не удается по-настоящему распараллелить вычисления.

Параллелизм моделей (SMP) SageMaker обрабатывает параллелизм моделей за вас с помощью нескольких строк кода. Он включает алгоритм разделения графа модели, который оптимизируется по скорости или памяти и поддерживает ручное разделение. Он обучается конвейерно по микропакетам, чтобы максимально использовать графический процессор даже при параллельном моделировании. С помощью SMP мы смогли разместить партию BERT-MEGA с 20 записями на 8 графических процессорах V100 одного экземпляра p3dn.24xlarge, и мы смогли разместить партию GPT-2 XL с 16 записями на 4 V100. Графические процессоры одного экземпляра p3.8xlarge, два достижения сложно реализовать без параллелизма моделей.

Использование параллелизма моделей SageMaker

Чтобы распространить код обучения модели PyTorch с помощью SMP, вам необходимо выполнить следующие шаги. Большинство из них аналогичны шагам для набора инструментов параллельного анализа данных выше, но мы собираемся добавить декораторы вместо шагов обучения и тестирования в самом скрипте модели.

  1. Импортируйте и инициализируйте SMP-клиент в своем скрипте.
  2. Оберните вашу модель DistributedModel методом SMP
  3. Поместите прямую и обратную логику в функцию и украсьте ее .smp.step
  4. Замените _8 _ / _ 9_ на DistributedModelObj.backward.
  5. Постобработка выходных данных микропакетов с использованием StepOutput методов, таких как reduce_mean.
  6. Вам также необходимо при необходимости написать или отредактировать код управления распределением, не зависящий от модели, например, какой работник сохраняет модель, и любые изменения размера вашего пакета и гиперпараметров, чтобы полностью занять обучающий кластер при сохранении конвергенции.

Приведенный ниже фрагмент кода суммирует типичные необходимые изменения в вымышленном учебном сценарии PyTorch:

# SMP: Import and initialize SMP API.
import smdistributed.modelparallel.torch as smp

smp.init()

model = load_my_model()

# Download dataset and create dataloader.
train_loader = load_my_data()

# SMP: Instantiate DistributedModel object using the model.
# This handles distributing the model among multiple ranks
# behind the scenes
# If horovod is enabled this will do an overlapping_all_reduce by
# default.
model = smp.DistributedModel(model)

# SMP: Define the smp.step consisting of forward and backward passes.
@smp.step()
def forward_backward(model, inputs, labels):
    loss = my_loss_function(model(inputs), labels)
    # SMP: Call backward on the model instead of the output tensor.
    # If loss is not a scalar or a 0d tensor, the backward call requires
    # out grad tensors in addition to the output tensors,
    # similar to torch.autograd.backward call.
    model.backward(loss)
    return loss

for input, label in train_loader:
    image, label = image.to(device), label.to(device)
    
    optimizer.zero_grad()
    loss_mb = forward_backward(model, inputs, labels)

    # SMP: Average the loss across microbatches.
    loss = loss_mb.reduce_mean()
    optimizer.step()

    # SMP: Print the loss only at rank 0.
    if smp.rank() == 0:
        print(f"Loss: {loss}")

Чтобы запустить распределенное выполнение этого обучающего скрипта, вы можете использовать высокоуровневый SageMaker SDK:

from sagemaker.pytorch.estimator import PyTorch

mpi_options = {
    "enabled": True,
   "processes_per_host": 4,
   "custom_mpi_options":"--mca btl_vader_single_copy_mechanism none "
  }
  
dist_options = {
    "modelparallel":{
       "enabled": True,
       "parameters": {
           "partitions": 4,  # we'll partition the model among the 4 GPUs 
           "microbatches": 8,  # Mini-batchs are split in micro-batch to increase parallelism
           "optimize": "memory" # The automatic model partitioning can optimize speed or memory
           }
       }
}


job = PyTorch(
    entry_point="my_script.py",
   role=role,  # IAM role for the training cluster
   instance_type="ml.p3.8xlarge",
   instance_count=1,
   framework_version=1.6,
   py_version="py36",
   distribution={"mpi": mpi_options, "smdistributed": dist_options}
)

job.fit()  # Can optionally use fit(wait=False) for asynchronous execution.

Мы включаем MPI, потому что SMP использует MPI как внутреннюю коммуникационную библиотеку. Количество processes_per_host должно соответствовать количеству графических процессоров, доступных в выбранном вами типе инстанса. Наконец, настраиваемый флаг MPI "--mca btl_vader_single_copy_mechanism none" помогает подавить определенные предупреждения, которые появляются, когда MPI используется с контейнерами SageMaker, и поддерживает чистоту ваших журналов обучения.

Заключительные мысли о новых запусках распределенного обучения SageMaker

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

Вы можете перейти к документации по набору инструментов для более глубокого погружения. Не стесняйтесь взглянуть и дайте нам знать, что вы строите!

Все изображения авторов