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

torch.distributed.optim предоставляет DistributedOptimizer, который принимает список удалённых параметров (RRef) и выполняет оптимизацию локально на узлах, где эти параметры находятся. Распределённый оптимизатор может использовать любой из локальных оптимизаторов Базовый класс для применения градиентов на каждом узле.

class torch.distributed.optim.DistributedOptimizer(optimizer_class, params_rref, *args, **kwargs) [source]

DistributedOptimizer принимает удалённые ссылки на параметры, распределённые по узлам, и применяет заданный оптимизатор локально для каждого параметра.

Этот класс использует get_gradients() для получения градиентов для конкретных параметров.

Одновременные вызовы step(), как от одного, так и от разных клиентов, будут сериализованы на каждом узле — так как оптимизатор каждого узла может работать только с одним набором градиентов за раз. Однако нет гарантии, что полная последовательность прямое распространение — обратное распространение — оптимизация будет выполняться для одного клиента за раз. Это означает, что применяемые градиенты могут не соответствовать последнему прямому распространению, выполненному на данном узле. Кроме того, нет гарантии порядка между узлами.

DistributedOptimizer создаёт локальный оптимизатор с включённой по умолчанию TorchScript, чтобы обновления оптимизатора не блокировались блокировкой интерпретатора Python (GIL) в случае многопоточного обучения (например, распределённого моделирования параллельно). Эта функция в настоящее время включена для большинства оптимизаторов. Вы также можете следовать инструкции в руководствах PyTorch, чтобы включить поддержку TorchScript для собственных пользовательских оптимизаторов.

Параметры

• optimizer_class (optim.Optimizer) – класс оптимизатора, который необходимо создать на каждом узле.
• params_rref (list[RRef]) – список RRefs к локальным или удалённым параметрам для оптимизации.
• args – аргументы для передачи в конструктор оптимизатора на каждом узле.
• kwargs – аргументы для передачи в конструктор оптимизатора на каждом узле.

Пример::

>>> import torch.distributed.autograd as dist_autograd
>>> import torch.distributed.rpc as rpc
>>> from torch import optim
>>> from torch.distributed.optim import DistributedOptimizer
>>>
>>> with dist_autograd.context() as context_id:
>>>   # Forward pass.
>>>   rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
>>>   rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
>>>   loss = rref1.to_here() + rref2.to_here()
>>>
>>>   # Backward pass.
>>>   dist_autograd.backward(context_id, [loss.sum()])
>>>
>>>   # Optimizer.
>>>   dist_optim = DistributedOptimizer(
>>>      optim.SGD,
>>>      [rref1, rref2],
>>>      lr=0.05,
>>>   )
>>>   dist_optim.step(context_id)

step(context_id) [source]

Выполняет один шаг оптимизации.

Этот метод вызовет torch.optim.Optimizer.step() на каждом узле, содержащем параметры, которые нужно оптимизировать, и будет заблокирован, пока все узлы не вернут результат. Предоставленный context_id будет использоваться для получения соответствующего context, который содержит градиенты, которые должны быть применены к параметрам.

Параметры

context_id – идентификатор контекста autograd, для которого нужно выполнить шаг оптимизатора.

class torch.distributed.optim.PostLocalSGDOptimizer(optim, averager) [source]

Оборачивает произвольный torch.optim.Optimizer и выполняет постоптимизацию локального SGD. Этот оптимизатор выполняет локальный оптимизатор на каждом шаге. После этапа прогрева он периодически усредняет параметры после применения локального оптимизатора.

Параметры

• optim (Optimizer) – локальный оптимизатор.
• averager (ModelAverager) – экземпляр усреднителя модели для выполнения алгоритма пост-локального SGD.

Пример:

>>> import torch
>>> import torch.distributed as dist
>>> import torch.distributed.algorithms.model_averaging.averagers as averagers
>>> import torch.nn as nn
>>> from torch.distributed.optim import PostLocalSGDOptimizer
>>> from torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook import (
>>>   PostLocalSGDState,
>>>   post_localSGD_hook,
>>> )
>>>
>>> model = nn.parallel.DistributedDataParallel(
>>>    module, device_ids=[rank], output_device=rank
>>> )
>>>
>>> # Register a post-localSGD communication hook.
>>> state = PostLocalSGDState(process_group=None, subgroup=None, start_localSGD_iter=100)
>>> model.register_comm_hook(state, post_localSGD_hook)
>>>
>>> # Create a post-localSGD optimizer that wraps a local optimizer.
>>> # Note that ``warmup_steps`` used in ``PostLocalSGDOptimizer`` must be the same as
>>> # ``start_localSGD_iter`` used in ``PostLocalSGDState``.
>>> local_optim = torch.optim.SGD(params=model.parameters(), lr=0.01)
>>> opt = PostLocalSGDOptimizer(
>>>     optim=local_optim,
>>>     averager=averagers.PeriodicModelAverager(period=4, warmup_steps=100)
>>> )
>>>
>>> # In the first 100 steps, DDP runs global gradient averaging at every step.
>>> # After 100 steps, DDP runs gradient averaging within each subgroup (intra-node by default),
>>> # and post-localSGD optimizer runs global model averaging every 4 steps after applying the local optimizer.
>>> for step in range(0, 200):
>>>    opt.zero_grad()
>>>    loss = loss_fn(output, labels)
>>>    loss.backward()
>>>    opt.step()

load_state_dict(state_dict) [source]

Это то же самое, что и для torch.optim.Optimizer load_state_dict(), но также восстанавливает значение шага усреднителя модели до значения, сохранённого в предоставленном state_dict.

Если в state_dict отсутствует запись "step", будет выведено предупреждение, и шаг усреднителя модели будет инициализирован значением 0.

state_dict() [source]

Это то же самое, что и для torch.optim.Optimizer state_dict(), но добавляет дополнительную запись для записи шага усреднителя модели в контрольную точку, чтобы гарантировать, что перезагрузка не вызовет ненужного прогрева.

step() [source]

Выполняет один шаг оптимизации (обновление параметров).

class torch.distributed.optim.ZeroRedundancyOptimizer(params, optimizer_class, process_group=None, parameters_as_bucket_view=False, overlap_with_ddp=False, **defaults) [source]

Этот класс оборачивает произвольный optim.Optimizer и фрагментирует его состояния по рангам в группе, как описано в ZeRO. Локальный экземпляр оптимизатора в каждом ранге отвечает только за обновление примерно 1 / world_size параметров и, следовательно, нуждается только в хранении 1 / world_size состояний оптимизатора. После локального обновления параметров каждый ранг будет транслировать свои параметры всем остальным узлам, чтобы все копии модели находились в одинаковом состоянии. ZeroRedundancyOptimizer может использоваться совместно с torch.nn.parallel.DistributedDataParallel для снижения пикового потребления памяти на ранг.

ZeroRedundancyOptimizer использует алгоритм сортированного жадного выбора для упаковки определённого числа параметров на каждом ранге. Каждый параметр принадлежит одному рангу и не делится между рангами. Разбиение произвольное и может не совпадать с порядком регистрации или использования параметров.

Параметры

params (Iterable) – список torch.Tensor или dict, содержащих все параметры, которые будут фрагментированы по рангам.

Ключевые аргументы

• optimizer_class (torch.nn.Optimizer) – класс локального оптимизатора.
• process_group (ProcessGroup, необязательно) – torch.distributed ProcessGroup (по умолчанию: dist.group.WORLD инициализирован torch.distributed.init_process_group()).
• parameters_as_bucket_view (bool, необязательно) – если True, параметры упаковываются в корзины для ускорения обмена, и param.data поля указывают на представления корзин с различными смещениями; если False, каждый отдельный параметр передаётся отдельно, и каждый params.data остаётся нетронутым (по умолчанию: False).
• overlap_with_ddp (bool, необязательно) – если True, step() перекрывается с синхронизацией градиентов DistributedDataParallel; это требует (1) либо функциональный оптимизатор для optimizer_class аргумента или одного с функциональным эквивалентом и (2) регистрации DDP коммуникационного хука, построенного из одной из функций в ddp_zero_hook.py; параметры упаковываются в корзины, соответствующие тем в DistributedDataParallel, что означает, что parameters_as_bucket_view аргумент игнорируется. Если False, step() выполняется независимо после прохождения обратного прохода (как обычно). (по умолчанию: False)
• **defaults – любые дополнительные аргументы, которые передаются локальному оптимизатору.

Пример:

>>> import torch.nn as nn
>>> from torch.distributed.optim import ZeroRedundancyOptimizer
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> model = nn.Sequential(*[nn.Linear(2000, 2000).to(rank) for _ in range(20)])
>>> ddp = DDP(model, device_ids=[rank])
>>> opt = ZeroRedundancyOptimizer(
>>>     ddp.parameters(),
>>>     optimizer_class=torch.optim.Adam,
>>>     lr=0.01
>>> )
>>> ddp(inputs).sum().backward()
>>> opt.step()

Предупреждение

В настоящее время ZeroRedundancyOptimizer требует, чтобы все переданные параметры были одного и того же плотного типа.

Предупреждение

Если вы передаёте overlap_with_ddp=True, будьте осторожны с следующим: Учитывая, как в настоящее время реализовано перекрытие DistributedDataParallel с ZeroRedundancyOptimizer, первые две или три итерации обучения не выполняют обновлений параметров в шаге оптимизатора, в зависимости от того, используется ли static_graph=False или static_graph=True соответственно. Это связано с тем, что ему нужна информация о стратегии букетирования градиента, используемой DistributedDataParallel, которая не окончательно формируется до второго прохода вперёд, если static_graph=False или до третьего прохода вперёд, если static_graph=True. Для корректировки можно использовать, например, добавление фиктивных входных данных.

Предупреждение

ZeroRedundancyOptimizer является экспериментальным и может быть изменён.

add_param_group(param_group) [source]

Добавляет группу параметров к Optimizer.

Это может быть полезно при дообучении предварительно обученной сети, так как замороженные слои можно сделать обучаемыми и добавить к Optimizer по мере продвижения обучения.

Параметры

param_group (dict) – указывает на параметры, которые необходимо оптимизировать, и опции оптимизации для группы.

Предупреждение

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

consolidate_state_dict(to=0) [source]

Объединяет список state_dict (по одному на ранг) на целевом ранге.

Параметры

to (int) – ранг, который получает состояния оптимизатора (по умолчанию: 0).

Исключения

RuntimeError – если overlap_with_ddp=True и этот метод вызывается до полной инициализации этого экземпляра ZeroRedundancyOptimizer, что происходит после перестроения корзин градиентов DistributedDataParallel.

Предупреждение

Этот метод необходимо вызывать на всех рангах.

join_hook(**kwargs) [source]

Возвращает хук объединения ZeRO, который позволяет проводить обучение на неравномерных входных данных путём маскирования коллективных коммуникаций в шаге оптимизатора.

Градиенты должны быть правильно установлены перед вызовом этого хука.

Параметры

kwargs (dict) – dict содержащий любые ключевые аргументы для изменения поведения хука объединения во время выполнения; все экземпляры Joinable с общим контекстом хука объединения получают то же значение для kwargs.

Этот хук не поддерживает никаких ключевых аргументов; т.е. kwargs не используется.

load_state_dict(state_dict) [source]

Загружает состояние, относящееся к данному рангу, из входного state_dict, обновляя локальный оптимизатор по мере необходимости.

Параметры

state_dict (dict) – состояние оптимизатора; должно быть объектом, возвращённым из вызова state_dict().

Исключения

RuntimeError – если overlap_with_ddp=True и этот метод вызывается до полной инициализации этого экземпляра ZeroRedundancyOptimizer, что происходит после перестроения корзин градиентов DistributedDataParallel.

state_dict() [source]

Возвращает последнее глобальное состояние оптимизатора, известное этому рангу.

Возбуждает

RuntimeError – если overlap_with_ddp=True и этот метод вызывается до того, как этот ZeroRedundancyOptimizer экземпляр полностью инициализирован, что происходит после того, как DistributedDataParallel буферы градиента были перестроены; или если этот метод вызывается без предварительного вызова consolidate_state_dict().

Тип возвращаемого значения

Dict[str, Any]

step(closure=None, **kwargs) [source]

Выполняет один шаг оптимизатора и синхронизирует параметры по всем рангам.

Параметры

closure (Callable) – замыкание, которое переоценивает модель и возвращает потерю; необязательно для большинства оптимизаторов.

Возвращает

Необязательные потери в зависимости от базового локального оптимизатора.

Тип возвращаемого значения

Optional[float]