评价此页
★ ★ ★ ★ ★
intermediate/dist_tuto

使用 PyTorch 编写分布式应用程序#

创建于：2017年10月06日 | 最近更新：2025年09月05日 | 最近验证：2024年11月05日

作者：Séb Arnold

注意

编辑 在 GitHub 中查看和编辑此教程。

先决条件

在这个简短的教程中，我们将介绍 PyTorch 的分布式软件包。我们将了解如何设置分布式环境、使用不同的通信策略，并深入探讨该软件包的一些内部原理。

设置#

PyTorch 中包含的分布式软件包（即 torch.distributed）使研究人员和从业人员能够轻松地在进程和机器集群之间并行化其计算。为此，它利用了消息传递语义，允许每个进程向任何其他进程通信数据。与多进程（torch.multiprocessing）软件包相反，进程可以使用不同的通信后端，并且不限制在同一台机器上执行。

为了开始使用，我们需要能够同时运行多个进程。如果您有权访问计算集群，则应向您的本地系统管理员咨询，或使用您最喜欢的协调工具（例如 pdsh、clustershell 或 slurm）。在本教程中，我们将使用单台机器并使用以下模板派生多个进程。

"""run.py:"""
#!/usr/bin/env python
import os
import sys
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def run(rank, size):
    """ Distributed function to be implemented later. """
    pass

def init_process(rank, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size)


if __name__ == "__main__":
    world_size = 2
    processes = []
    if "google.colab" in sys.modules:
        print("Running in Google Colab")
        mp.get_context("spawn")
    else:
        mp.set_start_method("spawn")
    for rank in range(world_size):
        p = mp.Process(target=init_process, args=(rank, world_size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

上面的脚本派生了两个进程，它们各自设置分布式环境、初始化进程组（dist.init_process_group），最后执行给定的 run 函数。

让我们来看看 init_process 函数。它确保每个进程都能够通过主节点进行协调，使用相同的 IP 地址和端口。请注意，我们使用的是 gloo 后端，但也可以使用其他后端（参见 第 5.1 节）。我们将在本教程的最后探讨 dist.init_process_group 中的奥秘，但它基本上允许进程通过共享彼此的位置来进行通信。

点对点通信#

Send and Recv

发送与接收 (Send and Recv)#

从一个进程到另一个进程的数据传输称为点对点通信。这些是通过 send 和 recv 函数或它们的即时（非阻塞）对应函数 isend 和 irecv 实现的。

"""Blocking point-to-point communication."""

def run(rank, size):
    tensor = torch.zeros(1)
    if rank == 0:
        tensor += 1
        # Send the tensor to process 1
        dist.send(tensor=tensor, dst=1)
    else:
        # Receive tensor from process 0
        dist.recv(tensor=tensor, src=0)
    print('Rank ', rank, ' has data ', tensor[0])

在上面的示例中，两个进程都从零张量开始，然后进程 0 递增该张量并将其发送给进程 1，从而使它们最终的值都为 1.0。请注意，进程 1 需要分配内存才能存储它将接收的数据。

还要注意，send/recv 是阻塞的：两个进程都会阻塞，直到通信完成。另一方面，即时操作是非阻塞的；脚本会继续执行，并且这些方法会返回一个 Work 对象，我们可以选择对其调用 wait()。

"""Non-blocking point-to-point communication."""

def run(rank, size):
    tensor = torch.zeros(1)
    req = None
    if rank == 0:
        tensor += 1
        # Send the tensor to process 1
        req = dist.isend(tensor=tensor, dst=1)
        print('Rank 0 started sending')
    else:
        # Receive tensor from process 0
        req = dist.irecv(tensor=tensor, src=0)
        print('Rank 1 started receiving')
    req.wait()
    print('Rank ', rank, ' has data ', tensor[0])

使用即时操作时，我们必须小心使用发送和接收的张量。因为我们不知道数据何时会通信到另一个进程，所以在 req.wait() 完成之前，我们不应该修改发送的张量，也不应该访问接收的张量。换句话说，

  • 在 dist.isend() 之后写入 tensor 将导致未定义行为。

  • 在执行 req.wait() 之前，在 dist.irecv() 之后读取 tensor 将导致未定义行为。

然而，在执行 req.wait() 之后，我们可以保证通信已经发生，并且存储在 tensor[0] 中的值为 1.0。

点对点通信在我们需要对进程的通信进行更细粒度的控制时非常有用。它们可以用来实现一些复杂的算法，例如 百度 DeepSpeech 或 Facebook 的大规模实验中使用的算法（参见 第 4.1 节）。

集合通信#

Scatter

Scatter (分发)#
Gather

Gather (收集)#
Reduce

Reduce (规约)#
All-Reduce

All-Reduce (全规约)#
Broadcast

Broadcast (广播)#
All-Gather

All-Gather (全收集)#

与点对点通信相反，集合通信允许在组 (group)中的所有进程之间进行通信。组是我们所有进程的一个子集。要创建一个组，我们可以向 dist.new_group(group) 传递一个进程秩（rank）列表。默认情况下，集合通信在所有进程上执行，这也称为 world。例如，为了获得所有进程上所有张量的总和，我们可以使用 dist.all_reduce(tensor, op, group) 集合通信。

""" All-Reduce example."""
def run(rank, size):
    """ Simple collective communication. """
    group = dist.new_group([0, 1])
    tensor = torch.ones(1)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group)
    print('Rank ', rank, ' has data ', tensor[0])

因为我们想要组中所有张量的总和，所以我们使用 dist.ReduceOp.SUM 作为规约操作符。一般来说，任何满足交换律的数学运算都可以用作操作符。PyTorch 开箱即用地提供了许多此类操作符，它们都工作在元素级（element-wise）：

  • dist.ReduceOp.SUM,

  • dist.ReduceOp.PRODUCT,

  • dist.ReduceOp.MAX,

  • dist.ReduceOp.MIN,

  • dist.ReduceOp.BAND,

  • dist.ReduceOp.BOR,

  • dist.ReduceOp.BXOR,

  • dist.ReduceOp.PREMUL_SUM.

支持的操作符的完整列表在此处。

除了 dist.all_reduce(tensor, op, group) 之外，PyTorch 当前还实现了许多其他集合通信操作。以下是一些支持的集合通信操作：

  • dist.broadcast(tensor, src, group)：将 tensor 从 src 复制到所有其他进程。

  • dist.reduce(tensor, dst, op, group)：将 op 应用于每个 tensor 并将结果存储在 dst 中。

  • dist.all_reduce(tensor, op, group)：与 reduce 相同，但结果会存储在所有进程中。

  • dist.scatter(tensor, scatter_list, src, group)：将第 \(i^{\text{th}}\) 个张量 scatter_list[i] 复制到第 \(i^{\text{th}}\) 个进程。

  • dist.gather(tensor, gather_list, dst, group)：将所有进程的 tensor 复制到 dst 中。

  • dist.all_gather(tensor_list, tensor, group)：在所有进程上，将所有进程的 tensor 复制到 tensor_list 中。

  • dist.barrier(group)：阻塞 group 中的所有进程，直到每个进程都进入了该函数。

  • dist.all_to_all(output_tensor_list, input_tensor_list, group)：将输入张量列表分发到组中的所有进程，并将收集到的张量列表返回到输出列表中。

可以通过查看 PyTorch Distributed 的最新文档来找到支持的集合通信操作的完整列表 （链接）。

分布式训练#

注意：您可以在这个 GitHub 仓库中找到本节的示例脚本。

现在我们了解了分布式模块的工作原理，让我们用它写一些有用的东西。我们的目标是复制 DistributedDataParallel 的功能。当然，这将是一个教学示例，在实际情况中，您应该使用上面链接的官方、经过良好测试且高度优化的版本。

简单来说，我们想要实现一个分布式版本的随机梯度下降（SGD）。我们的脚本将让所有进程在各自的数据批次（batch）上计算其模型的梯度，然后对它们的梯度进行平均。为了确保在更改进程数量时能获得相似的收敛结果，我们首先必须对数据集进行分区。（您也可以使用 torch.utils.data.random_split，而不是使用下面的代码片段。）

""" Dataset partitioning helper """
class Partition(object):

    def __init__(self, data, index):
        self.data = data
        self.index = index

    def __len__(self):
        return len(self.index)

    def __getitem__(self, index):
        data_idx = self.index[index]
        return self.data[data_idx]


class DataPartitioner(object):

    def __init__(self, data, sizes=[0.7, 0.2, 0.1], seed=1234):
        self.data = data
        self.partitions = []
        rng = Random()  # from random import Random
        rng.seed(seed)
        data_len = len(data)
        indexes = [x for x in range(0, data_len)]
        rng.shuffle(indexes)

        for frac in sizes:
            part_len = int(frac * data_len)
            self.partitions.append(indexes[0:part_len])
            indexes = indexes[part_len:]

    def use(self, partition):
        return Partition(self.data, self.partitions[partition])

利用上面的代码片段，我们现在只需使用以下几行代码即可轻松对任何数据集进行分区：

""" Partitioning MNIST """
def partition_dataset():
    dataset = datasets.MNIST('./data', train=True, download=True,
                             transform=transforms.Compose([
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.1307,), (0.3081,))
                             ]))
    size = dist.get_world_size()
    bsz = 128 // size
    partition_sizes = [1.0 / size for _ in range(size)]
    partition = DataPartitioner(dataset, partition_sizes)
    partition = partition.use(dist.get_rank())
    train_set = torch.utils.data.DataLoader(partition,
                                         batch_size=bsz,
                                         shuffle=True)
    return train_set, bsz

假设我们有 2 个副本，那么每个进程将拥有一个包含 60000 / 2 = 30000 个样本的 train_set。我们还将批量大小（batch size）除以副本数量，以保持 128 的整体批量大小。

现在我们可以编写通常的前向-反向-优化训练代码，并添加一个函数调用来平均模型的梯度。（以下内容在很大程度上启发自官方的 PyTorch MNIST 示例。）

""" Distributed Synchronous SGD Example """
def run(rank, size):
    torch.manual_seed(1234)
    train_set, bsz = partition_dataset()
    model = Net()
    optimizer = optim.SGD(model.parameters(),
                          lr=0.01, momentum=0.5)

    num_batches = ceil(len(train_set.dataset) / float(bsz))
    for epoch in range(10):
        epoch_loss = 0.0
        for data, target in train_set:
            optimizer.zero_grad()
            output = model(data)
            loss = F.nll_loss(output, target)
            epoch_loss += loss.item()
            loss.backward()
            average_gradients(model)
            optimizer.step()
        print('Rank ', dist.get_rank(), ', epoch ',
              epoch, ': ', epoch_loss / num_batches)

接下来需要实现 average_gradients(model) 函数，它只需接收一个模型并将其梯度在整个 world 中进行平均。

""" Gradient averaging. """
def average_gradients(model):
    size = float(dist.get_world_size())
    for param in model.parameters():
        dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
        param.grad.data /= size

瞧（Et voilà）！我们成功实现了分布式同步 SGD，并可以在大型计算机集群上训练任何模型。

注意：虽然最后一句话在技术上是正确的，但要实现生产级别的同步 SGD，还需要更多的技巧。同样，请使用已经过测试和优化的方法。

我们自己的 Ring-Allreduce#

作为一个额外的挑战，假设我们想要实现 DeepSpeech 的高效 ring allreduce。这通过使用点对点集合通信操作可以相当容易地实现。

""" Implementation of a ring-reduce with addition. """
def allreduce(send, recv):
   rank = dist.get_rank()
   size = dist.get_world_size()
   send_buff = send.clone()
   recv_buff = send.clone()
   accum = send.clone()

   left = ((rank - 1) + size) % size
   right = (rank + 1) % size

   for i in range(size - 1):
       if i % 2 == 0:
           # Send send_buff
           send_req = dist.isend(send_buff, right)
           dist.recv(recv_buff, left)
           accum[:] += recv_buff[:]
       else:
           # Send recv_buff
           send_req = dist.isend(recv_buff, right)
           dist.recv(send_buff, left)
           accum[:] += send_buff[:]
       send_req.wait()
   recv[:] = accum[:]

在上面的脚本中，allreduce(send, recv) 函数的签名与 PyTorch 中的签名略有不同。它接收一个 recv 张量，并将所有 send 张量的总和存储在其中。作为留给读者的练习，我们的版本与 DeepSpeech 中的版本之间仍然存在一个区别：他们的实现将梯度张量分割成块（chunks），以便最有效地利用通信带宽。（提示：torch.chunk）

高级主题#

我们现在准备发掘 torch.distributed 的一些更高级的功能。由于需要涵盖的内容很多，本节将分为两个子节：

  1. 通信后端：我们将学习如何使用 MPI 和 Gloo 进行 GPU 与 GPU 之间的通信。

  2. 初始化方法：我们将了解如何最好地设置 dist.init_process_group() 中的初始协调阶段。

通信后端#

torch.distributed 最优雅的方面之一是它能够对不同的后端进行抽象并在此之上构建。如前所述，PyTorch 中实现了多个后端。可以使用 Accelerator API 轻松选择这些后端，该 API 提供了与不同加速器类型配合使用的接口。一些最受欢迎的后端包括 Gloo、NCCL 和 MPI。根据所需的使用场景，它们各自具有不同的规格和权衡。支持功能的对比表可以在此处找到。

Gloo 后端

到目前为止，我们已经广泛使用了 Gloo 后端。作为一个开发平台，它非常方便，因为它包含在预编译的 PyTorch 二进制文件中，并且在 Linux（自 0.2 版本起）和 macOS（自 1.3 版本起）上均可工作。它支持 CPU 上的所有点对点和集合通信操作，以及 GPU 上的所有集合通信操作。CUDA 张量的集合通信操作实现没有 NCCL 后端提供的那么优化。

正如您一定注意到的那样，如果您将 model 放在 GPU 上，我们的分布式 SGD 示例将无法运行。为了使用多个 GPU，我们还需要进行以下修改：

  1. 使用 Accelerator API device_type = torch.accelerator.current_accelerator()

  2. 使用 torch.device(f"{device_type}:{rank}")

  3. model = Net() \(\rightarrow\) model = Net().to(device)

  4. 使用 data, target = data.to(device), target.to(device)

进行这些修改后，您的模型现在将在两个 GPU 上进行训练。如果您运行在 NVIDIA 硬件上，可以使用 watch nvidia-smi 来监控 GPU 利用率。

MPI 后端

消息传递接口（MPI）是高性能计算领域的标准化工具。它允许进行点对点和集合通信，并且是 torch.distributed API 的主要灵感来源。存在多种 MPI 实现（例如 Open-MPI、MVAPICH2、Intel MPI），每种实现都针对不同的目的进行了优化。使用 MPI 后端的优势在于 MPI 在大型计算机集群上的广泛可用性以及高水平的优化。一些最近的实现还能够利用 CUDA IPC 和 GPU Direct 技术，从而避免通过 CPU 进行内存复制。

不幸的是，PyTorch 的二进制文件无法包含 MPI 实现，我们必须手动重新编译。幸运的是，这个过程相当简单，因为在编译时，PyTorch 会自动寻找可用的 MPI 实现。以下步骤通过从源码安装 PyTorch 来安装 MPI 后端：

  1. 创建并激活您的 Anaconda环境，按照指南安装所有先决条件，但先不要运行 python setup.py install。

  2. 选择并安装您最喜欢的 MPI 实现。请注意，启用支持 CUDA 的 MPI 可能需要一些额外的步骤。在我们的案例中，我们将坚持使用不支持 GPU 的 Open-MPI：conda install -c conda-forge openmpi

  3. 现在，转到您克隆的 PyTorch 仓库并执行 python setup.py install。

为了测试我们新安装的后端，需要进行一些修改：

  1. 将 if __name__ == '__main__': 下的内容替换为 init_process(0, 0, run, backend='mpi')。

  2. 运行 mpirun -n 4 python myscript.py。

做出这些改变的原因是 MPI 需要在派生进程之前创建自己的环境。MPI 也会派生它自己的进程并执行 初始化方法 中描述的握手，这使得 init_process_group 的 rank 和 size 参数变得多余。这实际上非常强大，因为您可以向 mpirun 传递额外的参数，以便为每个进程定制计算资源。（例如每个进程的核心数、手动将机器分配给特定的 rank 以及 其他一些设置）。这样做，您应该会获得与其他通信后端相同的熟悉输出。

NCCL 后端

NCCL 后端针对 CUDA 张量提供了优化的集合通信操作实现。如果您在集合通信操作中仅使用 CUDA 张量，请考虑使用此后端以获得同类最佳的性能。NCCL 后端包含在支持 CUDA 的预构建二进制文件中。

XCCL 后端

XCCL 后端针对 XPU 张量提供了优化的集合通信操作实现。如果您的工作负载在集合通信操作中仅使用 XPU 张量，则此后端可提供同类最佳的性能。XCCL 后端包含在支持 XPU 的预构建二进制文件中。

初始化方法#

作为本教程的结尾，让我们来看看我们调用的初始函数：dist.init_process_group(backend, init_method)。具体来说，我们将讨论负责每个进程之间初步协调步骤的各种初始化方法。这些方法允许您定义如何完成这种协调。

初始化方法的选择取决于您的硬件设置，某种方法可能比其他方法更合适。除了以下各节之外，还请参阅官方文档以获取更多信息。

环境变量

在本教程中，我们一直在使用环境变量初始化方法。通过在所有机器上设置以下四个环境变量，所有进程都将能够正确连接到主节点，获取有关其他进程的信息，并最终与它们进行握手。

  • MASTER_PORT：托管 rank 0 进程的机器上的一个空闲端口。

  • MASTER_ADDR：托管 rank 0 进程的机器的 IP 地址。

  • WORLD_SIZE：进程总数，以便主节点知道需要等待多少个工作进程（worker）。

  • RANK：每个进程的 rank，以便它们知道自己是主节点还是工作进程。

共享文件系统

共享文件系统法要求所有进程都具有对共享文件系统的访问权限，并将通过共享文件协调它们。这意味着每个进程都将打开该文件、写入其信息，并等待直到所有进程都这样做。之后，所有必需的信息都将易于供所有进程使用。为了避免竞态条件，文件系统必须支持通过 fcntl 进行加锁。

dist.init_process_group(
    init_method='file:///mnt/nfs/sharedfile',
    rank=args.rank,
    world_size=4)

TCP

通过 TCP 进行初始化可以通过提供 rank 0 进程的 IP 地址和可达的端口号来实现。在这里，所有工作进程都将能够连接 to rank 0 进程并交换有关如何访问彼此的信息。

dist.init_process_group(
    init_method='tcp://10.1.1.20:23456',
    rank=args.rank,
    world_size=4)

致谢

我想感谢 PyTorch 的开发人员，他们在实现、文档和测试方面做了如此出色的工作。当代码不清晰时，我总是可以依靠文档或测试来寻找答案。特别是，我要感谢 Soumith Chintala、Adam Paszke 和 Natalia Gimelshein 在早期草稿中提供了极具洞察力的建议并解答了疑问。

在此页面上
PyTorch 库