泛型连接上下文管理器

通用连接上下文管理器有助于在不均匀上进行分布式训练 输入。本页概述了相关类的 API：、 和 。有关教程，请参阅使用 Join Context Manager 进行输入不均匀的分布式训练。JoinJoinableJoinHook

类 torch.distributed.algorithms 中。Join（joinables， enable=True， throw_on_early_termination=False， **kwargs）

此类定义了通用连接上下文管理器，它允许自定义 要在进程加入后调用的钩子。这些钩子应该隐藏 未加入的进程的集体通信，以防止挂起和 错误并确保算法的正确性。指JoinHook了解有关 Hook 定义的详细信息。

警告

上下文管理器要求每个参与Joinable自 调用方法notify_join_context()在它自己的 per- iteration collective 通信来确保正确性。

警告

上下文管理器要求将 这process_groupJoinHook对象是相同的。如果有多个JoinHookobjects，则使用第一个的 。 进程组和设备信息用于检查非 加入的进程和通知进程在启用时引发异常，这两者都使用全 减少。devicethrow_on_early_termination

参数

• joinables （List[Joinable]） – 参与Joinables;它们的钩子在给定的 次序。

• enable （bool） – 启用不均匀输入检测的标志;设置为 将禁用上下文管理器的功能，并且应该 仅当用户知道输入不会不均匀时才设置 （默认值：）。FalseTrue

• throw_on_early_termination （bool） – 控制是否抛出 检测到不均匀的输入时出现异常（默认值：）。False

例：

>>> import os
>>> import torch
>>> import torch.distributed as dist
>>> import torch.multiprocessing as mp
>>> import torch.nn.parallel.DistributedDataParallel as DDP
>>> import torch.distributed.optim.ZeroRedundancyOptimizer as ZeRO
>>> from torch.distributed.algorithms.join import Join
>>>
>>> # On each spawned worker
>>> def worker(rank):
>>>     dist.init_process_group("nccl", rank=rank, world_size=2)
>>>     model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank])
>>>     optim = ZeRO(model.parameters(), torch.optim.Adam, lr=0.01)
>>>     # Rank 1 gets one more input than rank 0
>>>     inputs = [torch.tensor([1.]).to(rank) for _ in range(10 + rank)]
>>>     with Join([model, optim]):
>>>         for input in inputs:
>>>             loss = model(input).sum()
>>>             loss.backward()
>>>             optim.step()
>>>     # All ranks reach here without hanging/erroring

静态notify_join_context（可加入）

通知连接上下文管理器调用进程尚未 加入;然后，如果 ， 检查是否不均匀 检测到输入（即，如果一个进程已加入）和 如果是这样，则引发异常。throw_on_early_termination=True

此方法应从Joinable对象 它的每次迭代集体通信。例如，这应该 在 中的前向传递开始时调用 。DistributedDataParallel

只有第一个Joinable对象传递到上下文中 manager 在此方法中执行集体通信，并且 对于其他人来说，这种方法是空洞的。

参数

joinable （Joinable） – 该Joinable对象调用此 方法。

结果

all-reduce 的异步工作句柄，用于通知上下文 进程尚未加入的 manager （如果为 第一个传递到上下文管理器中; 否则。joinableNone

类 torch.distributed.algorithms 中。可加入

这为可连接类定义了一个抽象基类。可加入的类 （继承自Joinable） 应实现join_hook(), ，它返回一个JoinHook实例，除了join_device()和join_process_group()返回设备并 process group 信息。

抽象属性 join_device： device

返回从中执行集体通信的设备 Join Context Manager 实现本身需要。

摘要 join_hook（**kwargs）

返回一个JoinHook实例Joinable.

参数

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

返回类型：

JoinHook 连接

抽象属性 join_process_group：Any

返回 所需的集体通信的进程组 连接上下文管理器本身。

类 torch.distributed.algorithms 中。JoinHook

这定义了一个 join 钩子，它在 join 中提供了两个入口点 上下文管理器：一个主钩子，当存在时被重复调用 一个未加入的进程，以及一个 post-hook，它被调用一次所有进程 已加入。

要为通用连接上下文管理器实现连接钩子，请定义一个 继承自JoinHook并酌情覆盖 和。main_hook()post_hook()

main_hook（

当存在未加入的进程时，将重复调用此 hook 在一次训练迭代中影子集体通信（即在 一个 FORWARD PASS、BACKWARD PASS 和 Optimizer 步骤）。

post_hook（is_last_joiner）

此 hook 在所有进程都加入后调用。它传递了一个 additional 参数，该参数指示 rank 是最后加入的 RANK 之一。boolis_last_joiner

参数

is_last_joiner （bool） – 如果排名是最后一个 加入; 否则。TrueFalse