PyTorch Training Acceleration

神经网络的训练往往是整个流程中最花时间的部分，尤其是在一个大型数据集上将随机初始化的神经网络训练到收敛。通常来说，使用更好的硬件平台可以让训练过程变得更快，例如使用更多更好的 GPU，使用更大更快的 SSD，以及在分布式训练的时候使用 InfiniBand 替换以太网等等。除了硬件平台之外，软件与算法也同样关键，好的软件平台应该尽可能的充分利用现有硬件平台的资源。本文从一个 baseline 出发，通过各种方法逐步对训练速度进行优化，我们选取的 baseline 是在 ImageNet 数据集上训练 ResNet18，总共训练 90 epochs，初始计算平台为如表 1 所示：

1、分析训练过程中的时间瓶颈

我们在这里使用了NVIDIA Tools Extension Library (NVTX)来测量训练过程中各个部分的时间开销，使用方式如下，只需要在训练代码中插入几行代码，然后使用 nsys profile python3 main.py 启动训练即可，最后会生成一个.qdrp 文件，我们在这里分别统计了一个 batch 中数据加载，前向传播，反向传播，以及梯度下降所花费的时间，为了节省时间我们对 300 个 batch 的训练过程进行了统计。

import torch.cuda.nvtx as nvtx
nvtx.range_push("Batch 0")
nvtx.range_push("Load Data")
for i, (input_data, target) in enumerate(train_loader):
    input_data = input_data.cuda(non_blocking=True)
    target = target.cuda(non_blocking=True)
    nvtx.range_pop(); nvtx.range_push("Forward")
    output = model(input_data)
    nvtx.range_pop(); nvtx.range_push("Calculate Loss/Sync")
    loss = criterion(output, target)
    prec1, prec5 = accuracy(output, target, topk=(1, 5))
    optimizer.zero_grad()
    nvtx.range_pop(); nvtx.range_push("Backward")
    loss.backward()
    nvtx.range_pop(); nvtx.range_push("SGD")
    optimizer.step()
    nvtx.range_pop(); nvtx.range_pop()
    nvtx.range_push("Batch " + str(i+1)); nvtx.range_push("Load Data")
nvtx.range_pop()
nvtx.range_pop()

在使用 baseline 代码进行训练的时候，我们使用 Nvidia Nsight Systems 对生成的.qdrp 文件进行可视化，结果如图 1 所示，可以看到训练 300 个 batch 所用时间为 336.524s，但值得注意的是，CUDA 内核实际运行的时间仅仅占用了总时间的约四分之一，而四分之三的时间都在进行数据加载，在这个过程中 CUDA 内核和 CPU 使用率几乎为 0，这成了我们使用 baseline 代码进行训练的过程中的性能瓶颈。

如果进一步查看每个 batch 里面的训练时间（例如 batch 122），可以看到将数据从 Host 内存传到 Device 显存所用时间为 13.6ms（数据大小为 224*224*3*256*4=154,140,672bytes，带宽 11.3GB/s）；forward 中调用 CUDA 内核所用时间为 15.4ms，CUDA 内核实际运行时间为 67.3ms；backward 中调用 CUDA 内核所用时间为 19.9ms，CUDA 内核实际运行时间为 164.7ms；SGD 中调用 CUDA 内核所用时间为 8.0ms，CUDA 内核实际运行时间为 1.8ms。

2、加速数据读取

从上一节的实验结果可以看出来，整个训练过程中四分之三的时间都花在了数据读取上，而这个过程中的 CPU/GPU 使用率都几乎为 0，造成了硬件资源上的浪费，因此我们首先对 I/O 部分进行优化。baseline 训练代码数据读取慢主要是因为 PyTorch 在训练的时候读取的是一张张.jpeg 等格式的原始图片，不仅需要解码，而且读取过程并不连续，cache miss 较多从而 I/O 时间会很长；而 Tenorflow 等框架会有自己的.tfrecord 格式，类似的还有 hdf5，lmdb 等，这些格式将整个数据集存储在一个大的二进制文件中，从而在读取数据的时候可以进行连续读取，效率会高很多。

在优化 I/O 的时候需要根据硬件环境的不同选取合理的方案，若计算平台的内存不足以容纳整个数据集可以使用 lmdb（hdf5 需要将整个文件加载进内存）。我们做了两组实验进行对比，分别是随机读取(shuffle 训练数据)和连续读取(不 shuffle 训练数据)的情况，结果如表 2 所示。如果不对数据集做 shuffle，那么用 lmdb 读取的速度会非常快，profile 的结果如图 2 所示，在训练了约 50 个 batch 之后，CUDA 内核几乎一直处于忙碌状态，这得益于连续读取所减小的 cache miss 率；但是在 shuffle 的情况下反而会变得更慢。

若内存足够大可以直接简单粗暴的将整个数据集直接放到内存中，具体则可以通过挂载 tmpfs 文件系统，然后将整个数据集放到该文件夹中来实现。如果没有挂载的权限也可以将数据集放到/dev/shm 目录下，同样也是 tmpfs 文件系统，通常 160G 的大小足够容纳整个 ImageNet 数据集。

# 有挂载权限
mkdir -p /userhome/memory_data/imagenet
mount -t tmpfs -o size=160G tmpfs /userhome/memory_data/imagenet
root_dir="/userhome/memory_data/imagenet"
# 无挂载权限
mkdir -p /dev/shm/imagenet
root_dir="/dev/shm/imagenet"
mkdir -p "${root_dir}/train"
mkdir -p "${root_dir}/val"
tar -xvf /userhome/data/ILSVRC2012_img_train.tar -C "${root_dir}/train"
tar -xvf /userhome/data/ILSVRC2012_img_val.tar -C "${root_dir}/val"

将数据集放到 tmpfs 文件系统后，我们重新用 baseline 代码进行了训练，结果如图 3 所示，CUDA 内核几乎一直处于满负荷的状态，整体的训练时间也从 baseline 的 336.524s 缩短到 87.989s（仅为 baseline 的 26%），提速效果非常明显。

然后同样的，我们进一步统计每个 batch 里面的训练时间，将数据从 Host 内存传到 Device 显存所用时间为 15.4ms(1.8ms ↑ )；forward 中调用 CUDA 内核所用时间为 7.9ms(7.5ms ↓ )，CUDA 内核实际运行时间为 67.5ms(0.2ms ↑ )；backward 中调用 CUDA 内核所用时间为 16.6ms(3.3ms ↑ )，CUDA 内核实际运行时间为 164.3ms(0.4ms ↑ )；SGD 中调用 CUDA 内核所用时间为 5.5ms(2.5ms ↑ )，CUDA 内核实际运行时间为 1.9ms(0.1ms ↑ )。可以看出计算上所花的时间与 baseline 几乎没有区别。

3、混合精度训练

从上一节可以看出，在优化数据读取之后 CUDA 内核可以处于满负荷的状态，因此要进一步加速训练只能从计算速度（主要是 forward 与 backward）本身来入手，但 PyTorch 使用了 CUDNN 的 kernel 进行计算，因此我们没法优化 kernel 本身，只能通过将部分 FP32 的计算转换为 FP16 来进行提速，即混合精度训练\cite{micikevicius2017mixed}。

混合精度训练一种比较简便的方案是使用Apex (A PyTorch Extension)，仅需要在原代码中添加几行代码即可，并且还可以选择不同的 FP16 训练方案。

from apex import amp, optimizers
# Allow Amp to perform casts as required by the opt_level
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
...
# loss.backward() becomes:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
...

其中的 opt_level 有多种选择，O0 是全部使用 FP32 进行训练，O1 在部分操作上使用 FP16 进行计算，O2 在绝大部分操作上使用 FP16 进行计算，O3 则是全部使用 FP16 进行计算。在使用 O1 与 O2 时，FP16 的算子（卷积，全连接等）使用 FP16 进行前向传播与反向传播，但维护了一个 FP32 的权重用于梯度下降时的更新（图 4），FP32 的算子（BN 等）则使用了 FP32 进行前向传播与反向传播，注意这两种算子都使用了 FP32 进行梯度下降。使用 FP16 进行训练有一定的风险会训练发散，具体选择哪一种方案需要根据实际情况在提速与防止训练发散之间进行权衡。我们在上一节中速度最快方案的基础上进一步对 FP16 训练做了实验，结果如表 3 所示，可以看出利用 FP16 进行训练的时候 forward 与 backward 的速度相比 FP32 有明显提升，总体训练时间也有较为明显的下降。

到本节为止，我们一直使用单个 GPU 进行训练，神经网络的训练速度被各种方法一步步加快，首先从 baseline 的 336.524s 开始（图 1），在优化 I/O 之后加速为 87.989s（图 3），之后再通过使用混合进度训练在不减小精度的情况下加速到 70.1s。

4、单机多卡并行训练

本节将在优化数据读取以及混合进度训练的基础上，使用单机多卡并行训练进一步加速，目前仅关注于数据并行，暂不考虑模型并行甚至算子并行的方法。

在 PyTorch 上实现单机多卡训练通常有两种方式，一种是使用 nn.DataParallel，另一种是使用 nn.parallel.DistributedDataParallel。前者只使用了单个进程；而后者使用了多个进程并行训练，同样也适用于多机多卡，除了 nn.parallel.DistributedDataParallel 之外还有一些第三方库可以用于 PyTorch 的多进程并行训练，例如APEX，Horovod等。

使用 nn.DataParallel 是最为简单的一种方式，仅用一行代码就可以实现：

model = torch.nn.DataParallel(model)

在使用 nn.DataParallel 进行并行训练时，首先会将整个 batch 的数据加载到一张主 GPU 上；然后再通过 PtoP 的拷贝将每份数据(BS/GPU_NUM)从主 GPU 依次拷贝到其他 GPU 上；之后再通过NVIDIA Collective Communications Library (NCCL)的 Broadcast(图 5)将神经网络的参数从主 GPU 广播给其他 GPU；然后每张 GPU 上各自进行前向传播与反向传播；最后再通过 Reduce(图 6)将其他 GPU 上的梯度归约到主 GPU 上，主 GPU 进行梯度下降得到优化后的神经网络。

我们首先在配置[1]的基础上进行实验，并使用相同的 batch size（例如总 batch size 为 256 时，使用双卡的时候每张卡 128，四卡的时候每张卡 64），结果如以表 4 所示（实验中迭代了 500 个 batch，中间 300 个 batch 的总时间作为表 4 中的总时间；DtoD 的时间包含两部分，分别是主 GPU 将每份数据 PtoP 传给其他 GPU 的时间，以及主 GPU Broadcast 神经网络参数和各个 GPU 将梯度 Reduce 到主 GPU 的时间总和）。其中 BS=512 时[7]的 profile 结果如图 7 所示

从上面的表 4 和图 7 中大概可以观察到以下几个现象：

1. 使用多 GPU 可以有效减小 forward 与 backward 的时间。
2. GPU 数量增加到一定程度时，GPU 运算完成之后会出现等待数据加载的情况。
3. 随着 GPU 数量的增加，GPU 之间通讯所花的时间占比变大，GPU 数量对于整体速度的提升越发有限。

对于其中的第二条，数据加载的时间主要由两部分构成，一方面是 I/O 的时间，这方面我们已经通过将数据集放进内存中解决掉了；另一方面则是数据预处理的时间，这方面我们可以通过增加 cpu 的线程数以及数据读取的线程数来解决，在 docker 容器中训练的时候可以尝试前者的方法，通过修改 docker run 时–cpus 的参数来实现，后者则可以通过在创建 torch.utils.data.DataLoader 对象时适当增加 num_workers 的大小来实现。我们在[7]的配置上进一步改进，在 BS=512 时的实验结果如表 5 所示，可以看出如果训练过程中 GPU 还存在等待数据的空闲，适当增加 CPU 线程数与数据加载的线程数可以有效减小总的时间（主要是通过减小 GPU 等待数据的空闲来实现的），但过度增大 CPU 线程数会造成一定程度的浪费，因为到了一定程度时 GPU 等待数据的空闲已经很小了，除此之外过度增加线程数也会由于线程本身的开销过大而变得更慢（例如下表的[10]与[11]）。

而对于第三条，我们注意到 nn.DataParallel 并行训练时会先将所有数据从内存传到一张主 GPU 上，然后再将数据从主 GPU 依次 PtoP 传到其他 GPU 上，该过程显然不如直接将每一份数据从内存传到对应的 GPU 上更加高效，这可以节省掉将每份数据从主 GPU 依次 PtoP 传到其他 GPU 上的时间（表 5 中 DtoD 耗时的前一部分，在配置[7]的情况下这部分时间占比接近 20%）。除此之外，nn.DataParallel 在每一次迭代时都会先将模型的参数从主 GPU 上 Broadcast 到其他 GPU 上，最后再将梯度 Reduce 到主 GPU 上。如果利用 All-Reduce 将归约之后的梯度传到每个 GPU 上（图 8），那么每个 GPU 上可以各自进行梯度下降，这样便省去了 nn.DataParallel 最开始 Broadcast 模型参数的时间开销（大概占比 2%）。

基于多进程并行的 nn.parallel.DistributedDataParallel 实现了这样的方法，具体如何使用可以参考官方的 example。我们在配置[10]的基础上进行了实验，结果如表 6 所示，使用 DistributedDataParallel 时的 DtoD 时间由于不同 GPU 之间的同步时间具有较大波动，因此没有进行统计，可以看出使用 DistributedDataParallel 可以减小约 1/3 的时间。从 profile 结果中（图 9）可以看出配置[12]的 CUDA 内核使用率远高于配置[7]，这一方面是由于我们使用了更多的 CPU 线程和数据加载线程，另一方面是由于使用了基于多进程实现的并行训练，相比多线程的实现可以突破 Python 的 GIL(Global Interpreter Lock)，从而进一步提升效率。

具体来看，图 10 与图 11 分别展示了使用 DataParallel 和 DistributedDataParallel 时训练某个 batch 的具体时间开销。在使用 DistributedDataParallel 时，图 9 中 HtoD 的数据大小从总 BS 减小到 BS/GPU_NUM，这部分时间理论上可以减小为 DataParallel 的 1/GPU_NUM（不同 GPU 的传输速度略有波动）；除此之外，DtoD 的时间可以完全节省掉；最后 Broadcast 和 Reduce 的时间被 All-Reduce 所替代（由于 All-Reduce 需要所有 GPU 的同步，耗时波动较大）。这两种方法在具体计算上的耗时（Forward/Backward/SGD）没有明显区别。

除了 DistributedDataParallel 之外，我们之前提到还有 APEX 和 Horovod 等第三方库实现了同样的多进程并行，我们同样也使用这两种第三方库进行了对比实验，结果如表 7 所示。可以看出三种实现方式的速度差别不大，值得注意的是 APEX 可以对 BN 进行同步运算，也就是在 BN 层计算 batch 的均值/标准差的时候对所有 GPU 上的数据进行同步，这样计算得到的均值/标准差是由整个 batch 的数据统计而来，而非单张 GPU 上的一份数据得到，通过这样的方式可以使训练更加稳定。但这样的方式需要每次计算均值/方差时所有 GPU 进行同步（sync），因此会带来一些额外的时间开销，但从 profile 的结果来看，GPU 之间计算不同步的主要原因是最初的 HtoD 时间不同，在经过第一个 BN 层同步之后，后续的 BN 层所需要的同步时间是很短的（前提是所有 GPU 的算力基本相同）。

最后再将上节的 FP16 混合精度训练与本节的多卡训练方法相结合，通过将这两种正交的加速方式相结合，我们可以将训练速度提升到极限。具体来说，从表 8 可以看出使用 DistributedDataParallel 是并行效率最高的，因此可以将其与 FP16(O2)相结合，实验结果表明该方法可以进一步加速到 38.8s，实验的 profile 结果如图 12 所示。

5、多机多卡并行训练

前面的实验都是在单个机器上进行的，但在实际情况中为了加速还需要进行多机多卡的分布式训练，例如双机四卡的情况下可以用以下的代码进行训练：

# 假设在2台机器上运行，每台可用卡数是4
#    机器1：
python -m torch.distributed.launch --nnodes=2 --node_rank=0 --nproc_per_node 4 \
  --master_adderss $my_address --master_port $my_port main.py
#    机器2：
python -m torch.distributed.launch --nnodes=2 --node_rank=1 --nproc_per_node 4 \
  --master_adderss $my_address --master_port $my_port main.py

我们分别在不同硬件配置下进行了训练速度的测试，主要区别在于是否使用 InfiniBand，实验结果如表 9 所示，由于在训练过程中不同机器上的数据需要同步，因此在训练耗时中有一部分来源于机器之间的通讯时间，使用 InfiniBand 可以达到比以太网更低的时延（但以太网通用性更好，InfiniBand 相对比较适用于并行计算的场景）。

6、总结

本文中通过调整配置一步步对训练进行优化，尽可能充分利用 gpu 的算力。首先在使用 baseline 代码，单个 GPU，BS=512 的情况下训练 300batch 的速度为 161.0s（配置【0】）；然后直接使用 DataParallel 在 4 张 GPU 上进行训练，可以加速为 81.7s（配置【7】）；之后再增加一定的 CPU 线程数与数据加载线程数，通过提升数据加载速度加速为 60.2s（配置【10】）；除此之外再将 DataParallel 替换为 DistributedDataParallel，通过减小 GPU 之间通讯时间加速为 41.6s（配置【12】）；然后再与混合精度训练相结合，加速到 38.8s（配置【16】）；最后再使用两台这样的节点通过 InfiniBand 通讯进行并行训练，加速到 22.1s（配置【18】）。