通用 Join 上下文管理器¶
通用 join 上下文管理器有助于在不均匀输入上进行分布式训练。此页面概述了相关类的 API:Join
、Joinable
和 JoinHook
。有关教程,请参阅 使用 Join 上下文管理器进行不均匀输入的分布式训练。
- class torch.distributed.algorithms.Join(joinables, enable=True, throw_on_early_termination=False, **kwargs)[source]¶
此类定义了通用 join 上下文管理器,该管理器允许在进程加入后调用自定义钩子。
这些钩子应覆盖未加入进程的集体通信,以防止挂起和出错,并确保算法的正确性。有关钩子定义的详细信息,请参阅
JoinHook
。警告
上下文管理器要求每个参与的
Joinable
在其自己的每次迭代集体通信之前调用方法notify_join_context()
以确保正确性。警告
上下文管理器要求
JoinHook
对象中的所有process_group
属性都相同。如果有多个JoinHook
对象,则使用第一个对象的device
。进程组和设备信息用于检查未加入的进程,以及在启用throw_on_early_termination
时通知进程抛出异常,这两者都使用全简并。- 参数
示例
>>> 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
- class torch.distributed.algorithms.Joinable[source]¶
这定义了一个可加入类的抽象基类。
可加入类(继承自
Joinable
)应该实现join_hook()
,它返回一个JoinHook
实例,以及join_device()
和join_process_group()
,分别返回设备和进程组信息。