关于torch.distributed.barrier()的作用

Reference

https://stackoverflow.com/questions/59760328/how-does-torch-distributed-barrier-work

http://t.csdn.cn/v6I0z

在Yolov5中我碰到了一个函数
torch_distributed_zero_first(LOCAL_RANK)

from contextlib import contextmanager

@contextmanager
def torch_distributed_zero_first(local_rank: int):
    # Decorator to make all processes in distributed training wait for each local_master to do something 确保分布式训练中主进程完毕
    if local_rank not in [-1, 0]:
        torch.distributed.barrier(device_ids=[local_rank])
    yield
    if local_rank == 0:
         torch.distributed.barrier(device_ids=[0])

1. 参数local_rank的理解

首先，如果我们使用了多进程处理任务时，Python通常假设等级 0 是第一个进程或基本进程。然后对其他进程进行不同的排序，例如1、2、3，总共四个进程。
如果我们使用单进程时，往往会设置进程号为-1。
总而言之，当local_rank=-1或者0时，我们认为它是主进程。
而对于多进程任务，我们往往只需要一个进程预处理数据或者读取数据，而为了与其他进程共享数据，我们需要在主进程处理数据的时候让其他进程全部停下来等待。

理解

基于这样的任务要求，也为了使用pytorch的多机多卡模式，我们需要通过torch.distributed.barrier()实现数据同步。

   if local_rank not in [-1, 0]: # 询问当前进程是否是主进程
   		# 不是就让pytorch对它进行阻塞，也就是暂停运行
        torch.distributed.barrier(device_ids=[local_rank]) 

yield # 如果他是主进程，他就暂时跳出函数执行其他任务

if local_rank == 0:
	# 此时主进程完成了其他任务进入了第二个if函数
	# 然后它也进入barrier（）函数里面
    torch.distributed.barrier(device_ids=[0]) 

当pytorch发现所有的进程都进入了barrier（），就会打开所有的barrier，所有的进程都可以继续进行。

总结

• 当进程遇到障碍时，它将阻塞
• 屏障的位置并不重要（例如，并非所有进程都必须输入相同的 if 语句）
• 一个进程被一个屏障阻塞，直到所有进程都遇到一个屏障，在这个屏障上为所有进程解除这些屏障