PyTorch 分布式通讯包-Torch.Distributed

一、后端

torch.distributed 支持三个后端：Gloo、MPI 和 NCCL，每个后端功能不同。下表展示了各后端在 CPU 和 GPU 上的支持情况。

后端	CPU 支持	GPU 支持
Gloo	✓	✘
MPI	✓	✘
NCCL	✘	✓

使用建议

分布式 GPU 训练 ：优先选择 NCCL 后端，它提供最佳的分布式 GPU 训练性能，尤其适用于多进程单节点或多节点场景。若使用 NCCL 遇到问题，可使用 Gloo 作为后备选项。
分布式 CPU 训练 ：建议使用 Gloo 后端。

环境变量

选择网络接口 ：
- NCCL：可使用 export NCCL_SOCKET_IFNAME=eth0 指定网络接口。
- Gloo：可使用 export GLOO_SOCKET_IFNAME=eth0 指定网络接口，还支持指定多个接口，如 export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3，后端会在这些接口间循环调度操作。

NCCL 调试 ：常用环境变量有 export NCCL_DEBUG=INFO 和 export NCCL_DEBUG_SUBSYS=ALL。

二、初始化

使用 torch.distributed.init_process_group() 函数初始化分布式包。该函数会阻塞，直到所有进程都加入。

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

初始化方法

TCP 初始化 ：需要指定一个地址（属于 rank 0 进程）和 world_size，要求所有进程都具有手动指定的 rank。
共享文件系统初始化 ：利用共享文件系统及 world_size 进行初始化， URL 应以 file:// 开头。注意，每次初始化需要一个全新的空文件，且有责任在训练结束时删除该文件。
环境变量初始化 ：从环境变量中读取配置，需设置的变量有 MASTER_PORT（必填，等级为 0 的计算机上的空闲端口）、MASTER_ADDR（必填，等级 0 节点的地址，等级 0 可以不填）、WORLD_SIZE（必填，可在环境变量或调用 init 函数时设置）、RANK（必填，可在环境变量或调用 init 函数时设置）。

三、通信原语

（一）点对点通信

同步发送张量 ：torch.distributed.send(tensor, dst, group=<object object>, tag=0)，其中 tensor 是要发送的张量，dst 是目标排名，group 是进程组（可选），tag 是与远程 recv 发送匹配的标签。
同步接收张量 ：torch.distributed.recv(tensor, src=None, group=<object object>, tag=0)，其中 tensor 用于填充接收到的数据，src 是源排名（可选），若未指定则从任何进程接收。
异步发送张量 ：torch.distributed.isend(tensor, dst, group=<object object>, tag=0)，返回分布式请求对象。
异步接收张量 ：torch.distributed.irecv(tensor, src, group=<object object>, tag=0)，返回分布式请求对象。

（二）同步和异步集体操作

每个集体操作功能都支持同步和异步两种操作：

同步操作 ：async_op 设置为 False 时的默认模式，函数返回时确保执行了集合操作。
异步操作 ：当 async_op 设置为 True 时，集合操作函数返回一个分布式请求对象，可通过 is_completed() 判断操作是否完成，wait() 阻塞进程直到操作完成。

（三）集体函数

广播：torch.distributed.broadcast(tensor, src, group=<object object>, async_op=False)，向整个组广播张量。
全部归约 ：torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=<object object>, async_op=False)，减少所有机器上的张量数据，使所有机器得到最终结果。
归约：torch.distributed.reduce(tensor, dst, op=ReduceOp.SUM, group=<object object>, async_op=False)，减少所有机器上的张量数据，只有目标进程收到最终结果。
全部聚集 ：torch.distributed.all_gather(tensor_list, tensor, group=<object object>, async_op=False)，收集整个组的张量到列表。
聚集：torch.distributed.gather(tensor, gather_list=None, dst=0, group=<object object>, async_op=False)，在单个过程中收集张量列表。
分散：torch.distributed.scatter(tensor, scatter_list=None, src=0, group=<object object>, async_op=False)，将张量列表分散到组中的所有进程。

四、多 GPU 集合功能

如果每个节点上有多个 GPU，在使用 NCCL 和 Gloo 后端时，支持在每个节点内的多个 GPU 之间进行分布式集体操作，如 broadcast_multigpu()、all_reduce_multigpu()、reduce_multigpu() 和 all_gather_multigpu()。这些功能可改善整体分布式训练性能，使用时需确保传递的张量列表中的每个张量位于调用该函数的主机的单独 GPU 设备上，且在所有分布式过程中，张量列表的长度相同。

五、启动工具

torch.distributed 包提供了启动工具 torch.distributed.launch，可用于为每个节点启动多个进程以进行分布式训练，支持 python2 和 python3。此外，torch.multiprocessing.spawn() 也可用于产生多个进程进行多进程分布式训练，但需要 Python 3.4 或更高版本。

六、实际案例

假设我们要开发一个分布式深度学习模型，用于图像分类任务。我们将展示如何利用分布式通讯包加速模型的训练过程。

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim


## 定义一个简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.fc1 = nn.Linear(32 * 26 * 26, 10)


    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = x.view(-1, 32 * 26 * 26)
        x = self.fc1(x)
        return x


## 初始化模型、损失函数和优化器
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)


## 将模型和数据移动到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)


## 初始化分布式进程组
dist.init_process_group(backend="nccl", init_method="env://")


## 假设我们有训练数据
## train_loader 是一个数据加载器，用于加载训练数据
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)


        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()


    print(f"Epoch {epoch + 1}, Loss: {loss.item()}")


## 保存模型
torch.save(model.state_dict(), "model.pth")

七、总结

本教程介绍了 PyTorch 分布式通讯包的基础知识，包括后端选择、初始化、通信原语以及多 GPU 集合功能等。掌握这些概念对于进行分布式深度学习开发非常重要，希望读者能通过这些知识充分利用分布式计算资源，加速深度学习项目。