评价此页

★ ★ ★ ★ ★

分布式优化器#

创建日期：2021年3月1日 | 最近更新日期：2025年6月16日

警告

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

torch.distributed.optim 提供了 DistributedOptimizer，它接受一组远程参数引用 (RRef)，并在参数所在的 worker 上本地运行优化器。分布式优化器可以使用任何本地优化器基类在每个 worker 上应用梯度。

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

DistributedOptimizer 接受散布在各个 worker 上的参数远程引用，并在本地为每个参数应用给定的优化器。

该类使用 get_gradients() 来检索特定参数的梯度。

对 step() 的并发调用（无论是来自相同还是不同的客户端）将在每个 worker 上进行串行化——因为每个 worker 的优化器一次只能处理一组梯度。但是，不保证为一个客户端完整执行前向-反向-优化器的操作序列。这意味着所应用的梯度可能不对应于给定 worker 上执行的最新前向传递。此外，不保证 worker 之间的顺序。

DistributedOptimizer 在创建本地优化器时默认启用了 TorchScript，因此在多线程训练（例如分布式模型并行）的情况下，优化器更新不会被 Python 全局解释器锁 (GIL) 阻塞。目前大多数优化器都支持此功能。您还可以按照 PyTorch 教程中的方案为您自己的自定义优化器启用 TorchScript 支持。

参数:

optimizer_class (optim.Optimizer) – 在每个 worker 上实例化的优化器类。
params_rref (list[RRef]) – 指向要优化的本地或远程参数的 RRef 列表。
args – 传递给每个 worker 上的优化器构造函数的参数。
kwargs – 传递给每个 worker 上的优化器构造函数的关键字参数。

示例：

>>> 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 上调用 torch.optim.Optimizer.step()，并阻塞直到所有 worker 返回。提供的 context_id 将用于检索对应的 上下文，该上下文包含应应用于参数的梯度。

参数:: context_id – 我们应该为其运行优化器步进的 autograd 上下文 ID。

class torch.distributed.optim.PostLocalSGDOptimizer(optim, averager)[源码]#

包装一个任意的 torch.optim.Optimizer 并运行后局部 SGD (post-local 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)[源码]#

这与 torch.optim.Optimizer 的 load_state_dict() 相同，但还会将模型平均器的步数值恢复为提供的 state_dict 中保存的值。

如果 state_dict 中没有 "step" 条目，它将发出警告并将模型平均器的步数初始化为 0。

state_dict()[源码]#: 这与 torch.optim.Optimizer 的 state_dict() 相同，但增加了一个额外条目，用于将模型平均器的步数记录到检查点中，以确保重新加载不会再次导致不必要的预热。

step()[源码]#: 执行单个优化步进（参数更新）。

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

包装一个任意的 optim.Optimizer 并在组内的各 rank 之间分片其状态。

共享方式如 ZeRO 中所述。

每个 rank 中的本地优化器实例仅负责更新大约 1 / world_size 的参数，因此仅需要保留 1 / world_size 的优化器状态。在本地更新参数后，每个 rank 将其参数广播给所有其他对等点，以使所有模型副本保持在相同状态。ZeroRedundancyOptimizer 可以与 torch.nn.parallel.DistributedDataParallel 结合使用，以减少每个 rank 的峰值显存消耗。

ZeroRedundancyOptimizer 使用排序贪婪算法在每个 rank 打包一定数量的参数。每个参数属于单个 rank，不在 rank 之间划分。划分是任意的，可能与参数注册或使用顺序不匹配。

参数:

params (Iterable) – 包含所有参数的 Iterable（torch.Tensor 或 dict），这些参数将在各 rank 间分片。

关键字参数:

optimizer_class (torch.nn.Optimizer) – 本地优化器的类。
process_group (ProcessGroup, 可选) – torch.distributed 进程组 (默认：由 torch.distributed.init_process_group() 初始化的 dist.group.WORLD)。
parameters_as_bucket_view (bool, 可选) – 如果为 True，参数将被打包到桶中以加速通信，并且 param.data 字段指向不同偏移处的桶视图；如果为 False，每个单独的参数将分别通信，且每个 params.data 保持不变（默认：False）。
overlap_with_ddp (bool, 可选) – 如果为 True，则 step() 与 DistributedDataParallel 的梯度同步重叠；这要求 (1) optimizer_class 参数是一个函数式优化器或具有等效函数式形式，以及 (2) 注册一个由 ddp_zero_hook.py 中的函数构建的 DDP 通信钩子；参数被打包到与 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)[源码]#

将参数组添加到 Optimizer 的 param_groups 中。

这在微调预训练网络时非常有用，因为冻结层可以随着训练的进行变为可训练并添加到 Optimizer 中。

参数:: param_group (dict) – 指定要优化的参数和特定于组的优化选项。

警告

此方法处理更新所有分区上的分片，但需要在所有 rank 上调用。在 rank 的子集上调用此方法将导致训练挂起，因为根据管理的参数调用通信原语，并期望所有 rank 参与同一组参数。

consolidate_state_dict(to=0)[源码]#

在目标 rank 上整合 state_dict 列表（每个 rank 一个）。

参数:: to (int) – 接收优化器状态的 rank（默认：0）。
抛出:: RuntimeError – 如果 overlap_with_ddp=True 且在 ZeroRedundancyOptimizer 实例完全初始化（即 DistributedDataParallel 梯度桶重建完成）之前调用此方法。

警告

这需要在所有 rank 上调用。

property join_device: device#: 返回默认设备。

join_hook(**_kwargs)[源码]#

返回 ZeRO 加入钩子 (join hook)。

它通过在优化器步进中模拟集体通信，实现在不均匀输入上的训练。

在调用此钩子之前必须正确设置梯度。

参数:: kwargs (dict) – 包含任何关键字参数的 dict，用于在运行时修改加入钩子的行为；共享同一个加入上下文管理器的所有 Joinable 实例都会被转发相同的 kwargs 值。
返回类型:: JoinHook

此钩子不支持任何关键字参数；即 kwargs 未被使用。

property join_process_group: Any#: 返回进程组。

load_state_dict(state_dict)[源码]#

从输入的 state_dict 中加载与给定 rank 相关的状态，并根据需要更新本地优化器。

参数:: state_dict (dict) – 优化器状态；应该是调用 state_dict() 返回的对象。
抛出:: RuntimeError – 如果 overlap_with_ddp=True 且在 ZeroRedundancyOptimizer 实例完全初始化（即 DistributedDataParallel 梯度桶重建完成）之前调用此方法。

state_dict()[源码]#

返回该 rank 已知的最后一个全局优化器状态。

抛出:: RuntimeError – 如果 overlap_with_ddp=True 且在 ZeroRedundancyOptimizer 实例完全初始化（即 DistributedDataParallel 梯度桶重建完成）之前调用此方法；或者在没有预先调用 consolidate_state_dict() 的情况下调用此方法。
返回类型:: dict[str, Any]

step(closure=None, **kwargs)[源码]#

执行单个优化器步进并在所有 rank 之间同步参数。

参数:: closure (Callable) – 重新评估模型并返回损失的闭包；对于大多数优化器是可选的。
返回:: 可选损失，取决于底层本地优化器。
返回类型:: float | None

注意

任何额外参数都会原样传递给基础优化器。

分布式优化器#

文档

教程

资源