在 Pytorch 中应该在哪里调用 torch.distributed.destroy_process_group()？

Question

我注意到文档没有那个功能。因此，不清楚应该在哪里调用它。是否必须：

1. 在每个工人代码的末尾调用它（即在 mp.spawn 内部）
2. 或者在 mp.spawn 之外调用它，即通过主进程

注意有一个 gitissue 要求将此功能放在文档中：https://github.com/pytorch/pytorch/issues/48203

这是 2 含义的示例：

def test_setup():
    print('test_setup')
    if torch.cuda.is_available():
        world_size = torch.cuda.device_count()
    else:
        world_size = 4
    master_port = find_free_port()
    mp.spawn(setup_process, args=(world_size, master_port), nprocs=world_size)
    dist.destroy_process_group()
    print('successful test_setup!')

Answer 1

dist.destroy_process_group()

函数应该在

mp.spawn()

函数之外调用，在所有派生的进程完成它们的任务并退出之后。此函数应由主进程调用，主进程使用

mp.spawn()

生成子进程。该功能的目的是在进程完成任务退出后清理分布式环境，防止资源泄漏，并确保程序的后续运行从零开始。

在您提供的示例中，在所有进程完成任务并退出后，

dist.destroy_process_group()

函数在

mp.spawn()

函数之外被调用。这是调用此函数的正确方法。

值得注意的是，这个函数并不是 PyTorch 特有的，而是

torch.distributed

模块的一个函数，它为 PyTorch 中的分布式训练提供了一个后端不可知的接口。