分布式优化器

警告

使用 CUDA 张量时，当前不支持分布式优化器

公开 DistributedOptimizer，它接受一个列表 的远程参数 （） 并运行 optimizer 在参数所在的 worker 上本地进行。分布式 optimizer 可以使用任何本地 optimizer 基类来 在每个工作程序上应用渐变。RRef

类 torch.distributed.optim。DistributedOptimizer（optimizer_class， params_rref， *args， **kwargs）

DistributedOptimizer 对分散的参数进行远程引用 跨工作程序，并为每个参数在本地应用给定的优化器。

此类按顺序使用 检索特定参数的梯度。

对 的并发调用 ， 来自相同或不同的客户端，将 在每个 worker 上序列化 —— 因为每个 worker 的优化器只能工作 一次在一组渐变上。但是，不能保证 完整的 forward-backward-optimizer 序列将针对一个客户端执行 一次。这意味着正在应用的渐变可能不对应 添加到在给定 worker 上执行的最新 forward pass 中。此外，没有 保证跨 worker 的排序。

DistributedOptimizer 在启用 TorchScript 的情况下创建本地优化器 默认情况下，优化器更新不会被 Python 全局 多线程训练（例如分布式 Model Parallel） 的 Model Parallel）。目前，大多数优化器都启用了此功能。你 也可以按照 PyTorch 教程中的配方启用 TorchScript 支持 对于您自己的自定义优化器。

参数

• optimizer_class （optim.Optimizer） – 要 instantiate 在每个 worker 上。

• params_rref （list[RRef]） – 本地或远程参数的 RRef 列表 进行优化。

• args – 要传递给每个工作程序上的 Optimizer 构造函数的参数。

• 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）

执行单个优化步骤。

这将调用每个 worker 包含要优化的参数，并将阻塞，直到所有 worker 返回。提供的将用于检索 对应的是 包含应应用于参数的渐变。context_id

参数

context_id – 我们应该为其运行 optimizer 步骤。

类 torch.distributed.optim。PostLocalSGDOptimizer（optim， averager）

包装任意值并在本地 SGD 后运行， 此优化器在每个步骤中运行本地优化器。 在预热阶段之后，它会在应用本地优化器后定期对参数进行平均。

参数

• optim （Optimizer） – 本地优化器。

• averager （ModelAverager） – 要运行 post-localSGD 算法的模型平均器实例。

例：

>>> 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）

这与 相同。 但也会将 Model Averager 的 step 值恢复为 1 保存在提供的 .state_dict

如果 中没有条目 ， 它将引发警告并将 Model Averager 的 step 初始化为 0。"step"state_dict

state_dict（）

这与 相同。 但添加了一个额外的条目来记录 Model Averager 的 step 到 checkpoint 确保 reload 不会导致不必要的再次预热。

step（）

执行单个优化步骤 （参数更新）。

类 torch.distributed.optim。ZeroRedundancyOptimizer（params， optimizer_class， process_group=无， parameters_as_bucket_view=False， overlap_with_ddp=False， **defaults）

包装任意值并将其状态分片组中的等级。

共享按照 ZeRO 的描述完成。

每个 rank 中的本地 optimizer 实例仅为 负责更新大约 parameters 和 因此只需要保持优化器状态。后 参数在本地更新，每个 rank 都会将其参数广播到 所有其他 Peer 节点保持所有模型副本处于相同状态。 可与 结合使用，以减少每等级峰值 内存消耗。1 / world_size1 / world_sizeZeroRedundancyOptimizer

ZeroRedundancyOptimizer使用排序贪婪算法打包数字 每个等级的参数。每个参数都属于一个等级，并且 不按等级划分。分区是任意的，可能与 参数注册或使用顺序。

参数

params （） – 一个 S 或 s 给出所有参数，这些参数将被分片 跨等级。IterableIterable

关键字参数

• optimizer_class （） – 本地 优化。torch.nn.Optimizer

• process_group （可选） – （默认值：初始化者 ）。ProcessGrouptorch.distributedProcessGroupdist.group.WORLD

• parameters_as_bucket_view （bool， optional） – 如果 ， 参数为 打包到 Bucket 中以加快通信速度，字段指向不同偏移量的 Bucket 视图;如果 每个单独的参数都单独通信，并且每个参数都保持不变（默认值：）。Trueparam.dataFalseparams.dataFalse

• overlap_with_ddp （bool， optional） – 如果 ， 为 与 的梯度重叠 同步;这需要 （1） 函数式优化器 对于参数或具有函数式 等效和 （2） 注册 DDP 通信挂钩 由 中的 函数之一构造; 参数被打包到与 中的参数匹配的存储桶中，这意味着该参数将被忽略。 如果 ，在向后传递后不相交地运行 （按正常值）。 （默认：TrueDistributedDataParalleloptimizer_classddp_zero_hook.pyDistributedDataParallelparameters_as_bucket_viewFalseFalse)

• 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()

警告

目前，要求所有 传入的参数是相同的 dense 类型。ZeroRedundancyOptimizer

警告

如果您传递 ，请注意以下情况：给定 目前实现 overlap with 的方式是第一个 2 次或 3 次训练迭代在 优化器步骤，分别取决于 if 或 。这是因为它需要 有关使用的梯度分桶策略的信息，该策略在 第二次向前传球（如果或直到第三次） 如果 .要对此进行调整，请选择一个选项 是预置虚拟输入。overlap_with_ddp=TrueDistributedDataParallelstatic_graph=Falsestatic_graph=TrueDistributedDataParallelstatic_graph=Falsestatic_graph=True

警告

ZeroRedundancyOptimizer 是实验性的，可能会发生变化。

add_param_group（param_group）

将参数组添加到 的 .Optimizerparam_groups

这在微调预先训练的网络时非常有用，因为 freeze 层可以设置为可训练并添加到 AS 培训进度。Optimizer

参数

param_group （dict） – 指定要优化的参数，并且 特定于组的优化选项。

警告

此方法处理更新所有分区上的分片 但需要被召集到所有级别。在 军衔会导致训练挂起，因为通信 根据托管参数调用 Primitives，并且 期望所有排名都参与同一组参数。

consolidate_state_dict（to=0）

合并目标排名上的 s 列表（每个等级一个）。state_dict

参数

to （int） – 接收优化器状态的排名（默认值：0）。

提高

RuntimeError – 如果且此方法为 在此实例之前调用 已完全初始化，这在梯度桶 重建。overlap_with_ddp=TrueDistributedDataParallel

警告

这需要在所有级别上调用。

属性 join_device： device

返回 default device。

join_hook（**kwargs）

返回 ZeRO join 钩子。

它通过以下方式对不均匀的输入进行训练 在 Optimizer 步骤中隐藏 Collective 通信。

在调用此钩子之前，必须正确设置渐变。

参数

kwargs （dict） – 包含任何关键字参数的 a 在运行时修改 Join 钩子的行为;共享同一联接上下文的所有实例 manager 将转发相同的值。Joinablekwargs

这个钩子不支持任何 keyword 参数;即 是 闲置。kwargs

属性join_process_group：任意

返回进程组。

load_state_dict（state_dict）

从 input 加载与给定排名相关的状态，根据需要更新本地优化器。state_dict

参数

state_dict （dict） – 优化器状态;应为返回的对象 从对 的调用。

提高

RuntimeError – 如果且此方法为 在此实例之前调用 已完全初始化，这在梯度桶 重建。overlap_with_ddp=TrueDistributedDataParallel

state_dict（）

返回此排名已知的最后一个全局优化器状态。

提高

RuntimeError – 如果且此方法为 在此实例之前调用 已完全初始化，这在梯度桶 重建;或者，如果调用此方法时没有前面的调用 到 .overlap_with_ddp=TrueDistributedDataParallel

返回类型

Dict[str， 任意]

step（closure=None， **kwargs）

执行单个优化器步骤并同步所有等级的参数。

参数

closure （Callable） – 重新评估模型的闭包，以及 返回损失;对于大多数优化器来说是可选的。

返回

可选的损失，具体取决于底层本地优化器。

返回类型

可选[float]