Синхронизация градиентов

Распределённый модуль PyTorch работает за счёт обмена данными между всеми GPU в вашей системе. Этот обмен занимает время, а обеспечение того, чтобы все процессы знали состояния друг друга, происходит в определённых точках срабатывания при использовании модуля ddp.

Эти точки срабатывания добавляются в модель PyTorch, а именно в её методы forward() и backward(). Это происходит, когда модель оборачивается в DistributedDataParallel:

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)

В Accelerate это преобразование выполняется автоматически при вызове prepare() с передачей вашей модели.

+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)

Замедление при накоплении градиентов

Теперь вы понимаете, что PyTorch добавляет хуки к методам forward и backward вашей модели PyTorch при обучении в распределённой конфигурации. Но как это может привести к замедлению вашего кода?

В DDP (распределённый параллелизм данных) ожидается, что определённые процессы выполняются в конкретном порядке и в определённые моменты времени, причём эти события должны происходить примерно одновременно перед переходом к следующему этапу.

Наиболее наглядный пример — обновление параметров модели через optimizer.step(). Без накопления градиентов все экземпляры модели должны вычислить, собрать и обновить свои градиенты, прежде чем переходить к следующей пачке данных. При накоплении градиентов вы накапливаете n градиентов потерь и пропускаете optimizer.step(), пока не будет обработано n пачек. Поскольку все процессы обучения нуждаются в синхронизации только к моменту вызова optimizer.step(), без каких-либо изменений в шаге обучения эта ненужная межпроцессная коммуникация может привести к значительному замедлению.

Как можно избежать этих накладных расходов?

Решение проблемы замедления

Поскольку при обучении на этих батчах вы пропускаете обновление параметров модели, их градиенты не需要同хронизироваться до момента, когда optimizer.step() действительно будет вызван. PyTorch не может автоматически определить, когда это необходимо, но предоставляет инструмент для помощи через контекстный менеджер no_sync, который добавляется к вашей модели после её преобразования в DDP.

В рамках этого контекстного менеджера PyTorch будет пропускать синхронизацию градиентов при вызове .backward(), а первый вызов .backward() вне этого контекстного менеджера инициирует синхронизацию. Пример приведён ниже:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()

В Accelerate это превращается в API, который можно вызывать независимо от устройства обучения (хотя в недистрибутивной системе он, возможно, ничего не будет делать): ddp_model.no_sync заменяется на no_sync() и работает так же:

  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()

Как и можно ожидать, функция accumulate() оборачивает эту условную проверку, отслеживая номер текущего батча, что в итоге даёт вам API накопления градиентов:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

В результате, при выборе API следует использовать либо accelerator.accumulate, либо accelerator.no_sync.

Насколько велико замедление и какие ошибки легко совершить

Для создания реалистичного примера рассмотрим следующую конфигурацию:

• Два узла с одним GPU T4 и один узел с двумя GPU
• Каждая GPU — это T4, и они размещены на GCP
• Используемый скрипт представляет собой модификацию скрипта NLP Example
• Размер батча на GPU составляет 16, а градиенты накапливаются каждые 4 шага

Все скрипты доступны в этом репозитории.

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

Насколько?

Справочник:

• Базовый вариант: не использует ни одну из описанных здесь практик синхронизации
• no_sync некорректно: no_sync только вокруг вызова backward, а не forward
• no_sync: правильное использование паттерна no_sync
• accumulate: правильное использование accumulate()

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

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

Если вы хотите убедиться, что всё выполняется корректно, мы настоятельно рекомендуем использовать функцию accumulate() и передавать gradient_accumulation_steps или gradient_accumulation_plugin в объект Accelerator, чтобы Accelerate мог управлять этим процессом за вас.

no_sync требует дополнительной памяти GPU при использовании FSDP

Обратите внимание, что отсутствие синхронизации градиентов может негативно сказаться на процессе обучения с использованием FSDP. Как предупреждалось в torch, контекстный менеджер no_sync для FSDP потребует дополнительного объема памяти.

Поэтому в ситуациях с высокой нагрузкой на память при использовании FSDP мы рекомендуем установить sync_each_batch в значение True в GradientAccumulationPlugin, чтобы отключить no_sync.

См. пример ниже, в котором мы дообучаем Mixtral (47 млрд параметров) на 8 GPU A100-80GB. Мы видим, что даже при скромном gradient_accumulation_steps=2 мы быстро сталкиваемся с нехваткой памяти (OOM), если включён no_sync. Снова это связано с дополнительными накладными расходами памяти из-за no_sync в FSDP. Однако, если no_sync отключён через sync_each_batch=True, то потребление памяти при gradient_accumulation_steps=16 возвращается к уровню gradient_accumulation_steps=1.

Отключение no_sync означает, что производительность снизится из-за дополнительных синхронизаций данных, как объяснялось в предыдущих разделах этого руководства.