torch.distributed.barrier() 是如何工作的

How does torch.distributed.barrier() work

我已经阅读了我能找到的关于 torch.distributed.barrier() 的所有文档,但仍然无法理解它在 this script 中的使用方式,非常感谢您的帮助。

所以 official doc of torch.distributed.barrier 说 "Synchronizes all processes.This collective blocks processes until the whole group enters this function, if async_op is False, or if async work handle is called on wait()."

它在脚本中的两个地方使用:

First place

    if args.local_rank not in [-1, 0] and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

        ... (preprocesses the data and save the preprocessed data)

    if args.local_rank == 0 and not evaluate:
        torch.distributed.barrier() 

Second place

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

        ... (loads the model and the vocabulary)

    if args.local_rank == 0:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

我无法将代码中的注释与官方文档中所述的此函数的功能相关联。它如何确保只有第一个进程在 torch.distributed.barrier() 的两次调用之间执行代码,以及为什么它只在第二次调用之前检查本地等级是否为 0?

提前致谢!

首先你需要了解行列。简而言之:在多处理上下文中,我们通常假设等级 0 是第一个进程或基础进程。然后对其他过程进行不同的排序,例如1、2、3,共四个进程。

有些操作不需要并行完成,或者您只需要一个进程进行一些预处理或缓存,以便其他进程可以使用该数据。

在您的示例中,如果第一个 if 语句由非基础进程(等级 1、2、3)输入,它们将阻塞(或“等待”),因为它们 运行 进入屏障.他们在那里等待,因为 barrier() 阻塞,直到 所有 进程达到障碍,但基本进程尚未达到障碍。

所以此时非基本进程(1、2、3)被阻止,但基本进程(0)继续。基本进程将执行一些操作(在本例中为预处理和缓存数据),直到它到达第二个 if 语句。在那里,基础进程将 运行 变成障碍。此时,所有进程都停在了一个障碍上,这意味着所有当前的障碍都可以解除,所有进程都可以继续。因为基础进程准备了数据,所以其他进程现在可以使用该数据。

也许最重要的是要理解:

  • 当进程遇到障碍时它会阻塞
  • 屏障的位置并不重要(例如,并非所有进程都必须输入相同的 if 语句)
  • 一个进程被屏障阻塞,直到所有进程都遇到屏障,所有进程的屏障都会被解除