torch.distributed

参考 torch.distributed.init_process_group() - 云+社区 - 腾讯云

目录

后端

PyTorch自带的后端

使用哪个后端？

常见的环境变量

基础

初始化

torch.distributed.is_available()[source]

torch.distributed.init_process_group(backend, init_method=None, timeout=datetime.timedelta(0, 1800), world_size=-1, rank=-1, store=None, group_name='')[source]

class torch.distributed.Backend

torch.distributed.get_backend(group=)[source]

torch.distributed.get_rank(group=)[source]

torch.distributed.get_world_size(group=)[source]

torch.distributed.is_initialized()[source]

torch.distributed.is_mpi_available()[source]

torch.distributed.is_nccl_available()[source]

TCP初始化

共享文件系统初始化

环境变量初始化

组

torch.distributed.new_group(ranks=None, timeout=datetime.timedelta(0, 1800), backend=None)[source]

点对点通信

torch.distributed.send(tensor, dst, group=, tag=0)[source]

torch.distributed.recv(tensor, src=None, group=, tag=0)[source]

torch.distributed.isend(tensor, dst, group=, tag=0)[source]

torch.distributed.irecv(tensor, src, group=, tag=0)[source]

同步和异步收集操作

收集函数

torch.distributed.broadcast(tensor, src, group=, async_op=False)[source]

torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=, async_op=False)[source]

torch.distributed.reduce(tensor, dst, op=ReduceOp.SUM, group=, async_op=False)[source]

torch.distributed.all_gather(tensor_list, tensor, group=, async_op=False)[source]

torch.distributed.gather(tensor, gather_list=None, dst=0, group=, async_op=False)[source]

torch.distributed.scatter(tensor, scatter_list=None, src=0, group=, async_op=False)[source]

torch.distributed.reduce_scatter(output, input_list, op=ReduceOp.SUM, group=, async_op=False)[source]

torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=, async_op=False)[source]

torch.distributed.barrier(group=, async_op=False)[source]

class torch.distributed.ReduceOp

class torch.distributed.reduce_op[source]

多GPU收集函数

torch.distributed.broadcast_multigpu(tensor_list, src, group=, async_op=False, src_tensor=0)[source]

torch.distributed.all_reduce_multigpu(tensor_list, op=ReduceOp.SUM, group=, async_op=False)[source]

torch.distributed.reduce_multigpu(tensor_list, dst, op=ReduceOp.SUM, group=, async_op=False, dst_tensor=0)[source]

torch.distributed.all_gather_multigpu(output_tensor_lists, input_tensor_list, group=, async_op=False)[source]

torch.distributed.reduce_scatter_multigpu(output_tensor_list, input_tensor_lists, op=ReduceOp.SUM, group=, async_op=False)[source]

---

后端

torch.distributed支持三种后端，每个都有不同的能力，下表展示了那个函数对使用CPU/CUDA张量是可以得到的。MPI只有在构建PyTorch的实现支持CUDA时才支持它。

PyTorch自带的后端

Pytorch分布式现在仅仅支持Linux。默认地，用Gloo和NCLL后端用来构建，并且包含在Pytorch分布式中(NCCL仅仅使用CUDA时才构筑)。MPI是一个可选的后端，并且仅仅在如果你从Pytorch源中构建时才包含。(例如，在一个MPI已经安装的Host上构建Pytorch。)

使用哪个后端？

过去，我们经常问：“我应该使用哪个后端？”

经验法则

GPU的分布式训练使用NCCL后端

CPU分布式训练使用Gloo后端

GPU host使用无限制宽带互连：

• 使用NCCL，因为它是目前唯一支持无限制宽带和GPUDirect的后端。

GPU hosts使用以太网互连：

• 使用NCCL，因为它目前提供最佳分布式GPU训练性能，特别对于多进程单点或者多点分布式训练。如果你遇到任何NCCL的问题，使用Gloo作为应急计划选项。(注意对于GPU目前Gloo比NCCL运行的慢)

CPU hosts用无限制带宽连接：

• 如果无限制带宽支持IB上的IP，使用Gloo，否则，使用MPI作为替代，我们计划在即将更新的版本中加入Gloo的无限制带宽支持。
• CPU hosts with Ethernet interconnect

用以太网连接的GPU：

• Use Gloo, unless you have specific reasons to use MPI.
• 使用Gloo，除非你有特殊的原因需要使用MPI。

常见的环境变量

选择要使用的网络接口

默认的，NCCL和Gloo后端都会尝试尽力选择正确的网络接口来使用。如果自动检测到的接口不正确，你可以使用下面的环境变量来重写：

• NCCL_SOCKET_IFNAME, for example export NCCL_SOCKET_IFNAME=eth0
• GLOO_SOCKET_IFNAME, for example export GLOO_SOCKET_IFNAME=eth0

如果你使用Gloo后端，你可以通过用逗号分隔它们来指定一个多借口而，比如export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3。后端将会在这些接口上以循环的方式派遣操作。在这些变量中所有进程指定相同的接口号很重要。

Other NCCL environment variables

其他NCCL环境变量。

NCCL提供了大量的环境变量来微调这个目的。

通常使用下面用来debugging目的中的一个：

• export NCCL_DEBUG=INFO
• export NCCL_DEBUG_SUBSYS=ALL

对NCCL环境变量的全部列表，请参考NVIDIA NCCL’s official documentation

基础

torch.distributed包提供Pytorch支持和通信基元，对多进程并行，在一个或多个机器上运行的若干个计算阶段。类torch.nn.parallel.DistributedDataParallel() 基于此功能来提供同步式分布式训练，作为任何Pytorch模块的包裹器。这和Multiprocessing package - torch.multiprocessing 和torch.nn.DataParallel()提供的并行式不同，它支持多网络连接的机器，对每个进程使用者必须明确的启动一个主训练脚本的一个分开的副本。

在单个机器同步的情况下，torch.distributed或者torch.nn.parallel.DistributedDataParallel()包裹器，相对于其他数据并行的方法包括torch.nn.DataParallel()依然有很大的优势：

• 在每次迭代时每个进程保持自己的优化器并且执行一个完全的优化步骤。同时这也许会出现冗余，因为梯度已经聚合在一起，并且平均到进程上，因此对每个进程都是一样的，这意味着不需要无参数广播步骤，减少了张量和节点之间的时间花费。
• 每个进程进程包含一个独立的Pytorch解释器，消除额外的解释器超出。来自执行若干执行线程的GIL-thrashing，模型复制或来自单个Python进程的GPUs。这对模型特别重要使得Python的运行时间能够大量使用，包括使用循环成的模型或者许多小的组件。

初始化

在调用任何方法之前，这个包需要使用torch.distributed.init_process_group()来进行初始化。这将阻塞，直到所有进程都已加入。

torch.distributed.is_available()[source]

如果分布式包可以获得的话就返回True。否则，torch.distributed对任何API都是不可见的。目前，torch.distributed在Linux和MacOS上都可以得到。设置USE_DISTRIBUTED=1来启动它，当从源中构建PyTorch时。目前，对Linux系统，默认值是USE_DISTRIBUTED=1，对MacOS系统，默认值是USE_DISTRIBUTED=0。

torch.distributed.init_process_group(backend, init_method=None, timeout=datetime.timedelta(0, 1800), world_size=-1, rank=-1, store=None, group_name='')[source]

初始化默认的分布式进程组，并且也会初始化分布式包。

There are 2 main ways to initialize a process group:

有两种主要的方式来进行初始化过程：

1. 明确的指定store，rank和world_size
2. 指定init_method (一个URL字符串)，这指明在哪里和如何发现同类。选择性的指定rank和world_size，或者编码所有在URL中需要的参数，并且忽略它们。

如果两者都没有指定，init_method假定为"env://"。

参数：

• backend(str或Backend)：--要使用的backend。依赖于构筑时间配置，mpi, gloo, nccl作为有效值。这个参数以小写字符串表示(例如，“gloo”)，这也能通过Backend属性(例如，Backend.GLOO)来访问。如果利用nccl在每个机器上使用多进程，每个进程必须独占访问它使用的每个GPU，因为在进程间共享GPU将会导致停滞。
• init_method(str，可选)--URL指定如何初始化进程组。默认是“env://”，如果init_method和store都没有指定的时候。和store是互斥的。
• world_size (python:int, optional) – Number of processes participating in the job. Required if store is specified.
• world_size(int，可选的)--在工作中参与的进程数，如果store指定时需要设置。
• rank (python:int, optional) – Rank of the current process. Required if store is specified.
• rank(int, 可选的)--当前进程的等级。如果store指定时就需要设置。
• store(store，可选的)--所有工作可访问的键/值存储，用于交换连接/地址信息。与init_method互斥。
• timeout(timedelta,可选的)--对流程组执行的操作的超时。默认值为30分钟。这适用于gloo后端。对于nccl，只有在环境变量NCCL_BLOCKING_WAIT被设置为1时才适用。

• group_name(str，可选的，丢弃)--组名字。

为了使得backend == Backend.MPI，PyTorch需要在支持MPI的系统上从源代码构建。这对NCCL同样适用。

class torch.distributed.Backend

一个可获得后端的类似枚举的类：GLOO，NCLL和MPI。

这个类的值是小写字母，例如"gloo"。它们可以作为属性来访问，例如Backend,NCCL。

这个类能直接调用来解析字符串，例如Backend(backend_str)，如果back_ste是空时将会检查。如果是这样的话，返回解析后的小写字母字符串。它也接受大写字符串，例如Backend("GLOO")，返回"gloo"。

注意

入口端。UNDEFINED存在，但仅用作某些字段的初始值。用户不应该直接使用它，也不应该假设它的存在。

torch.distributed.get_backend(group=<object object>)[source]

返回给定进程组的后端。

参数：

• group (ProcessGroup, optional) – 要处理的过程组。默认是通用的主进程组。如果其他指定的组指定后，调用的过程必须是group的 一部分。

返回值：

• 给定进程组的后端，作为一个小写字母的字符串。

torch.distributed.get_rank(group=<object object>)[source]

返回当前进程组的等级。在一个分布式进程组内，rank是一个独特的识别器分配到每个进程。它们通常是从0到world_size的连续的整数。

参数：

• group (ProcessGroup, optional) – 要处理的过程组。

返回值：

• 如果不是组的一部分，等级就是进程组-1。

torch.distributed.get_world_size(group=<object object>)[source]

返回当前进程组的进程数。

参数：

• group (ProcessGroup, optional) – 要处理的过程组。

返回值：

• 如果不是组的一部分，进程组的world_size就为-1。

torch.distributed.is_initialized()[source]

• 检查默认的进程数是否已经被初始化

torch.distributed.is_mpi_available()[source]

• 检查MPI后端是否可以得到

torch.distributed.is_nccl_available()[source]

• 检查NCCL后端是否可以得到。

目前支持三种初始化方式：

TCP初始化

使用TCP时有两种初始化方法，两者都需要一个所有进程都可以访问的网络地址和一个想要的world_size。第一种方法需要指定一个属于rank 0集成的地址。这个初始化方法需要对所有进程手动指定等级。注意多路广播地址在最近的分布式包中已经不再支持。group_name也已经丢弃了。

   
   

     
     

      
      

     
     
     
     

      
      

       
       import torch.distributed 
       
       as dist
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       # Use address of one of the machines
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       dist.init_process_group(backend, init_method=
       
       'tcp://10.1.1.20:23456',
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               rank=args.rank, world_size=
       
       4)
      
      
     
     
   
   

共享文件系统初始化

其他初始化方法使用一个文件系统，这个文件系统是共享的，并且在一个组内所有的机器都是能看见的，伴随着需要的world_size。URL应该以file://开始，并且包含共享文件系统中不存在文件的路径(在一个存在的目录中)。如果它不存在的话，文件系统初始化会自动创建一个文件，但是将会杉树文件。因此，你必须确定这个文件是否已在下一次init_process_group()调用相同文件路径名之前已经删除。注意多路广播地址在最近的分布式包中已经不再支持。group_name也已经丢弃了。

警告

这个方法假定文件系统支持用fant1锁定 - 大多数局部系统和NFS支持它。

警告

这个方法通常创建一个文件并且尽力去删除还在程序结束时移除这个文件。换句话说，用文件初始化方法初始的每个初始化将会需要一个需要一个全新的空文件，以便初始化成功。如果通过以前初始化后的文件再一次使用(没有清楚干净时将会发生)，这是不希望发生的行为，并且进程会导致僵局和失败。因此，如果自动删除没有发生的话，尽管这个方法将会尽力清除这个文件，你必须保证文件在训练结束时被移除，来阻止相同的文件在下一次中再次使用。如果你计划在同样的文件系统中调用init_process_group()时，这非常重要。换句话说，如果文件没有移除和清楚，并且你再一次在那个文件上调用init_process_group()，将会发生失败。这里的经验法则是，确保在调用文件 init_process_group()不存在或者是空的。

   
   

     
     

      
      

     
     
     
     

      
      

       
       import torch.distributed 
       
       as dist
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       # rank should always be specified
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       dist.init_process_group(backend, init_method=
       
       'file:///mnt/nfs/sharedfile',
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               world_size=
       
       4, rank=args.rank)
      
      
     
     
   
   

环境变量初始化

这个方法将会从环境变量中读取配置，允许人们完全定做信息如何获得。变量设置为：

• MASTER_PORT - 需要，在rank 0的机器上必须是一个空闲接口

• MASTER_ADDR - 需要(除了rank 0)；rank 0节点的地址

• WORLD_SIZE - 需要；可以在这里设置或者在初始化函数中调

• RANK - 需要；可以在这里设置或者在初始化函数中调用

级别为0的机器将用于设置所有连接。这是默认方法，意味着不需要指定init_method(也可以是env://)。

组

在默认组上收集默认操作(也称为world)并且需要要求所有进程进入分布式函数调用。然而，一些工作量也从更多细粒度通信中受益。这就是分布式组发挥作用的地方。new_group()函数可以用来闯进新的组，利用所有进程的任意子集。它返回一个不透明的组句柄，可以将其作为组参数提供给所有集合(集合是用于在某些众所周知的编程模式中交换信息的分布式函数)。

torch.distributed.new_group(ranks=None, timeout=datetime.timedelta(0, 1800), backend=None)[source]

创建一个分布式组。这个函数需要所有在主要组的进程(例如，所有为分布式工作部分的进程)进入这个函数。额外的，在所有进程上组应该以相同的顺序创建。

参数：

• ranks(list[python:int]) - 组员的进程列表。

• timeout(timedelta,optinal) - 对流程组执行的操作的超时。默认值为30分钟。这只适用于gloo后端。

• backend (str or Backend, optional) –

• The backend to use. Depending on build-time configurations, valid values are gloo and nccl. By default uses the same backend as the global group. This field should be given as a lowercase string (e.g., "gloo"), which can also be accessed via Backend attributes (e.g., Backend.GLOO).

• 使用的backend。依赖于构建时间配置，有效值是gloo和nccl。默认使用相同的后端作为全局组。这个域应该是小写字母的字符串(例如，“gloo”)，这也可以通过backend属性来进行访问。

返回值：

• 分布式组的句柄，可用于集体调用。

点对点通信

torch.distributed.send(tensor, dst, group=<object object>, tag=0)[source]

以同步方式传送一个张量。

参数：

• tensor (Tensor) – Tensor to send.

• tensor (Tensor) – 要传送的张量。

• dst (python:int) – Destination rank.

• dst (python:int) – 要到达的等级。

• group (ProcessGroup, optional) – 要处理的过程组

• tag (python:int, optional) – 标签匹配发送与远程recv

torch.distributed.recv(tensor, src=None, group=<object object>, tag=0)[source]

同步接收一个张量。

参数：

• tensor (Tensor) – Tensor to send.

• tensor (Tensor) – 要传送的张量。

• dst (python:int) – Destination rank.

• dst (python:int) – 要到达的等级。

• group (ProcessGroup, optional) – 要处理的过程组

• tag (python:int, optional) – 标签匹配发送与远程recv

返回值：

• 如果发件人不属于组，则其级别为-1。

isend() and irecv() return distributed request objects when used. In general, the type of this object is unspecified as they should never be created manually, but they are guaranteed to support two methods:

• is_completed() - returns True if the operation has finished

• wait() - will block the process until the operation is finished. is_completed() is guaranteed to return True once it returns.

torch.distributed.isend(tensor, dst, group=<object object>, tag=0)[source]

同步发送一个张量。

参数：

• tensor (Tensor) – Tensor to send.

• tensor (Tensor) – 要传送的张量。

• dst (python:int) – Destination rank.

• dst (python:int) – 要到达的等级。

• group (ProcessGroup, optional) – 要处理的过程组

• tag (python:int, optional) – 标签匹配发送与远程recv

返回值：

• 分布式请求对象。没有，如果不是组的一部分

torch.distributed.irecv(tensor, src, group=<object object>, tag=0)[source]

同步接收一个张量。

参数：

• tensor (Tensor) – Tensor to send.

• tensor (Tensor) – 要传送的张量。

• dst (python:int) – Destination rank.

• dst (python:int) – 要到达的等级。

• group (ProcessGroup, optional) – 要处理的过程组

• tag (python:int, optional) – 标签匹配发送与远程recv

返回值：

• 分布式请求对象。没有，如果不是组的一部分

同步和异步收集操作

Every collective operation function supports the following two kinds of operations:

每个收集操作都支持下面的两种收集方式：

同步操作 - 默认模式，当async_op被设置为False时。当函数返回时，它保证收集操作已经执行(如果他是一个CUDA操作的话，没有必要必须完成，因为所有的CUDA操作都是异步的)，根据集合操作的数据，可以调用任何进一步的函数调用。在同步模式下，集合函数不返回任何内容。

异步操作 - 当async_op设置为True时。收集操作函数返回一个分布式请求对象。通常情况下，你不需要手动创建它，而它保证支持两种方法：

• is_completed() - 如果操作完成的话，返回真

• wait() - 将会阻塞这个进程，知道操作完成。

收集函数

torch.distributed.broadcast(tensor, src, group=<object object>, async_op=False)[source]

将张量广播到整个组。在收集中所有参与的进程，张量必须有相同的元素数。

参数：

• tensor (Tensor) – Data to be sent if src is the rank of current process, and tensor to be used to save received data otherwise.

• tensor (Tensor) – 如果src是当前进程的号，就发送数据，并且张量保存接受数据。

• src (python:int) – 源等级。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

Returns：

返回值：

• 如果async_op是真的话，就返回异步工作句柄。如果async_op为空或者组的部分为空，就为空。

torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=<object object>, async_op=False)[source]

在所有机器上减少张量数据，通过获得最终的结果。在调用之后张量在所有过程中都是按位相同的。

参数：

• tensor (Tensor) – 输入和输出的集合。该函数在适当的位置运行。

• op (optional) – torch.distributed.ReduceOp枚举的价值之一。指定一个操作用来进行逐元素降低。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op时候应该是一个异步操作

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.reduce(tensor, dst, op=ReduceOp.SUM, group=<object object>, async_op=False)[source]

在所有机器上减少张量数据。仅仅具有dst的等级接收最终的数据。

参数：

• tensor (Tensor) – 输入和输出的集合。该函数在适当的位置运行。

• op (optional) – torch.distributed.ReduceOp枚举的价值之一。指定一个操作用来进行逐元素降低。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op时候应该是一个异步操作

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.all_gather(tensor_list, tensor, group=<object object>, async_op=False)[source]

Gathers tensors from the whole group in a list.

用类表的方式再真个组中收集张量。

参数：

• tensor_list (list[Tensor]) – Output list. It should contain correctly-sized tensors to be used for output of the collective.

• tensor_list (list[Tensor]) – 输出的列表。它应该包含当前尺寸的张量用来进行收集的输出。

• tensor (Tensor) – 从当前进程广播的张量。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op时候是async_op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.gather(tensor, gather_list=None, dst=0, group=<object object>, async_op=False)[source]

以单一的进程收集张量的列表。

参数：

• tensor (Tensor) – 输入张量。

• gather_list (list[Tensor], optional) – 用于收集数据的适当大小的张量列表(默认为None，必须在目标级别上指定)

• dst (python:int, optional) – 目的等级，默认是0

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.scatter(tensor, scatter_list=None, src=0, group=<object object>, async_op=False)[source]

将张量列表分散到一组中的所有进程。每个进程将只接收到一个张量，并将其数据存储在张量参数中。

参数：

• tensor (Tensor) – 输入张量。

• gather_list (list[Tensor], optional) – 用于收集数据的适当大小的张量列表(默认为None，必须在目标级别上指定)

• dst (python:int, optional) – 目的等级，默认是0

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.reduce_scatter(output, input_list, op=ReduceOp.SUM, group=<object object>, async_op=False)[source]

缩减，然后分散一列张量到一组中的所有进程。

Parameters

• output (Tensor) – 输出张量

• input_list (list[Tensor]) –要减少和打散的张量列表

• group (ProcessGroup, optional) – 要处理的过程租

• async_op (bool, optional) – 这个op是否应该是一个异步op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=<object object>, async_op=False)[source]

每个进程将输入张量列表分散到一组中所有进程，并在输出列表中返回收集到的张量列表。

参数：

• output_tensor_list (list[Tensor]) – 每个rank收集一个张量的列表

• input_tensor_list (list[Tensor]) – 每个rank打散一个张量列表

• group (ProcessGroup, optional) – 要处理的过程租

• async_op (bool, optional) – 这个op是否应该是一个异步op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

Warning

all_to_all是实验性的，可能会发生变化。

例

   
   

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       input = torch.arange(
       
       4) + rank * 
       
       4
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       input = 
       
       list(
       
       input.chunk(
       
       4))
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       input
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       0]), tensor([
       
       1]), tensor([
       
       2]), tensor([
       
       3])]     
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       4]), tensor([
       
       5]), tensor([
       
       6]), tensor([
       
       7])]     
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       8]), tensor([
       
       9]), tensor([
       
       10]), tensor([
       
       11])]   
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       12]), tensor([
       
       13]), tensor([
       
       14]), tensor([
       
       15])] 
       
       # Rank 3
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> output = 
       
       list(torch.empty([
       
       4], dtype=torch.int64).chunk(
       
       4))
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> dist.all_to_all(output, 
       
       input)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> output
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       0]), tensor([
       
       4]), tensor([
       
       8]), tensor([
       
       12])]    
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       1]), tensor([
       
       5]), tensor([
       
       9]), tensor([
       
       13])]    
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       2]), tensor([
       
       6]), tensor([
       
       10]), tensor([
       
       14])]   
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       3]), tensor([
       
       7]), tensor([
       
       11]), tensor([
       
       15])]   
       
       # Rank 3
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       # Essentially, it is similar to following operation:
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> scatter_list = 
       
       input
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> gather_list  = output
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       for i 
       
       in 
       
       range(world_size):
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>>   dist.scatter(gather_list[i], scatter_list 
       
       if i == rank 
       
       else [], src = i)
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       input
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       tensor([
       
       0, 
       
       1, 
       
       2, 
       
       3, 
       
       4, 
       
       5])                                       
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       tensor([
       
       10, 
       
       11, 
       
       12, 
       
       13, 
       
       14, 
       
       15, 
       
       16, 
       
       17, 
       
       18])                     
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       tensor([
       
       20, 
       
       21, 
       
       22, 
       
       23, 
       
       24])                                     
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       tensor([
       
       30, 
       
       31, 
       
       32, 
       
       33, 
       
       34, 
       
       35, 
       
       36])                             
       
       # Rank 3
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> input_splits
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       2, 
       
       2, 
       
       1, 
       
       1]                                                     
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       3, 
       
       2, 
       
       2, 
       
       2]                                                     
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       2, 
       
       1, 
       
       1, 
       
       1]                                                     
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       2, 
       
       2, 
       
       2, 
       
       1]                                                     
       
       # Rank 3
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> output_splits
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       2, 
       
       3, 
       
       2, 
       
       2]                                                     
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       2, 
       
       2, 
       
       1, 
       
       2]                                                     
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       1, 
       
       2, 
       
       1, 
       
       2]                                                     
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [
       
       1, 
       
       2, 
       
       1, 
       
       1]                                                     
       
       # Rank 3
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       input = 
       
       list(
       
       input.split(input_splits))
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> 
       
       input
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       0, 
       
       1]), tensor([
       
       2, 
       
       3]), tensor([
       
       4]), tensor([
       
       5])]                   
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       10, 
       
       11, 
       
       12]), tensor([
       
       13, 
       
       14]), tensor([
       
       15, 
       
       16]), tensor([
       
       17, 
       
       18])] 
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       20, 
       
       21]), tensor([
       
       22]), tensor([
       
       23]), tensor([
       
       24])]                 
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       30, 
       
       31]), tensor([
       
       32, 
       
       33]), tensor([
       
       34, 
       
       35]), tensor([
       
       36])]         
       
       # Rank 3
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> output = ...
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> dist.all_to_all(output, 
       
       input)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       >>> output
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       0, 
       
       1]), tensor([
       
       10, 
       
       11, 
       
       12]), tensor([
       
       20, 
       
       21]), tensor([
       
       30, 
       
       31])]   
       
       # Rank 0
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       2, 
       
       3]), tensor([
       
       13, 
       
       14]), tensor([
       
       22]), tensor([
       
       32, 
       
       33])]           
       
       # Rank 1
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       4]), tensor([
       
       15, 
       
       16]), tensor([
       
       23]), tensor([
       
       34, 
       
       35])]              
       
       # Rank 2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       [tensor([
       
       5]), tensor([
       
       17, 
       
       18]), tensor([
       
       24]), tensor([
       
       36])]                  
       
       # Rank 3
      
      
     
     
   
   

    
    
   
   
   
   

torch.distributed.barrier(group=<object object>, async_op=False)[source]

对所有进程进行同步。这个收集阻塞知道整个组进入这个函数，如果async_op为Fasle或者，如果在wait()上调用异步工作句柄。

参数：

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

class torch.distributed.ReduceOp

一个枚举类型的类用来获得可以减少的操作：SUM, PRODUCT, MIN, MAX, BAND, BOR, and BXOR。这个类的值能通过属性访问，例如ReduceOp.SUM.，它们被用来指定减少集体的战略，例如 reduce(), all_reduce_multigpu()。

成员：

• SUM
• PRODUCT
• MIN
• MAX
• BAND
• BOR
• BXOR

class torch.distributed.reduce_op[source]

Deprecated enum-like class for reduction operations: SUM, PRODUCT, MIN, and MAX.

ReduceOp is recommended to use instead.

多GPU收集函数

如果在每个节点上超过一个GPU，当使用NCCL或者Gloo后端时，broadcast_multigpu() all_reduce_multigpu() reduce_multigpu() 和all_gather_multigpu()在每个节点内所有多GPU都支持分布式收集操作。在过去张量列表中的每个张量需要在hosts的分开的GPU设备上，在这里调用函数。注意张量列表的长度需要在所有分布式进程上都相同。同时我们注意到目前的多GPU收集函数仅仅支持NCCL后端。例如，如果用来进行分布式训练的系统有两个节点，每个有8个GPU。在这16个GPU的每个上面，这个张量我们都需要减少。接下来的代码可以作为一个参考：

在节点0上运行代码

   
   

     
     

      
      

     
     
     
     

      
      

       
       import torch
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       import torch.distributed 
       
       as dist
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       dist.init_process_group(backend=
       
       "nccl",
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               init_method=
       
       "file:///distributed_test",
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               world_size=
       
       2,
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               rank=
       
       0)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       tensor_list = []
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       for dev_idx 
       
       in 
       
       range(torch.cuda.device_count()):
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           tensor_list.append(torch.FloatTensor([
       
       1]).cuda(dev_idx))
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       dist.all_reduce_multigpu(tensor_list)
      
      
     
     
   
   

在节点1上运行代码

   
   

     
     

      
      

     
     
     
     

      
      

       
       import torch
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       import torch.distributed 
       
       as dist
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       dist.init_process_group(backend=
       
       "nccl",
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               init_method=
       
       "file:///distributed_test",
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               world_size=
       
       2,
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                               rank=
       
       1)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       tensor_list = []
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       for dev_idx 
       
       in 
       
       range(torch.cuda.device_count()):
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           tensor_list.append(torch.FloatTensor([
       
       1]).cuda(dev_idx))
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       dist.all_reduce_multigpu(tensor_list)
      
      
     
     
   
   

调用完成后，两个节点上的所有16张量的全部减少值为16。

torch.distributed.broadcast_multigpu(tensor_list, src, group=<object object>, async_op=False, src_tensor=0)[source]

将张量广播给每节点有多个GPU张量的整个组。张量在所有gpu中必须有相同数量的元素，这些元素来自于参与集合的所有进程。列表中的每个张量必须在不同的GPUOnly nccl和gloo后端是当前支持的张量，只能是GPU张量。

参数：

• tensor_list (List[Tensor]) – 在收集操作中参与的张量。如果src是rank，指定tensor_list(tensor_list[src_tensor])的src_tensor元素将会广播到其他张量(在不同的GPU上)，在src进程和所有其他non-src进程的tensor_list。您还需要确保len(tensor_list)对于调用该函数的所有分布式进程是相同的。

• src (python:int) – 源层级。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

• src_tensor (python:int, optional) – 源张量在tensor_list中的秩

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.all_reduce_multigpu(tensor_list, op=ReduceOp.SUM, group=<object object>, async_op=False)[source]

在所有机器上减少张量数据，以得到最终结果的方式。这个函数在每个节点处减少大量的张量，同时每个张量存在于不同的GPUs上。因此，张量列表中的输出张量需要是GPU张量。同时，张量列表中的灭个张量需要在不同的GPU上存在。调用以后，张量列表中的所有张量将会逐比特相同。仅仅nccl和gloo后端目前支持张量应该成为GPU张量。

参数：

• 收集到的输入和输出张量的列表。函数在原地操作，并且需要在不同的GPU上每个张量是一个GPU张量。你也需要确保len(tensor_list)对所有调用这个函数的分布式训练是相同的。

• op (optional) – torch.distributed.ReduceOp枚举来的张量之一。指定一个操作用来进行逐元素的减少。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.reduce_multigpu(tensor_list, dst, op=ReduceOp.SUM, group=<object object>, async_op=False, dst_tensor=0)[source]

在所有机器上减少多GPU的张量数据。每个在tensor_list中的张量应该在单独的GPU中存在。只有dst为秩的进程的tensor_list[dst_张量]的GPU会接收最终结果。目前只支持nccl后端张量，只能是GPU张量。

参数：

• tensor_list (List[Tensor]) – Input and output GPU tensors of the collective. The function operates in-place. You also need to make sure that len(tensor_list) is the same for all the distributed processes calling this function.

• tensor_list (List[Tensor]) – 收集的输入和输出GPU张量。函数操作在原地完成。你也需要确保len(tensor_list)对所有调用这个函数的分布式进程都是相同的。

• dst (python:int) – Destination rank

• dst (python:int) – 目的等级

• op (optional) – One of the values from torch.distributed.ReduceOp enum. Specifies an operation used for element-wise reductions.

• op (optional) – 从torch.distributed.ReduceOp枚举来的值之一。指定一个用来逐元素减少的操作。

• group (ProcessGroup, optional) – 要处理的过程组

• async_op (bool, optional) – 这个op是否应该是一个异步op

• dst_tensor (python:int, optional) – 目的张量在张量表中的rank

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.all_gather_multigpu(output_tensor_lists, input_tensor_list, group=<object object>, async_op=False)[source]

从列表中的整个组中收集张量。每个张量在tensor_list应该驻留在一个单独的GPUOnly nccl后端，当前支持的张量应该是GPU张量。

参数：

• output_tensor_lists (List[List[Tensor]]) –  输出列表，在每个GPU上应该包含当前尺寸的张量，用来进行收集的输出，例如output_tensor_lists[i]包含了input_tensor_list[i]的GPU上的all_gather结果。注意output_tensor_lists 的每个元素的尺寸都为world_size * len(input_tensor_list)，因为该函数在组中都收集来自每个单GPU的结果。为了解释output_tensor_lists[i]的灭个元素，注意rank k的input_tensor_list[j]将会出现在output_tensor_lists[i][k * world_size + j]。也注意到len(output_tensor_lists)，output_tensor_lists 中的每个元素的尺寸(每个元素是一个列表，因此len(output_tensor_lists[i]))需要在调用这个函数的所有进程上相等)

• input_tensor_list (List[Tensor]) – 从当前进程广播的张量列表(在不同的gpu上)。注意，调用该函数的所有分布式进程的len(input_tensor_list)需要相同。

• async_op (bool, optional) – 这个op是否应该是一个异步op

• dst_tensor (python:int, optional) – 目的张量在张量表中的rank

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。

torch.distributed.reduce_scatter_multigpu(output_tensor_list, input_tensor_lists, op=ReduceOp.SUM, group=<object object>, async_op=False)[source]

将一个张量列表压缩并分散到整个组中。目前只支持nccl后端。output_tensor_list中的每个张量都应该驻留在一个单独的GPU上，就像input_tensor_lists中的每个张量列表一样。

Parameters

• output_tensor_list (List[Tensor]) – 输出张量(在不同的gpu上)来接收操作的结果。注意，调用该函数的所有分布式进程的len(output_tensor_list)需要相同。

• input_tensor_lists (List[List[Tensor]]) – 输入列表。它应该在每个GPU上包含大小正确的张量来用于集合的输入，例如input_tensor_lists[i]包含了位于output_tensor_list[i] GPU上的reduce_scatter输入。注意，input_tensor_lists的每个元素的大小为world_size * len(output_tensor_list)，因为该函数分散了组中每个单独GPU的结果。要解释input_tensor_lists[i]的每个元素，请注意，排名k的output_tensor_list[j]接收到来自input_tensor_lists[i][k * world_size + j]的reduce-scattered结果。还要注意，len(input_tensor_lists)和input_tensor_lists中每个元素的大小(每个元素都是一个列表，因此len(input_tensor_lists[i])对于调用该函数的所有分布式进程都需要相同。

• g

• async_op (bool, optional) – 这个op是否应该是一个异步op

• dst_tensor (python:int, optional) – 目的张量在张量表中的rank

返回值：

• 如果async_op设置为真的话，是一个异步工作句柄。如果async_op为空或者部分组为空，就为None 。