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

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 (список[RRef]) – список RRef к локальным или удалённым параметрам для оптимизации.
• 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 – идентификатор контекста автографа, для которого должен быть выполнен шаг оптимизатора.

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.

Если записи "step" в state_dict нет, будет выведено предупреждение и значение шага усреднителя модели будет инициализировано в 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]