【分布式训练】基于Pytorch的分布式数据并行训练

简介： 在PyTorch中使用DistributedDataParallel进行多GPU分布式模型训练

动机

加速神经网络训练的最简单方法是使用GPU，它在神经网络中常见的计算类型（矩阵乘法和加法）上提供了比CPU更大的加速。随着模型或数据集变得越来越大，一个GPU很快就会变得不足。例如，像BERT和GPT-2这样的大型语言模型是在数百个GPU上训练的。要执行多GPU训练，我们必须有一种方法在不同的GPU之间分割模型和数据，并协调训练。

为什么要并行分布数据？

很多人喜欢在Pytorch中实现自己的深度学习模型，因为它在神经网络框架的控制和易用性之间具有最佳平衡。Pytorch有两种方法可以跨多个GPU拆分模型和数据：nn.DataParallel和nn.DistributedDataParallel。

nn.DataParallel更易于使用（只需包装模型并运行训练脚本）。然而，由于它使用一个process来计算模型权重，然后在每个批次中将其分配给每个GPU，因此网络很快成为瓶颈，GPU利用率通常很低。此外，nn.DataParallel要求所有GPU都在同一个节点上，并且不能与Apex一起用于mixed-precision训练。

因此，nn.DataParallel和nn.DistributedDataParallel的主要差异可以总结为以下几点：
1. DistributedDataParallel支持模型并行，而DataParallel不支持，这意味着如果模型太大单卡显存不足时只能使用前者；
2. DataParallel是单进程多线程，只用于单机情况，而DistributedDataParallel是多进程的，适用于单机和多机情况，真正实现分布式训练；
3. DistributedDataParallel的训练更高效，因为每个进程都是独立的Python解释器，避免GIL问题，而且通信成本低，其训练速度更快，基本上DataParallel已经被弃用；
4. 必须要说明的是DistributedDataParallel中每个进程都有独立的优化器，执行自己的更新过程，但是梯度通过通信传递到每个进程，所有执行的内容是相同的。

现有资料的不足

总的来说，Pytorch文档是完整且清晰的，但是，当试图弄清楚如何使用DistributedDataParallel时，发现所有的示例和教程都是不可访问、不完整或重载了不相关的功能的组合。

Pytorch提供了一个关于使用AWS进行分布式培训的教程，它很好地展示了如何在AWS方面进行设置。然而，它的其余部分有点混乱，因为出于某种原因，它花了很多时间来展示如何计算指标，然后再回到展示如何包装模型和启动流程。它也没有描述nn.DistributedDataParallel的作用，这使得相关的代码块很难遵循。

关于用Pytorch编写分布式应用程序的教程比第一遍所需的要详细得多，而且对于没有Python多处理背景的人来说是无法访问的。它花费了大量时间复制nn.DistributedDataParallel中的功能。然而，它没有给出它所做工作的高级概述，也没有提供如何使用它的见解(https://pytorch.org/tutorials/intermediate/ddp_tutorial.html)

还有一个Pytorch教程，介绍如何开始使用分布式数据并行。这一个展示了如何进行一些设置，但没有解释设置的目的，然后展示了一些代码，以在GPU之间拆分模型并进行一个优化步骤。不幸的是，我很确定编写的代码不会运行（函数名不匹配），而且它也没有告诉你如何运行代码。与前面的教程一样，它也没有对分布式培训的工作原理进行高层次的概述。

Pytorch提供的最接近MWE示例的是Imagenet训练示例。不幸的是，这个例子还展示了Pytorch几乎所有的其他功能，因此很难找出分布式多GPU训练的相关内容。

Apex提供了他们自己版本的Pytorch Imagenet示例。他们的nn.DistributedDataParallel版本是Pytorch的替代品，只有在学习如何使用Pytork之后才有帮助。

这个教程很好地描述了引擎盖下发生的事情以及它与nn.DataParallel的区别。但是，它没有关于如何使用nn.DataParallel的代码示例。

Outline

本教程真正针对的是那些已经熟悉在Pytorch中训练神经网络模型的人。首先概述整体思想。然后，展示了在GPU上使用MNIST进行训练的最小工作示例。我修改了这个例子，在多个GPU上进行训练，可能跨越多个节点，并逐行解释这些变化。重要的是，还解释了如何运行代码。作为奖励，还演示了如何使用Apex进行简单的混合精度分布式训练。

整体框架图

使用DistributedDataParallel进行Multiprocessing会在多个GPU上复制模型，每个GPU都由一个进程控制。（进程是在计算机上运行的python的一个实例；通过让多个进程并行运行，我们可以利用具有多个CPU核心的procressor。如果你愿意，可以让每个进程控制多个GPU，但这显然比每个进程有一个GPU要慢。也可以让多个工作进程为每个GPU获取数据，但为了简单起见，将省略这一点。）GPU可以全部位于同一节点上，也可以分布在多个节点上。（一个节点是一台“计算机”，包括它的所有CPU和GPU。如果你使用AWS，一个节点就是一个EC2实例。）每个进程都执行相同的任务，每个进程都与所有其他进程通信。只有梯度在进程/GPU之间传递，这样网络通信就不会成为瓶颈。

在训练过程中，每个进程从磁盘加载自己的mini-batches，并将其传递给GPU。每个GPU都有自己的前向通道，然后GPU之间的梯度都会减小。每一层的梯度不依赖于前一层，因此梯度all-reduce与反向传递同时计算，以进一步缓解网络瓶颈。在反向过程结束时，每个节点都具有平均梯度，确保模型权重保持同步。

所有这些都要求多个进程（可能在多个节点上）同步并进行通信。Pytorch通过其distributed.init_process_group函数来实现这一点。此函数需要知道在哪里可以找到进程0，以便所有进程都可以同步，以及预期的进程总数。每个单独的进程还需要知道进程的总数及其在进程中的排名以及使用哪个GPU。将进程的总数称为world size是很常见的。最后，每个进程都需要知道要处理哪一部分数据，以便批处理不重叠。Pytorch提供了nn.utils.data.DistributedSampler来实现这一点，即为各个进程切分数据，以保证训练数据不重叠。

更详细的DDP的内部机理见官方的文档介绍：DISTRIBUTED DATA PARALLEL

带解释的最小demo示例

为了演示如何做到这一点，将创建一个在MNIST上训练的示例，然后将其修改为在多个节点的多个GPU上运行，最后还允许混合精度训练（mixed-precision training）。

没有multiprocessing

首先，导入所需要的依赖库：

import os
import argparse
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms
from datetime import datetime
from apex.parallel import  DistributedDataParallel as DDP
from apex import amp

我们定义了一个非常简单的卷积模型来预测MNIST。

class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc = nn.Linear(7 * 7 * 32, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

下面是训练过程：

def train(gpu, args):
    torch.manual_seed(0)
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    # model = nn.DataParallel(model, device_ids=device_ids)
    # model = model.cuda(device=gpu)
    batch_size = 100
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(),
                                               download=True)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
                                               num_workers=0, pin_memory=True)

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i + 1) % 100 == 0 and gpu == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                         loss.item()))
    if gpu == 0:
        print("Training complete in: " + str(datetime.now() - start))

main()函数将接受一些参数并运行训练函数。

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--nodes", default=1, type=int, metavar='N')
    parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
    parser.add_argument('--epochs', default=2, type=int, metavar='N', help='number of total epochs to run')
    args = parser.parse_args()
    train(0, args)

最后，要确保main()函数被调用。

if __name__ == '__main__':
    main()

可以通过打开一个终端并键入python src/mnist.py-n 1-g 1-nr 0来运行此代码，它将在单个节点上的单个gpu上进行训练。

开启multiprocessing

要使用multiprocessing来做到这一点，我们需要一个脚本来为每个GPU启动一个进程。每个进程都需要知道要使用哪个GPU，以及它在所有正在运行的进程中的排名。需要在每个节点上运行脚本。

来看看对每个函数的更改。为了便于查找，已将新代码隔离开来：

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--nodes", default=1, type=int, metavar='N')
    parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
    parser.add_argument('--epochs', default=2, type=int, metavar='N', help='number of total epochs to run')
    args = parser.parse_args()
    #########################################################
    args.world_size = args.gpus * args.nodes
    os.environ['MASTER_ADDR'] = '172.20.109.105'
    os.environ['MASTER_PORT'] = '8888'
    mp.spawn(train, nprocs=args.gpus, args=(args,))
    #########################################################
    # train(0, args)

其中：

• args.nodes表示节点总数，
• args.gpus表示每个节点的GPU总数（每个节点GPU数是一样的）
• args.nr表示当前节点在所有节点中的序号。

根据节点总数和每个节点的GPU数，可以计算world_size，即要运行的进程总数，所有的进程需要知道进程0的IP地址以及端口，这样所有进程可以在开始时同步，一般情况下称进程0是master进程，比如我们会在进程0中打印信息或者保存模型。

PyTorch提供了mp.spawn来在一个节点启动该节点所有进程，每个进程运行train(i, args)，其中 i 从0到args.gpus - 1。请记住，在每个节点上运行main()函数，这样总共会有args.nodes*args.gpus=args.world_size进程。

同样，我们要修改训练函数：

def train(gpu, args):
    ############################################################
    rank = args.nr * args.gpus + gpu
    dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    ############################################################
    torch.manual_seed(0)
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    batch_size = 100
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    ############################################################
    # Wrap the model
    model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
    ############################################################

    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(),
                                               download=True)
    ############################################################
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_dataset, num_replicas=args.world_size,
                                                                    rank=rank)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=False,
                                               num_workers=0, pin_memory=True, sampler=train_sampler)
    ############################################################

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i + 1) % 100 == 0 and gpu == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                         loss.item()))
    if gpu == 0:
        print("Training complete in: " + str(datetime.now() - start))

这里，首先计算出当前进程序号：rank = args.nr * args.gpus + gpu，然后就是通过dist.init_process_group初始化分布式环境，其中

• backend参数指定通信后端，包括mpi、gloo、nccl，这里选择nccl，它是Nvidia提供的官方多卡通信框架，相对比较高效。mpi也是高性能计算常用的通信协议，不过需要自己安装MPI实现框架，比如OpenMPI。gloo倒是内置通信后端，但是不够高效。
• init_method指的是如何初始化，以完成刚开始的进程同步，这里我们设置的是env://，指的是环境变量初始化方式，需要在环境变量中配置4个参数：MASTER_PORT，MASTER_ADDR，WORLD_SIZE，RANK，前面两个参数我们已经配置，后面两个参数也可以通过dist.init_process_group函数中world_size和rank参数配置。
  其它的初始化方式还包括共享文件系统以及TCP，比如采用TCP作为初始化方法init_method='tcp://10.1.1.20:23456'，其实也是要提供master的IP地址和端口。注意这个调用是阻塞的，必须等待所有进程来同步，如果任何一个进程出错，就会失败。
• 对于模型侧，只需要用DistributedDataParallel包装一下原来的model即可，将模型复制到GPU上以进行处理，在背后它会支持梯度的All-Reduce操作。
• 对于数据侧，使用nn.utils.data.DistributedSampler来给各个进程切分数据，只需要在dataloader中使用这个sampler就好，值得注意的一点是要在训练循环过程的每个epoch开始时调用train_sampler.set_epoch(epoch)，（主要是为了保证每个epoch的划分是不同的）其它的训练代码都保持不变。

最后就可以执行代码了，比如我们有4节点，每个节点是8张显卡，那么需要在4个节点终端上分别执行：

python src/mnist-distributed.py -n 4 -g 8 -nr i

例如在节点0上执行：

python src/mnist-distributed.py -n 4 -g 8 -nr 0

换句话说，在每个节点上运行这个脚本，告诉它在训练开始前启动彼此同步的args.gpus进程。

要注意的是，此时的有效batch_size其实是batch_size_per_gpu * world_size，对于有BN的模型还可以采用同步BN获取更好的效果：

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

上述讲述的是分布式训练过程，其实同样适用于评估或者测试过程，比如我们把数据划分到不同的进程中进行预测，这样可以加速预测过程。实现代码和上述过程完全一样，不过我们想计算某个指标，那就需要从各个进程的统计结果进行All-Reduce，因为每个进程仅是计算的部分数据的内容。比如我们要计算分类准确度，我们可以统计每个进程的数据总数total和分类正确的数量count，然后进行聚合。

这里要提的一点，当用dist.init_process_group初始化分布式环境时，其实就是建立一个默认的分布式进程组（distributed process group），这个group同时会初始化Pytorch的torch.distributed包。这样我们可以直接用torch.distributed的API就可以进行分布式基本操作了，下面是具体实现：

# define tensor on GPU, count and total is the result at each GPU
t = torch.tensor([count, total], dtype=torch.float64, device='cuda')
dist.barrier()  # synchronizes all processes
dist.all_reduce(t, op=torch.distributed.ReduceOp.SUM,)  # Reduces the tensor data across all machines in such a way that all get the final result.
t = t.tolist()
all_count = int(t[0])
all_total = int(t[1])
acc = all_count / all_total

分布式训练启动方式

上述过程中，采用PyTorch的torch.multiprocessing包（Multiprocessing package - torch.multiprocessing）来启动分布式训练，目前官方给出的ImageNet训练例子是采用这种方式的，detectron2库也是采用这种方式启动：https://github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py。

如果使用torch.multiprocessing.spawn启动，要注意送入的训练function必须是fn(i,*args) 这种格式，其中第一个参数 i 指代的是当前节点的进程编号，这个参数其实就充当了local_rank, 所谓的local_rank是指的训练进程在当前节点的序号，前面说的 rank 其实是全局的进程序号，这个参数很重要，因为要根据这个参数来设置每个进程所使用的 device 设备，一般情况下，直接认为local_rank即为所采用的GPU编号，设置如下：

torch.cuda.set_device(args.local_rank)  # before your code runs

# set DDP
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)

# 或者
with torch.cuda.device(args.local_rank):
    # your code to run

除了采用mp.spawn，还可以采用torch.distributed.launch来启动程序（Distributed communication package - torch.distributed），这个是更常用的启动方式。比如对于单机多卡训练，其启动方式如下：

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
           YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

其中NUM_GPUS_YOU_HAVE是GPU的总量，而YOUR_TRAINING_SCRIPT.py是训练的脚本，其和上述基本一致，不过区别是采用torch.distributed.launch启动，会自动设置一些环境变量（https://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py#L211），比如我们需要的RANK和WORLD_SIZE 就直接可以从环境变量中获取：

rank = int(os.environ["RANK"])
world_size = int(os.environ['WORLD_SIZE'])

对于local_rank的获取有两种方式：
1）一种是在训练脚本添加一个命令行参数，程序启动时会对其自动赋值：

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

local_rank = args.local_rank

2）另外一种方式采用torch.distributed.launch启动加上--use_env=True，此时情况下会设置LOCAL_RANK这个环境变量，就可以从环境变量中获取local_rank：

"""
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --use_env=True
           YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
"""
import os
local_rank = int(os.environ["LOCAL_RANK"]) 

对于多机多卡训练，比如2个node，其启动命令如下所示：

# Node 1: (IP: 192.168.1.1, and has a free port: 1234)
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

# Node 2
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
           --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

这里

• --nnodes表示传入node数目
• --node_rank表示传入node的编号
• world_size=nnodes*nproc_per_node。

不过最新版本的PyTorch推出了torchrun来替代torch.distributed.launch。 torchrun和torch.distributed.launch的用法基本一致，不过弃用了--use_env命令，直接将local_rank设置在环境变量中，目前最新版本的torchvision是采用torchrun启动方式，具体见 vision/references/classification at main · pytorch/vision。

混合精度训练（采用apex）

安装Apex：

git clone https://github.com/NVIDIA/apex
cd apex
# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key... 
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
# otherwise
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./

Apex官方文档：Apex (A PyTorch Extension)

混合精度训练（浮点（FP32）和半精度（FP16）相结合的训练）使我们能够使用更大的batch_size，并利用NVIDIA Tensor Cores实现更快的计算。AWS p3实例使用带有Tensor内核的NVIDIA Tesla V100 GPU。采用NVIDIA的apex进行混合精度训练非常简单，只需要修改部分代码：

def train(gpu, args):
    ############################################################
    rank = args.nr * args.gpus + gpu
    dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    ############################################################
    torch.manual_seed(0)
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    batch_size = 100
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    ############################################################
    # Wrap the model
    model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
    model = DDP(model)
    ############################################################

    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(),
                                               download=True)
    ############################################################
    # train_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_dataset, num_replicas=args.world_size,
    #                                                                 rank=rank)
    # train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=False,
    #                                            num_workers=0, pin_memory=True, sampler=train_sampler)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
                                               num_workers=0, pin_memory=True)
    ############################################################

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            ############################################################
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
            ############################################################
            optimizer.step()
            if (i + 1) % 100 == 0 and gpu == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                         loss.item()))
    if gpu == 0:
        print("Training complete in: " + str(datetime.now() - start))

其实就两处变化:

• 首先是采用amp.initialize来包装model和optimizer以支持混合精度训练，其中opt_level指的是优化级别，如果为O0（使用all floats）或者O3（使用half-precision throughout）不是真正的混合精度，但是可以用来确定模型效果和速度的baseline，而O1和O2是混合精度的两种设置，可以选择某个进行混合精度训练，其详细信息可以在Apex文档中找到。
• 另外一处是在进行根据梯度更新参数前，要先通过amp.scale_loss对梯度进行scale以防止梯度下溢（underflowing）。此外，你还可以用apex.parallel.DistributedDataParallel替换nn.DistributedDataParallel。

是的，所有这些代码中的第一个字符是大写字母“O”，而第二个字符是数字。是的，如果用零来代替，会得到一条令人费解的错误消息。

apex.parallel.distributedDataParallel是nn.distributedDataParallear的一个替换。不再需要指定GPU，因为Apex每个process只允许一个GPU。它还假设脚本在将模型移动到GPU之前调用torch.cuda.set_device(local_rank)。

混合精度训练要求对损失进行缩放，以防止梯度下溢。Apex会自动做到这一点。

此脚本的运行方式与分布式训练脚本相同。

python without_multiprocessing.py -n 1 -g 4 -nr 0

另外，新版本的PyTorch已经内置混合精度训练，具体见AUTOMATIC MIXED PRECISION PACKAGE - TORCH.AMP添加链接描述。而且PyTorch官方的分布式实现现在已经比较完善，而且性能和效果都不错，可以替代的方案是horovod，不仅支持PyTorch还支持TensorFlow和MXNet框架，实现起来也是比较容易的，速度方面也不错。

参考资料

1. Distributed data parallel training in Pytorch
2. 在PyTorch中使用DistributedDataParallel进行多GPU分布式模型训练
3. PyTorch分布式训练简明教程(2022更新版)
4. 分布式训练框架Horovod