分布式数据并行 ¶

class torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, init_sync=True, process_group=None, bucket_cap_mb=None, find_unused_parameters=False, check_reduction=False, gradient_as_bucket_view=False, static_graph=False, delay_all_reduce_named_params=None, param_to_hook_all_reduce=None, mixed_precision=None, device_mesh=None)[source][source]¶

基于模块级别实现基于 torch.distributed 的分布式数据并行性。

此容器通过在每个模型副本之间同步梯度来提供数据并行性。要同步的设备由输入 process_group 指定，默认情况下为整个世界。请注意， DistributedDataParallel 不会将输入分块或以其他方式分片到参与 GPU 上；用户负责定义如何进行此操作，例如通过使用 DistributedSampler 。

参见：基础知识和使用 nn.parallel.DistributedDataParallel 替代 multiprocessing 或 nn.DataParallel。输入约束与 torch.nn.DataParallel 相同。

创建此类需要先通过调用 torch.distributed.init_process_group() 初始化 torch.distributed 。

DistributedDataParallel 已被证明在单节点多 GPU 数据并行训练中比 torch.nn.DataParallel 快得多。

要在具有 N 个 GPU 的主机上使用 DistributedDataParallel ，您应该启动 N 个进程，确保每个进程仅独占使用从 0 到 N-1 的单个 GPU。这可以通过为每个进程设置 CUDA_VISIBLE_DEVICES 或通过以下方式实现：

>>> torch.cuda.set_device(i)

i 的取值范围从 0 到 N-1。在每个进程中，您应参考以下内容来构建此模块：

>>> torch.distributed.init_process_group(
>>>     backend='nccl', world_size=N, init_method='...'
>>> )
>>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)

为了在每个节点上启动多个进程，您可以使用 torch.distributed.launch 或 torch.multiprocessing.spawn 。

注意

请参阅 PyTorch 分布式概述以了解有关分布式训练的所有相关功能的简要介绍。

注意

DistributedDataParallel 可以与 torch.distributed.optim.ZeroRedundancyOptimizer 结合使用，以减少每个 rank 的优化器状态内存占用。请参阅 ZeroRedundancyOptimizer 配方以获取更多详细信息。

注意

nccl 后端是目前使用 GPU 时最快且强烈推荐的后端。这适用于单节点和多节点分布式训练。

注意

此模块还支持混合精度分布式训练。这意味着您的模型可以具有不同类型的参数，例如混合类型的 fp16 和 fp32 ，对这些混合类型参数的梯度下降将正常工作。

注意

如果您在一个进程中使用 torch.save 来检查点模块，而在其他进程中使用 torch.load 来恢复它，请确保为每个进程正确配置 map_location 。如果没有 map_location ， torch.load 将恢复到保存模块的设备上。

注意

当在 M 节点上使用 batch=N 进行模型训练时，如果将批次的实例（不是通常的平均）的损失求和（而不是平均），与在单个节点上使用 batch=M*N 训练的相同模型相比，梯度将小 M 倍。当您想要获得与本地训练相对应的数学等效的训练过程时，应考虑这一点。但在大多数情况下，您可以将 DistributedDataParallel 包装的模型、DataParallel 包装的模型以及单个 GPU 上的普通模型视为相同（例如，使用等效批次的相同学习率）。

注意

参数永远不会在进程之间广播。该模块对梯度执行全归约步骤，并假设它们将被优化器以相同的方式修改。缓冲区（例如 BatchNorm 统计信息）在每个迭代中从进程 0 的模块广播到系统中的所有其他副本。

注意

如果您将 DistributedDataParallel 与 Distributed RPC 框架结合使用，您应该始终使用 torch.distributed.autograd.backward() 来计算梯度，并使用 torch.distributed.optim.DistributedOptimizer 来优化参数。

示例：

>>> import torch.distributed.autograd as dist_autograd
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> import torch
>>> from torch import optim
>>> from torch.distributed.optim import DistributedOptimizer
>>> import torch.distributed.rpc as rpc
>>> from torch.distributed.rpc import RRef
>>>
>>> t1 = torch.rand((3, 3), requires_grad=True)
>>> t2 = torch.rand((3, 3), requires_grad=True)
>>> rref = rpc.remote("worker1", torch.add, args=(t1, t2))
>>> ddp_model = DDP(my_model)
>>>
>>> # Setup optimizer
>>> optimizer_params = [rref]
>>> for param in ddp_model.parameters():
>>>     optimizer_params.append(RRef(param))
>>>
>>> dist_optim = DistributedOptimizer(
>>>     optim.SGD,
>>>     optimizer_params,
>>>     lr=0.05,
>>> )
>>>
>>> with dist_autograd.context() as context_id:
>>>     pred = ddp_model(rref.to_here())
>>>     loss = loss_func(pred, target)
>>>     dist_autograd.backward(context_id, [loss])
>>>     dist_optim.step(context_id)

注意

DistributedDataParallel 目前对梯度检查点提供有限支持，使用 torch.utils.checkpoint() 。如果使用 use_reentrant=False（推荐）进行检查点，DDP 将按预期工作，没有任何限制。然而，如果使用 use_reentrant=True（默认值）进行检查点，当模型中没有未使用的参数并且每个层最多检查点一次时（确保您没有将 find_unused_parameters=True 传递给 DDP），DDP 将按预期工作。我们目前不支持层被多次检查点的情况，或者当检查点模型中有未使用参数的情况。

注意

要让非 DDP 模型加载 DDP 模型的状态字典，需要在加载前应用 consume_prefix_in_state_dict_if_present() 来去除 DDP 状态字典中的前缀“module.”。

警告

构造函数、前向方法和输出（或该模块输出的函数）的微分是分布式同步点。在可能存在不同进程执行不同代码的情况下，请考虑这一点。

警告

此模块假设在创建模型时所有参数都已注册在模型中。之后不应添加或删除任何参数。同样适用于缓冲区。

警告

此模块假设每个分布式进程的模型中的所有参数注册顺序相同。模块本身将按照模型注册参数的逆序执行梯度 allreduce 。换句话说，确保每个分布式进程具有完全相同的模型以及完全相同的参数注册顺序是用户的责任。

警告

此模块允许具有非连续行主序步长的参数。例如，您的模型可能包含一些参数的 torch.memory_format 是 torch.contiguous_format ，而其他参数的格式是 torch.channels_last 。然而，不同进程中的对应参数必须具有相同的步长。

警告

此模块不支持 torch.autograd.grad() （即，它仅在需要将梯度累积在参数的 .grad 属性中时才工作）。

警告

如果您计划使用此模块与 nccl 后端或使用 Infiniband 的 gloo 后端（与使用多个工作进程的数据加载器一起使用），请将多进程启动方法更改为 forkserver （仅限 Python 3）或 spawn 。不幸的是，Gloo（使用 Infiniband）和 NCCL2 不是 Fork 安全的，如果您不更改此设置，可能会遇到死锁。

警告

您永远不应该在用 DistributedDataParallel 包装您的模型后尝试更改模型的参数。因为，当用 DistributedDataParallel 包装模型时，构造函数将在构造时在模型的全部参数上注册额外的梯度减少函数。如果您之后更改模型的参数，梯度减少函数将不再与正确的参数集匹配。

警告

使用 DistributedDataParallel 与分布式 RPC 框架结合是实验性的，可能会发生变化。

参数:

模块（Module）- 要并行化的模块
device_ids（python:int 或 torch.device 列表） –
CUDA 设备。1）对于单设备模块， device_ids 可以包含一个设备 ID，代表对应此进程的输入模块所在的唯一 CUDA 设备。或者， device_ids 也可以是 None 。2）对于多设备模块和 CPU 模块， device_ids 必须是 None 。

当两种情况下 device_ids 都是 None 时，前向传递的输入数据和实际模块都必须放置在正确的设备上。（默认： None ）

输出设备（int 或 torch.device）- 单设备 CUDA 模块的输出位置。对于多设备模块和 CPU 模块，必须是 None ，模块本身决定输出位置。（默认：单设备模块为 device_ids[0] ）

broadcast_buffers（bool）- 标志，用于在 forward 函数开始时同步（广播）模块的缓冲区。默认： True

init_sync（bool）- 是否在初始化时同步，以验证参数形状和广播参数及缓冲区。警告：如果设置为 False，则用户必须自行确保所有 rank 上的权重相同。（默认： True ）

process_group – 用于分布式数据全归约的进程组。如果为 None ，将使用默认进程组，该进程组由 torch.distributed.init_process_group() 创建。（默认： None ）

bucket_cap_mb – DistributedDataParallel 将参数分配到多个桶中，以便每个桶的梯度下降可以与反向计算重叠。 bucket_cap_mb 控制桶的大小，单位为兆字节（MiB）。如果 None ，则使用默认大小 25 MiB。（默认： None ）

find_unused_parameters (bool) – 从被包装模块的 forward 函数的返回值中包含的所有张量遍历 autograd 图。在此图中没有接收梯度的参数将被预先标记为准备减少。此外，可能被包装模块的 forward 函数使用但未包含在损失计算中且因此也不会接收梯度的参数也将被预先标记为准备减少。（默认： False ）

check_reduction – 此参数已弃用。

gradient_as_bucket_view (bool) – 当设置为 True 时，梯度将指向不同的 allreduce 通信桶的偏移量视图。这可以减少峰值内存使用，节省的内存大小将与总梯度大小相等。此外，它避免了梯度与 allreduce 通信桶之间复制的开销。当梯度是视图时， detach_() 不能在梯度上调用。如果遇到此类错误，请参考 torch/optim/optimizer.py 中的 zero_grad() 函数作为解决方案。请注意，梯度在第一次迭代后将是视图，因此应在第一次迭代后检查峰值内存节省。

static_graph (bool) –

当设置为 True 时，DDP 知道训练图是静态的。静态图意味着 1）在整个训练循环中，使用的和未使用的参数集不会改变；在这种情况下，用户是否设置 find_unused_parameters = True 无关紧要。2）在整个训练循环中，图的训练方式不会改变（意味着没有依赖于迭代的控制流）。当 static_graph 设置为 True 时，DDP 将支持过去无法支持的案例：1）重入反向。2）多次激活检查点。3）在模型有未使用参数时进行激活检查点。4）存在位于前向函数之外的模型参数。5）当存在未使用参数时，可能提高性能，因为当 static_graph 设置为 True 时，DDP 将不会在每次迭代中搜索图以检测未使用参数。要检查您是否可以将 static_graph 设置为 True ，一种方法是在您之前模型训练结束时检查 ddp 日志数据，如果 ddp_logging_data.get("can_set_static_graph") == True ，通常您也可以设置 static_graph = True 。

示例::

>>> model_DDP = torch.nn.parallel.DistributedDataParallel(model)
>>> # Training loop
>>> ...
>>> ddp_logging_data = model_DDP._get_ddp_logging_data()
>>> static_graph = ddp_logging_data.get("can_set_static_graph")

delay_all_reduce_named_params (list of tuple of str and torch.nn.Parameter) – 当指定在 param_to_hook_all_reduce 中的参数的梯度准备好时，将延迟所有 reduce 的命名参数列表。DDP 的其他参数不适用于此参数中指定的命名参数，因为这些命名参数将被 DDP reducer 忽略。

param_to_hook_all_reduce (torch.nn.Parameter) – 用于钩子延迟所有 reduce 的参数。

变量:

module (Module) – 要并行化的模块。

示例：

>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
>>> net = torch.nn.parallel.DistributedDataParallel(model)

join(divide_by_initial_world_size=True, enable=True, throw_on_early_termination=False)[source][source]¶

DDP 训练中处理进程间输入不均匀的上下文管理器。

此上下文管理器将跟踪已加入的 DDP 进程，并通过插入集体通信操作来“阴影”前向和反向传递，以匹配非加入的 DDP 进程创建的操作。这将确保每个集体调用都有已加入的 DDP 进程的相应调用，防止在进程间输入不均匀的情况下训练时出现挂起或错误。或者，如果指定标志 throw_on_early_termination 为 True ，则所有训练器一旦某个 rank 耗尽输入，就会抛出错误，允许根据应用程序逻辑捕获和处理这些错误。

一旦所有 DDP 进程都已加入，上下文管理器将向所有进程广播对应于最后加入进程的模型，以确保所有进程中的模型相同（这是由 DDP 保证的）。

要使用此上下文管理器启用进程间输入不均匀的训练，只需将其包装在您的训练循环中即可。无需对模型或数据加载进行任何其他修改。

警告

如果这个上下文管理器包裹的模型或训练循环有额外的分布式集体操作，例如模型前向传递中的 SyncBatchNorm ，则必须启用 throw_on_early_termination 标志。这是因为此上下文管理器不了解非-DDP 集体通信。此标志将在任何一级耗尽输入时导致所有级别抛出异常，从而允许在所有级别捕获并恢复这些错误。

参数:

divide_by_initial_world_size (bool) – 如果 True ，将梯度除以初始的 world_size DDP 训练启动时的世界大小。如果 False ，将计算有效的世界大小（尚未耗尽输入的进程数）并在 allreduce 过程中将梯度除以该值。将 divide_by_initial_world_size=True 设置为确保每个输入样本（包括不均匀的输入）在贡献全局梯度方面的权重相等。这是通过始终将梯度除以初始的 world_size 来实现的，即使我们遇到不均匀的输入。如果您将其设置为 False ，我们将梯度除以剩余的节点数。这确保了与在较小的 world_size 上训练的等价性，尽管这也意味着不均匀的输入会对全局梯度贡献更多。通常，您会希望将其设置为 True ，以处理训练作业的最后几个输入不均匀的情况。在极端情况下，如果输入数量差异很大，将其设置为 False 可能会提供更好的结果。
enable (bool) – 是否启用不均匀输入检测。在已知参与进程的输入都是偶数的情况下，传入 enable=False 以禁用。默认为 True 。
抛出早期终止错误（布尔值）- 当至少一个 rank 耗尽输入时，是抛出错误还是继续训练。如果 True ，将在第一个 rank 到达数据末尾时抛出。如果 False ，将使用更小的有效 world size 继续训练，直到所有 rank 都加入。请注意，如果指定了此标志，则将忽略标志 divide_by_initial_world_size 。默认为 False 。

示例：

>>> import torch
>>> import torch.distributed as dist
>>> import os
>>> import torch.multiprocessing as mp
>>> import torch.nn as nn
>>> # On each spawned worker
>>> def worker(rank):
>>>     dist.init_process_group("nccl", rank=rank, world_size=2)
>>>     torch.cuda.set_device(rank)
>>>     model = nn.Linear(1, 1, bias=False).to(rank)
>>>     model = torch.nn.parallel.DistributedDataParallel(
>>>         model, device_ids=[rank], output_device=rank
>>>     )
>>>     # Rank 1 gets one more input than rank 0.
>>>     inputs = [torch.tensor([1]).float() for _ in range(10 + rank)]
>>>     with model.join():
>>>         for _ in range(5):
>>>             for inp in inputs:
>>>                 loss = model(inp).sum()
>>>                 loss.backward()
>>>     # Without the join() API, the below synchronization will hang
>>>     # blocking for rank 1's allreduce to complete.
>>>     torch.cuda.synchronize(device=rank)

join_hook(**kwargs)[source][source]¶

no_sync()[source][source]

用于禁用 DDP 进程间梯度同步的上下文管理器。

在此上下文中，梯度将累积在模块变量上，稍后将在退出上下文的第一个正向反向传递中进行同步。

示例：

>>> ddp = torch.nn.parallel.DistributedDataParallel(model, pg)
>>> with ddp.no_sync():
>>>     for input in inputs:
>>>         ddp(input).backward()  # no synchronization, accumulate grads
>>> ddp(another_input).backward()  # synchronize grads

警告

正向传递应包含在上下文管理器内，否则梯度仍会进行同步。

register_comm_hook(state, hook)[source][source]¶

分布式数据并行 ¶

文档

教程

资源