Архитектура распределённого автограда

Общие сведения

Предположим, у вас есть два узла и очень простая модель, разделённая между ними. Её можно реализовать с помощью torch.distributed.rpc следующим образом:

import torch
import torch.distributed.rpc as rpc

def my_add(t1, t2):
  return torch.add(t1, t2)

# On worker 0:
t1 = torch.rand((3, 3), requires_grad=True)
t2 = torch.rand((3, 3), requires_grad=True)

# Perform some computation remotely.
t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))

# Perform some computation locally based on remote result.
t4 = torch.rand((3, 3), requires_grad=True)
t5 = torch.mul(t3, t4)

# Compute some loss.
loss = t5.sum()

Основная цель распределённого автограда — обеспечить выполнение обратного прохода для таких распределённых моделей с использованием loss, которые мы вычислили, и запись соответствующих градиентов для всех тензоров, которым требуются градиенты.

Запись автограда во время прямого прохода

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

Для распределённого автограда необходимо отслеживать все RPC во время прямого прохода, чтобы обеспечить корректное выполнение обратного прохода. Для этого при выполнении RPC мы добавляем в граф автограда функции send и recv.

• Функция send добавляется в источник RPC, а её выходные рёбра указывают на функцию автограда для входных тензоров RPC. На вход этой функции во время обратного прохода поступает значение, полученное от узла назначения как результат соответствующей функции recv.
• Функция recv добавляется в узел назначения, а её входные данные извлекаются из операций, выполненных на узле назначения с использованием входных тензоров. Выходные градиенты этой функции во время обратного прохода отправляются на исходный узел соответствующей функции send.
• Каждой паре send-recv присваивается глобально уникальный autograd_message_id, позволяющий однозначно идентифицировать эту пару. Это удобно для поиска соответствующей функции на удалённом узле во время обратного прохода.
• Для RRef каждый раз, когда мы вызываем torch.distributed.rpc.RRef.to_here(), для задействованных тензоров добавляется соответствующая пара send-recv.

Например, граф автограда для приведённого выше примера будет выглядеть так (для простоты t5.sum() не показан):

Контекст распределённого автограда

Каждому прямому и обратному проходу, использующему распределённый автоград, назначается уникальный torch.distributed.autograd.context, у которого есть глобально уникальный autograd_context_id. Этот контекст создаётся на каждом узле по мере необходимости.

Этот контекст выполняет следующие задачи:

1. Несколько узлов, выполняющих распределённые обратные проходы, могут накапливать градиенты для одного и того же тензора, поэтому поле .grad тензора будет содержать градиенты из различных распределённых обратных проходов, прежде чем мы успеем запустить оптимизатор. Это аналогично многократному локальному вызову torch.autograd.backward(). Чтобы разделять градиенты для каждого обратного прохода, они накапливаются в torch.distributed.autograd.context соответствующего прохода.
2. Во время прямого прохода мы сохраняем в этом контексте функции send и recv для каждого прохода автограда. Это позволяет удерживать ссылки на соответствующие узлы графа автограда, чтобы он оставался доступным. Кроме того, это упрощает поиск нужных функций send и recv во время обратного прохода.
3. В целом мы также используем этот контекст для хранения метаданных каждого распределённого прохода автограда.

С точки зрения пользователя контекст автограда настраивается следующим образом:

import torch.distributed.autograd as dist_autograd
with dist_autograd.context() as context_id:
  loss = model.forward()
  dist_autograd.backward(context_id, loss)

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

Распределённый обратный проход

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

import torch
a = torch.rand((3, 3), requires_grad=True)
b = torch.rand((3, 3), requires_grad=True)
c = torch.rand((3, 3), requires_grad=True)
d = a + b
e = b * c
d.sum().backward()

import torch
import torch.distributed.rpc as rpc

a = torch.rand((3, 3), requires_grad=True)
b = torch.rand((3, 3), requires_grad=True)
c = torch.rand((3, 3), requires_grad=True)

d = rpc.rpc_sync("worker1", torch.add, args=(a, b))
e = rpc.rpc_sync("worker1", torch.mul, args=(b, c))
loss = d.sum()

import torch
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc

def my_add(t1, t2):
  return torch.add(t1, t2)

# On worker 0:

# Setup the autograd context. Computations that take
# part in the distributed backward pass must be within
# the distributed autograd context manager.
with dist_autograd.context() as context_id:
  t1 = torch.rand((3, 3), requires_grad=True)
  t2 = torch.rand((3, 3), requires_grad=True)

  # Perform some computation remotely.
  t3 = rpc.rpc_sync("worker1", my_add, args=(t1, t2))

  # Perform some computation locally based on remote result.
  t4 = torch.rand((3, 3), requires_grad=True)
  t5 = torch.mul(t3, t4)

  # Compute some loss.
  loss = t5.sum()

  # Run the backward pass.
  dist_autograd.backward(context_id, [loss])

  # Retrieve the gradients from the context.
  dist_autograd.get_gradients(context_id)

Распределённый оптимизатор

DistributedOptimizer работает следующим образом:

1. Принимает список удалённых параметров (RRef) для оптимизации. Это также могут быть локальные параметры, обёрнутые в локальный RRef.
2. Принимает класс Optimizer в качестве локального оптимизатора, который будет запущен для всех уникальных владельцев RRef.
3. Распределённый оптимизатор создаёт экземпляр локального Optimizer на каждом рабочем узле и хранит RRef на эти экземпляры.
4. При вызове torch.distributed.optim.DistributedOptimizer.step() распределённый оптимизатор использует RPC для удалённого запуска всех локальных оптимизаторов на соответствующих рабочих узлах. В качестве входных данных для torch.distributed.optim.DistributedOptimizer.step() необходимо передать контекст распределённого автограда context_id. Локальные оптимизаторы используют его для применения градиентов, сохранённых в соответствующем контексте.
5. Если несколько параллельно работающих распределённых оптимизаторов обновляют одни и те же параметры на рабочем узле, эти обновления сериализуются с помощью блокировки.

Простой сквозной пример

Объединив всё вместе, получим простой сквозной пример использования распределённого автограда и распределённого оптимизатора. Если сохранить код в файл с именем dist_autograd_simple.py, его можно запустить командой MASTER_ADDR="localhost" MASTER_PORT=29500 python dist_autograd_simple.py:

import torch
import torch.multiprocessing as mp
import torch.distributed.autograd as dist_autograd
from torch.distributed import rpc
from torch import optim
from torch.distributed.optim import DistributedOptimizer

def random_tensor():
    return torch.rand((3, 3), requires_grad=True)

def _run_process(rank, dst_rank, world_size):
    name = "worker{}".format(rank)
    dst_name = "worker{}".format(dst_rank)

    # Initialize RPC.
    rpc.init_rpc(
        name=name,
        rank=rank,
        world_size=world_size
    )

    # Use a distributed autograd context.
    with dist_autograd.context() as context_id:
        # Forward pass (create references on remote nodes).
        rref1 = rpc.remote(dst_name, random_tensor)
        rref2 = rpc.remote(dst_name, random_tensor)
        loss = rref1.to_here() + rref2.to_here()

        # Backward pass (run distributed autograd).
        dist_autograd.backward(context_id, [loss.sum()])

        # Build DistributedOptimizer.
        dist_optim = DistributedOptimizer(
        optim.SGD,
        [rref1, rref2],
        lr=0.05,
        )

        # Run the distributed optimizer step.
        dist_optim.step(context_id)

def run_process(rank, world_size):
    dst_rank = (rank + 1) % world_size
    _run_process(rank, dst_rank, world_size)
    rpc.shutdown()

if __name__ == '__main__':
  # Run world_size workers
  world_size = 2
  mp.spawn(run_process, args=(world_size,), nprocs=world_size)