PyTorch实战（26）——PyTorch分布式训练

PyTorch实战（26）------PyTorch分布式训练

- [0. 前言](#0. 前言)
- [1. 使用 PyTorch 进行分布式训练](#1. 使用 PyTorch 进行分布式训练)
- - [1.1 以常规方式训练模型](#1.1 以常规方式训练模型)
  - [1.2 分布式训练模型](#1.2 分布式训练模型)
- [2. GPU 分布式训练](#2. GPU 分布式训练)
- 小结
- 系列链接

0. 前言

在将预训练的机器学习模型投入生产环境之前，模型训练是不可或缺的关键环节。随着深度学习的发展，大模型往往具有数百万乃至数十亿参数。使用反向传播来调整这些参数需要大量的内存和计算资源。即便如此，模型训练仍然可能需要数天甚至数月时间才能完成。

在本节中，我们将探讨如何通过跨机器和机器内多进程的分布式训练来加速模型训练过程。我们将系统学习 PyTorch 提供的三大分布式训练 API------torch.distributed、torch.multiprocessing 以及 torch.utils.data.distributed.DistributedSampler，使用这些 API 能够极大的简化分布式训练，介绍如何使用 PyTorch 的分布式训练工具，在 CPU 和 GPU 上加速训练。

通过本节学习，将能够充分释放硬件设备的训练潜力。对于超大规模模型训练而言，本节所探讨的工具不仅至关重要，在某些情况下甚至是不可或缺的。

1. 使用 PyTorch 进行分布式训练

在本节中，我们将模型训练过程从常规训练转换为分布式训练，探讨 PyTorch 提供的分布式训练工具，这些工具能显著提升训练速度并优化硬件使用效率。

1.1 以常规方式训练模型

(1) 首先导入所需库：

python 复制代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import time
import argparse

device = torch.device("cpu")

(2) 接下来，定义卷积神经网络 (Convolutional Neural Network, CNN)模型架构：

python 复制代码

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout(0.10)
        self.dp2 = nn.Dropout(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op

(3) 然后，定义模型的训练过程：

python 复制代码

def train(args):
    train_dataloader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1302,), (0.3069,))])),
        batch_size=128, shuffle=True)  
    model = ConvNet()
    optimizer = optim.Adadelta(model.parameters(), lr=0.5)
    model.train()

在函数的前半部分，使用 PyTorch 训练数据集定义了 PyTorch 的训练数据加载器。实例化卷积神经网络 (ConvNet)，并定义了优化器。

python 复制代码

    for epoch in range(args.epochs):
        for b_i, (X, y) in enumerate(train_dataloader):
            X, y = X.to(device), y.to(device)
            pred_prob = model(X)
            loss = F.nll_loss(pred_prob, y) # nll is the negative likelihood loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if b_i % 10 == 0:
                print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(
                    epoch, b_i, len(train_dataloader),
                    100. * b_i / len(train_dataloader), loss.item()))

在函数的后半部分，运行训练循环预定义的 epoch 数。在循环内，通过批数据的方式遍历整个训练数据集，本节中批大小为 128。对于每个包含 128 个训练数据点的批次，使用模型进行前向传播，以计算预测概率。然后，我们将预测结果结合真实标签计算批次损失，并通过反向传播利用该损失梯度来调整模型参数。

(4) 将所有组件整合在 main() 函数中：

python 复制代码

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', default=1, type=int)
    args = parser.parse_args()
    start = time.time()
    train(args)
    print(f"Finished training in {time.time()-start} secs")

使用参数解析器，它可以帮助我们在运行 Python 训练程序时从命令行输入超参数，例如 epoch 数。我们还对训练过程进行了计时，以便可以将它与分布式训练过程进行比较。

(5) 最后，确保通过命令行执行脚本时能运行 main() 函数：

python 复制代码

if __name__ == '__main__':
    main()

(6) 在命令行中执行以下命令来运行该 Python 脚本：

shell 复制代码

$ python convnet_undistributed.py -- epoches 1

本节中我们仅设置训练一个 epoch，因为当前重点不在于模型精度，而在于模型训练耗时。可以看到输出结果如下所示：

训练 1 个 epoch 大约花费了 28 秒，一个 epoch 共包含 469 个批次，实际训练时间会随硬件配置差异而波动。

1.2 分布式训练模型

通过使用 PyTorch 提供的分布式处理 API，即使存在跨进程或跨机器重复传递数据的额外开销，模型训练速度也能显著提升。

(1) 首先，导入所需库，增加几个与分布式训练相关的模块：

python 复制代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import torch.multiprocessing as mp
import torch.distributed as dist

import os
import time
import argparse

torch.multiprocessing 用于在单台机器上生成多个 Python 进程(通常根据 CPU 核心数生成对应数量的进程)，而torch.distributed则实现不同机器间的通信协作，使它们能共同完成模型训练。执行时，我们需要在每台参与训练的机器上显式启动训练脚本。
PyTorch 内置的通信后端(如 Gloo )会自动处理机器间的通信协调。在每台机器内部，多进程机制会进一步将训练任务并行分配到各个进程。

(2) 模型架构定义部分保持不变：

python 复制代码

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout2d(0.10)
        self.dp2 = nn.Dropout2d(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op

(3) 定义 train() 函数：

python 复制代码

def train(cpu_num, args):
    rank = args.machine_id * args.num_processes + cpu_num                        
    dist.init_process_group(                                   
    backend='gloo',                                         
    init_method='env://',                                   
    world_size=args.world_size,                              
    rank=rank                                               
    ) 
    torch.manual_seed(0)
    device = torch.device("cpu")

可以看到，代码开头新增了两条关键语句。首先是计算进程的 rank 值------这本质上是该进程在整个分布式系统中的顺序标识符。举例来说，若使用 2 台各配备 4 核 CPU 的机器进行训练，为充分利用硬件资源可能需要启动 8 个进程(每台机器 4 个)。此时就需要为这些进程建立标识体系：先为两台机器分配 ID 0 和 1，再为每台机器内的 4 个进程分配子 ID 0 至 3。最终，第 n 台机器上第 i 个进程的全局 rank 值可通过以下公式确定：
r a n k = n × 4 + k rank=n\times 4+k rank=n×4+k

第二行代码使用了 torch.distributed 模块中的 init_process_group，该方法为每个启动的进程配置以下关键参数：

用于机器间通信的后端(在节使用 Gloo)
参与分布式训练的进程总量(由 args.world_size 指定)，亦称 world_size
当前启动进程的全局 rank 值

init_process_group 方法会阻塞所有进程，直到跨机器的全部进程都完成初始化才会继续执行。
PyTorch 提供了三种内置的分布式训练后端：

Gloo
NCCL
MPI

简而言之，对于 CPU 上的分布式训练，使用 Gloo，对于 GPU，使用 NCCL。

python 复制代码

    train_dataset = datasets.MNIST('./data', train=True, download=True,
                                   transform=transforms.Compose([
                                       transforms.ToTensor(),
                                       transforms.Normalize((0.1302,), (0.3069,))]))  
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )
    train_dataloader = torch.utils.data.DataLoader(
       dataset=train_dataset,
       batch_size=args.batch_size,
       shuffle=False,            
       num_workers=0,
       sampler=train_sampler)
    model = ConvNet()
    optimizer = optim.Adadelta(model.parameters(), lr=0.5)
    model = nn.parallel.DistributedDataParallel(model)
    model.train()

与单机训练相比，分布式训练的关键改进体现在数据加载与模型封装两个层面。我们将 MNIST 数据集实例化与数据加载器拆分为独立步骤，其间插入 DistributedSampler 采样器。该采样器将训练数据均分为 world_size 个分区，确保每个进程处理等量数据。注意数据加载器的 shuffle 参数需设为 False，因为数据分配已由采样器控制。

代码中的另一个新增部分是 nn.parallel.DistributedDataParallel 函数，它应用于模型对象。这部分可能是代码中最重要的部分，因为 DistributedDataParallel 是实现分布式梯度下降算法的关键组件。其底层运行机制如下：

分布式环境中的每个派生进程都会获得独立的模型副本
每个进程的模型都维护自己的优化器，并与全局迭代保持同步的局部优化步骤
在每次分布式训练迭代时，各进程独立计算损失值及梯度，随后跨进程对这些梯度求取平均值
平均后的梯度将通过全局反向传播机制同步到所有模型副本，用于调整参数
由于全局反向传播步骤的存在，所有模型参数在每次迭代时都保持一致，从而实现自动同步

DistributedDataParallel 通过让每个 Python 进程运行在独立的解释器上，有效规避了在单一解释器下多线程实例化多个模型可能引发的全局解释器锁 (Global Interpreter Lock, GIL) 限制问题。这进一步提升了性能表现，特别是对于那些需要大量Python专属运算的模型而言。

python 复制代码

    for epoch in range(args.epochs):
        for b_i, (X, y) in enumerate(train_dataloader):
            X, y = X.to(device), y.to(device)
            pred_prob = model(X)
            loss = F.nll_loss(pred_prob, y) # nll is the negative likelihood loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if b_i % 10 == 0 and cpu_num==0:
                print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(
                    epoch, b_i, len(train_dataloader),
                    100. * b_i / len(train_dataloader), loss.item()))

最后，训练循环几乎和单机训练一样。唯一的区别在于我们限制只有排名为0的进程才能获取日志信息。这是因为排名为 0 的机器用于建立所有通信连接。因此，我们通常将排名为 0 的进程作为参考来跟踪模型训练性能。如果不加以限制，每个模型训练迭代都会产生与进程数量相同的日志行数。

(4) 将所有组件整合在 main() 函数中：

python 复制代码

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-machines', default=1, type=int,)
    parser.add_argument('--num-processes', default=1, type=int)
    parser.add_argument('--machine-id', default=0, type=int)
    parser.add_argument('--epochs', default=1, type=int)
    parser.add_argument('--batch-size', default=128, type=int)
    args = parser.parse_args()
    
    args.world_size = args.num_processes * args.num_machines                
    os.environ['MASTER_ADDR'] = '127.0.0.1'              
    os.environ['MASTER_PORT'] = '8892'      
    start = time.time()
    mp.spawn(train, nprocs=args.num_processes, args=(args,))
    print(f"Finished training in {time.time()-start} secs")

首先，我们新增了以下参数：

num_machines：机器总数量
num_processes：每台机器上要启动的进程数量
machine_id：当前机器的序号 ID。需要注意的是，这个 Python 脚本需要在每台机器上单独启动
batch_size：每个批次的数据样数量。参数作用如下：
- 所有进程将各自计算梯度，这些梯度将在每次迭代中被平均，从而获得整体梯度
- 完整的训练数据集被分割为 world_size 个独立的数据集

因此，在每次迭代时，完整的批次数据需要被分割成 world_size 个子批次，每个进程处理一个子批次。因为 batch_size 现在与 world_size 相关联，所以我们将其作为输入参数提供，目的是为了简化训练接口。

参数定义后，计算 world_size 作为派生参数。接着，我们定义两个重要的环境变量：

MASTER_ADDR：运行 rank 0 进程的主机 IP 地址
MASTER_PORT：运行 rank 0 进程的主机上可用的端口号

rank 0 机器负责建立所有后端通信连接，因此整个系统必须能随时定位到该主机，这就是为什么需要提供其 IP 地址和端口号。本节中训练任务将在单台本地机器上运行，因此使用 localhost 地址即可，但在跨服务器的多机训练场景中，则需要提供 rank 0 服务器的真实 IP 地址及空闲端口号。

最后一个变化是使用多进程 (multiprocessing) 来在每台机器上启动 num_processes 个进程，而非仅运行单个训练进程。分布式参数会传递给每个派生进程，确保模型训练过程中各进程与机器之间能自主协调。

(5) 分布式训练：

python 复制代码

if __name__ == '__main__':
    main()

(6) 启动分布式训练脚本。首先使用分布式脚本进行类非分布式运行，将机器数量和进程数量都设置为 1：

shell 复制代码

$ python convnet_distributed.py --num-machines 1 --num-processes 1 --machine-id 0 --epochs 1 --batch-size 128

需要注意的是，由于本次训练只使用单个进程，batch_size 与之前非分布式训练时保持一致(仍为 128)。运行结果如下输出：

若将此结果与上一节非分布式训练的输出对比，可发现训练时间基本相当(约 30 秒)，损失值变化也十分相似。

(7) 接下来，运行一个真正的分布式训练，使用 2 个进程而不是 1 个进程。相应地，将 batch_size 从 128 降至 64：

shell 复制代码

$ python convnet_distributed.py --num-machines 1 --num-processes 2 --machine-id 0 --epochs 1 --batch-size 64

输出结果如下所示：

可以看到，训练时间从 30 秒减少到了 20 秒。训练损失的变化趋势没有受到影响，这表明分布式训练可以加速训练过程，同时保持模型的准确性。

(8) 接下来，使用 4 个进程，并相应地将批大小从 64 降至 32：

shell 复制代码

$ python convnet_distributed.py --num-machines 1 --num-processes 4 --machine-id 0 --epochs 1 --batch-size 32

输出结果如下所示：

可以看到，训练时间进一步减少，从 20 秒降至 15 秒。训练损失的变化趋势仍然与之前的训练相似。通过分布式训练，我们已经将训练时间从 30 秒缩短到了 15 秒，减少了 2 倍。

(9) 进一步增加进程数，使用 8 个进程代替 4 个进程，并相应地将批次大小从 32 降至 16：

shell 复制代码

$ python convnet_distributed.py --num-machines 1 --num-processes 8 --machine-id 0 --epochs 1 --batch-size 16

输出结果如下所示：

与预期相反，训练时间不仅没有进一步缩短，反而从 15 秒略微增加至 18 秒。由于代码在本地机器执行，系统还存在其他进程(如浏览器)会与部分分布式训练进程争夺资源。如果分布式训练模型是在远程机器上进行的，同时这些机器的唯一任务就是进行模型训练，在这样的机器上，建议使用与 CPU 核心数相等甚至更多的进程数。

(10) 最后需要指出的是，由于在本节中我们只使用了一台机器，因此我们只需要启动一个 Python 脚本来开始训练。然而，如果是在多台机器上进行训练，那么除了修改 MASTER_ADDR 和 MASTER_PORT 外，还需要在每台机器上启动一个 Python 脚本。例如，如果有 2 台机器，在机器 1 上执行：

shell 复制代码

$ python distributed_script.py --num_machines=2 --num-processes 8 --machine_id=0 --epochs 1 --batch-size 16

在机器 2 上执行：

shell 复制代码

$ python distributed_script.py --num_machines=2 --num-processes 8 --machine_id=1 --epochs 1 --batch-size 16

至此，我们完成了关于使用 PyTorch 在 CPU 上实施分布式训练深度学习模型的实践探讨，这种方法能带来显著的加速效果。仅需添加少量代码，就能将常规 PyTorch 模型训练脚本升级为分布式训练模式。虽然上述实验基于简单的卷积网络，但由于我们完全无需修改模型架构代码，因此这套方案可直接扩展到更复杂的模型训练场景。接下来，我们将简要讨论如何应用类似的代码更改，实现 GPU 环境下的分布式训练。

2. GPU 分布式训练

我们通常使用以下 PyTorch 代码定义模型训练设备：

python 复制代码

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

这行代码的作用是自动检测可用计算设备，并优先选择 CUDA (GPU)。这种优先选择源于 GPU 通过并行化处理神经网络常规运算(如矩阵乘法和加法)所能提供的显著加速优势。本节我们将探讨如何通过 GPU 分布式训练进一步加速模型训练。

(1) 虽然导入语句和模型架构定义代码与使用 CPU 分布式训练一节完全一致，但 train() 函数中有几处关键修改：

python 复制代码

def train(gpu_num, args):
    rank = args.machine_id * args.num_gpu_processes + gpu_num                        
    dist.init_process_group(                                   
        backend='nccl',                                         
        init_method='env://',                                   
        world_size=args.world_size,                              
        rank=rank                      
    ) 
    model = ConvNet()
    torch.cuda.set_device(gpu_num)
    model.cuda(gpu_num)
    criterion = nn.NLLLoss().cuda(gpu_num)

在使用 GPU 时，NCCL 是首选的通信后端。同时，模型和损失函数都必须部署到 GPU 设备上，以确保充分利用 GPU 提供的并行矩阵运算加速能力：

python 复制代码

    train_dataset = datasets.MNIST('./data', train=True, download=True,
                                   transform=transforms.Compose([
                                       transforms.ToTensor(),
                                       transforms.Normalize((0.1302,), (0.3069,))]))  
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )
    train_dataloader = torch.utils.data.DataLoader(
       dataset=train_dataset,
       batch_size=args.batch_size,
       shuffle=False,            
       num_workers=0,
       pin_memory=True,
       sampler=train_sampler)
    optimizer = optim.Adadelta(model.parameters(), lr=0.5)
    model = nn.parallel.DistributedDataParallel(model,
                                                device_ids=[gpu_num])
    model.train()

DistributedDataParallel API 包含一个关键参数------device_ids，用于指定调用该 API 的 GPU 进程 ID。此外可以看到，数据加载器 (dataloader) 中新增了 pin_memory 参数并设为 True，该参数能显著加速训练过程中从主机(此处指加载数据集的 CPU )到各设备 (GPU) 的数据传输。
pin_memory 机制的工作原理是将数据"锁定" (pin) 在 CPU 内存中，即把数据样本分配到固定的页锁定内存区域。训练时，这些内存区域的数据会被高效地拷贝到对应 GPU。该机制需与 non_blocking=True 参数配合使用：

python 复制代码

    for epoch in range(args.epochs):
        for b_i, (X, y) in enumerate(train_dataloader):
            X, y = X.cuda(non_blocking=True), y.cuda(non_blocking=True)
            pred_prob = model(X)
            loss = criterion(pred_prob, y) 
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if b_i % 10 == 0 and gpu_num==0:
                print('epoch: {} [{}/{} ({:.0f}%)]\t training loss: {:.6f}'.format(
                    epoch, b_i, len(train_dataloader),
                    100. * b_i / len(train_dataloader), loss.item()))

通过调用参数 pin_memory 和 non_blocking，使得以下两者之间的操作得以重叠：

从 CPU 到 GPU 数据(真实标签)的传输
GPU 模型训练计算(或 GPU 内核执行)

这从根本上提升了整体 GPU 训练流程的效率。

(2) 除了 train() 函数的修改外，main() 函数同样有所调整调整：

python 复制代码

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num-machines', default=1, type=int,)
    parser.add_argument('--num-gpu-processes', default=1, type=int)
    parser.add_argument('--machine-id', default=0, type=int)
    parser.add_argument('--epochs', default=1, type=int)
    parser.add_argument('--batch-size', default=64, type=int)
    args = parser.parse_args()
    
    args.world_size = args.num_gpu_processes * args.num_machines                
    os.environ['MASTER_ADDR'] = '127.0.0.1'              
    os.environ['MASTER_PORT'] = '8892'      
    start = time.time()
    mp.spawn(train, nprocs=args.num_gpu_processes, args=(args,))
    print(f"Finished training in {time.time()-start} secs")

用 num_gpu_processes 替代了原来的 num_process 参数，该参数值通过 torch.cuda.device_count() 自动获取可用 GPU 数量。其余代码相应调整，但 GPU 版本的核心逻辑与之前保持一致。执行以下命令即可启动 GPU 分布式训练：

shell 复制代码

$ python convnet_distributed_cuda.py --num-machines 1 --num-gpu-processes 1 --machine-id 0 --epochs 1 --batch-size 128

至此，我们已完成关于使用 PyTorch 进行 GPU 分布式模型训练的简要探讨。上述代码同样适用于其他深度学习模型，当前深度学习模型大多采用 GPU 分布式训练方案。此外，Horovod、DeepSpeed 和 PyTorch Lightning 等库都提供了更简洁的 API 来简化 PyTorch 模型的分布式训练流程。

小结

在本节中，我们探讨了机器学习中一个重要的实践方面------如何优化模型训练过程，介绍了使用 PyTorch 在 CPU 和 GPU 上进行分布式训练的适用范围与强大效能。

系列链接

PyTorch实战（1）------深度学习（Deep Learning）
PyTorch实战（2）------使用PyTorch构建神经网络
 PyTorch实战（3）------PyTorch vs. TensorFlow详解
 PyTorch实战（4）------卷积神经网络（Convolutional Neural Network，CNN）
PyTorch实战（5）------深度卷积神经网络
 PyTorch实战（6）------模型微调详解
 PyTorch实战（7）------循环神经网络
 PyTorch实战（8）------图像描述生成
 PyTorch实战（9）------从零开始实现Transformer
PyTorch实战（10）------从零开始实现GPT模型
 PyTorch实战（11）------随机连接神经网络（RandWireNN）
PyTorch实战（12）------图神经网络（Graph Neural Network，GNN）
PyTorch实战（13）------图卷积网络（Graph Convolutional Network，GCN）
PyTorch实战（14）------图注意力网络（Graph Attention Network，GAT）
PyTorch实战（15）------基于Transformer的文本生成技术
 PyTorch实战（16）------基于LSTM实现音乐生成
 PyTorch实战（17）------神经风格迁移
 PyTorch实战（18）------自编码器（Autoencoder，AE）
PyTorch实战（19）------变分自编码器（Variational Autoencoder，VAE）
PyTorch实战（20）------生成对抗网络（Generative Adversarial Network，GAN）
PyTorch实战（21）------扩散模型（Diffusion Model）
PyTorch实战（22）------MuseGAN详解与实现
 PyTorch实战（23）------基于Transformer生成音乐
 PyTorch实战（24）------深度强化学习
 PyTorch实战（25）------使用PyTorch构建DQN模型