Использование Local SGD с Accelerate

Local SGD — это техника распределённого обучения, при которой градиенты не синхронизируются на каждом шаге. Таким образом, каждый процесс обновляет собственную версию весов модели, а после заданного числа шагов эти веса синхронизируются путём усреднения по всем процессам. Это повышает эффективность коммуникации и может привести к значительному ускорению обучения, особенно если компьютер не оснащён более быстрым межпроцессорным соединением, таким как NVLink. В отличие от накопления градиентов (где для повышения эффективности коммуникации необходимо увеличивать эффективный размер батча), Local SGD не требует изменения размера батча или скорости обучения / расписания. Однако при необходимости Local SGD можно комбинировать с накоплением градиентов.

В этом руководстве вы узнаете, как быстро настроить Local SGD Accelerate. По сравнению со стандартной настройкой Accelerate, для этого требуется всего две дополнительные строки кода.

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

device = "cuda"
model.to(device)

gradient_accumulation_steps = 2

for index, batch in enumerate(training_dataloader):
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if (index + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

Преобразование в Accelerate

Сначала код, показанный выше, будет преобразован для использования Accelerate без помощника LocalSGD и без помощника накопления градиентов:

+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for index, batch in enumerate(training_dataloader):
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
      loss = loss / gradient_accumulation_steps
+     accelerator.backward(loss)
      if (index+1) % gradient_accumulation_steps == 0:
          optimizer.step()
          scheduler.step()

Делегирование синхронизации модели Accelerate

Осталось лишь позволить Accelerate самостоятельно управлять синхронизацией параметров модели и накоплением градиентов. Для простоты предположим, что синхронизация требуется каждые 8 шагов. Это достигается добавлением одного оператора with LocalSGD и одного вызова local_sgd.step() после каждого шага оптимизатора:

+local_sgd_steps=8

+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
    for batch in training_dataloader:
        with accelerator.accumulate(model):
            inputs, targets = batch
            outputs = model(inputs)
            loss = loss_function(outputs, targets)
            accelerator.backward(loss)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
+           local_sgd.step()

Под капотом код Local SGD отключает автоматическую синхронизацию градиентов (но накопление по-прежнему работает как ожидается!). Вместо этого он усредняет параметры модели каждые local_sgd_steps шагов (а также в конце цикла обучения).

Ограничения

Текущая реализация работает только с базовым обучением на нескольких GPU (или CPU) без, например, DeepSpeed.

Ссылки

Хотя нам не известны истинные истоки этого простого подхода, идея локального SGD довольно стара и восходит как минимум к:

Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). Parallel SGD: When does averaging help?. arXiv preprint arXiv:1606.07365.

Термин Local SGD мы приписываем следующей статье (хотя могут существовать более ранние источники, о которых нам не известно).

Stich, Sebastian Urban. «Local SGD Converges Fast and Communicates Little». ICLR 2019-International Conference on Learning Representations. No. CONF. 2019.