在 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

您可以注意到，在PyTorch 教程中，该函数cleanup()在每个进程结束时调用（即在 mp.spawn 内）

def cleanup():
    dist.destroy_process_group()