通用 Join 上下文管理器

创建于：2025 年 6 月 6 日 | 最后更新于：2025 年 6 月 6 日

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

class torch.distributed.algorithms.Join(joinables, enable=True, throw_on_early_termination=False, **kwargs)

此类定义了通用的 Join 上下文管理器，它允许在进程加入后调用自定义钩子。

这些钩子应该模拟未加入进程的集体通信，以防止挂起和错误，并确保算法的正确性。有关钩子定义的详细信息，请参阅 JoinHook。

警告

上下文管理器要求每个参与的 Joinable 在其自身的每个迭代集体通信之前调用 notify_join_context() 方法，以确保正确性。

警告

上下文管理器要求 JoinHook 对象中的所有 process_group 属性都相同。如果存在多个 JoinHook 对象，则使用第一个对象的 device。进程组和设备信息用于检查未加入的进程，并用于在启用 throw_on_early_termination 时通知进程抛出异常，这两者都使用 all-reduce。

参数

• joinables (List[Joinable]) – 参与的 Joinable s 的列表；它们的钩子按给定顺序迭代。

• enable (bool) – 一个启用不均匀输入检测的标志；将其设置为 False 会禁用上下文管理器的功能，仅应在用户知道输入不会不均匀时才设置（默认值：True）。

• 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

static notify_join_context(joinable)

通知 Join 上下文管理器调用进程尚未加入。

然后，如果 throw_on_early_termination=True，则检查是否检测到不均匀输入（即是否有进程已加入），并在此情况下抛出异常。

此方法应从 Joinable 对象在其每个迭代集体通信的开始处调用。例如，这应该在 DistributedDataParallel 的 forward pass 开始时调用。

只有传递给上下文管理器的第一个 Joinable 对象在此方法中执行集体通信，而对于其他对象，此方法是空的。

参数

joinable (Joinable) – 调用此方法的 Joinable 对象。

返回

如果 joinable 是第一个传递给上下文管理器的对象，则为 all-reduce 的异步工作句柄，用于通知上下文管理器该进程尚未加入；否则为 None。

class torch.distributed.algorithms.Joinable

这定义了一个可连接类的抽象基类。

可连接类（继承自 Joinable）应实现 join_hook()，它返回一个 JoinHook 实例，此外还需要实现 join_device() 和 join_process_group()，它们分别返回设备和进程组信息。

abstract property join_device: device

返回用于执行 Join 上下文管理器所需的集体通信的设备。

abstract join_hook(**kwargs)

为给定的 Joinable 返回一个 JoinHook 实例。

参数

kwargs (dict) – 一个包含任何关键字参数的 dict，用于在运行时修改 JoinHook 的行为；共享相同 Join 上下文管理器的所有 Joinable 实例都会收到相同的 kwargs 值。

返回类型

JoinHook

abstract property join_process_group: Any

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

class torch.distributed.algorithms.JoinHook

这定义了一个 Join Hook，它在 Join 上下文管理器中提供了两个入口点。

入口点：一个主钩子，在存在未加入的进程时会反复调用；一个后钩子，在所有进程都加入后调用一次。

要为通用 Join 上下文管理器实现 Join Hook，请定义一个继承自 JoinHook 的类，并根据需要重写 main_hook() 和 post_hook()。

main_hook()

在存在未加入进程时调用此钩子，以模拟训练迭代中的集体通信。

训练迭代，即一次前向传播、后向传播和优化器步骤。

post_hook(is_last_joiner)

在所有进程都加入后调用钩子。

它会传入一个额外的 bool 参数 is_last_joiner，指示该 rank 是否是最后加入的 rank 之一。

参数

is_last_joiner (bool) – 如果 rank 是最后加入的 rank 之一，则为 True；否则为 False。