gpu并发和数据量的关系

文章目录

• 1 编译器和解释器

• 1.1 符号式编程
• 1.2 混合式编程
• 1.3 Sequential的混合式编程

• 1.3.1 通过混合式编程加速

• 1.4 小结

• 2 异步计算

• 2.1 通过后端异步处理

• 3 自动并行

• 3.1 基于GPU的并行计算

• torch.cuda.synchronize() 函数记录：

• 3.2 并行计算与通信
• 3.3 小结

• 4 硬件

• 4.1 计算机
• 4.2 内存
• 4.3 存储器
• 4.8 总结

• 5 多GPU训练

• 5.1 问题拆分
• 5.2 数据并行性
• 5.3 简单网络
• 5.4 数据同步
• 5.5 数据分发

• nn.parallel 中的函数记录

• 5.6 训练

• 6 多GPU的简洁实现

• 6.1 简单网络
• 6.2 网络初始化
• 6.3 训练
• 6.4 小结

• 7 参数服务器

• 7.1 数据并行训练
• 7.2 环同步（Ring Synchronization）
• 7.3. 多机训练

1 编译器和解释器

def add(a, b):
    return a + b

def fancy_func(a, b, c, d):
    e = add(a, b)
    f = add(c, d)
    g = add(e, f)
    return g

print(fancy_func(1, 2, 3, 4))

Python是一种解释型语言（interpreted language）。因此，当对上面的 fancy_func 函数求值时，它按顺序执行函数体的操作。 图12.1.1 说明了数据流。

尽管命令式编程很方便，但可能效率不高。一方面原因，Python 会单独执行这三个函数的调用，而没有考虑 add 函数在 fancy_func 中被重复调用。如果在一个 GPU（甚至多个 GPU）上执行这些命令，那么 Python 解释器产生的开销可能会非常大。此外，它需要保存 e 和 f 的变量值，直到 fancy_func 中的所有语句都执行完毕。这是因为程序不知道在执行语句 e = add(a, b) 和 f = add(c, d) 之后，其他部分是否会使用变量 e 和 f。

1.1 符号式编程

考虑另一种选择符号式编程（symbolic programming），即代码通常只在完全定义了过程之后才执行计算。
一般包括以下步骤：
1.定义计算流程。
2.将流程编译成可执行的程序。
3.给定输入，调用编译好的程序执行。

编译器在将其转换为机器指令之前可以看到完整的代码，所以这种优化是可以实现的。例如，只要某个变量不再需要，编译器就可以释放内存（或者从不分配内存），或者将代码转换为一个完全等价的片段。

def add_():
    return '''
def add(a, b):
    return a + b
'''

def fancy_func_():
    return '''
def fancy_func(a, b, c, d):
    e = add(a, b)
    f = add(c, d)
    g = add(e, f)
    return g
'''

def evoke_():
    return add_() + fancy_func_() + 'print(fancy_func(1, 2, 3, 4))'

prog = evoke_()
print(prog)
y = compile(prog, '', 'exec')
exec(y)

1.2 混合式编程

1.3 Sequential的混合式编程

import torch
from torch import nn
from d2l import torch as d2l


# 生产网络的工厂模式
def get_net():
    net = nn.Sequential(nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 2))
    return net

x = torch.randn(size=(1, 512))
net = get_net()
print(net(x))

通过使用 torch.jit.script 函数来转换模型，我们就有能力编译和优化多层感知机中的计算，而模型的计算结果保持不变。

net = torch.jit.script(net)
print(net(x))

我们编写与之前相同的代码，再使用 torch.jit.script 简单地转换模型，当完成这些任务后，网络就将得到优化（我们将在下面对性能进行基准测试）。

1.3.1 通过混合式编程加速

为了证明通过编译获得了性能改进，我们比较了混合编程前后执行 net(x) 所需的时间。先定义一个度量时间的函数。

#@save
class Benchmark:
    def __init__(self, description='Done'):
        self.description = description

    def __enter__(self):
        self.timer = d2l.Timer()
        return self

    def __exit__(self, *args):
        print(f'{self.description}: {self.timer.stop():.4f} sec')

net = get_net()
with Benchmark('无torchscript'):
    for i in range(1000): net(x)

net = torch.jit.script(net)
with Benchmark('有torchscript'):
    for i in range(1000): net(x)
############
无torchscript: 0.1698 sec
有torchscript: 0.1050 sec

如以上结果所示，在 nn.Sequential 的实例被函数 torch.jit.script 脚本化后，通过使用符号式编程提高了计算性能。

1.4 小结

• 命令式编程使得新模型的设计变得容易，因为可以依据控制流编写代码，并拥有相对成熟的 Python 软件生态。
• 符号式编程要求我们先定义并且编译程序，然后再执行程序，其好处是提高了计算性能。

2 异步计算

今天的计算机是高度并行的系统，由多个 CPU 核、多个 GPU、多个处理单元组成。通常每个 CPU 核有多个线程，每个设备通常有多个 GPU ，每个 GPU 有多个处理单元。 总之，我们可以同时处理许多不同的事情，并且通常是在不同的设备上。
不幸的是，Python 并不善于编写并行和异步代码，至少在没有额外帮助的情况下不是好选择。归根结底，Python 是单线程的，将来也是不太可能改变的。
**对于 PyTorch 来说 GPU 操作在默认情况下是异步的。**当你调用一个使用 GPU 的函数时，操作会排队到特定的设备上，但不一定要等到以后才执行。这允许我们并行执行更多的计算，包括在 CPU 或其他 GPU 上的操作。

2.1 通过后端异步处理

我们要生成一个随机矩阵并将其相乘。让我们在 NumPy 和 PyTorch 张量中都这样做，看看它们的区别。请注意，PyTorch 的 tensor 是在 GPU 上定义的。

# GPU 计算热身
device = d2l.try_gpu()
a = torch.randn(size=(1000, 1000), device=device)
b = torch.mm(a, a)

with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)

with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
############
numpy: 0.4015 sec
torch: 0.0100 sec

通过 PyTorch 的基准输出比较快了几个数量级。NumPy 点积是在 CPU 上执行的，而 PyTorch 矩阵乘法是在 GPU 上执行的，后者的速度要快得多。

with d2l.Benchmark():
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
    torch.cuda.synchronize(device)  # 同步
##############
Done: 0.0207 sec

广义上说，PyTorch有一个用于与用户直接交互的前端（例如通过Python），还有一个由系统用来执行计算的后端。 由前端语言发出的操作被传递到后端执行。后端管理自己的线程，这些线程不断收集和执行排队的任务。

import torch
from d2l import torch as d2l

device = d2l.try_gpu()
x = torch.ones((1, 2), device=device)
y = torch.ones((1, 2), device=device)
z = x * y + 2
print(z)
#################
tensor([[3., 3.]], device='cuda:0')

每当 Python前端线程执行前三条语句中的一条语句时，它只是将任务返回到后端队列。当最后一个语句的结果需要被打印出来时，Python前端线程将等待C++ 后端线程完成变量 z 的结果计算。这种设计的一个好处是 Python 前端线程不需要执行实际的计算。因此，不管 Python的性能如何，对程序的整体性能几乎没有影响。

3 自动并行

深度学习框架（例如，MxNet 和PyTorch）会在后端自动构建计算图。利用计算图，系统可以了解所有依赖关系，并且可以选择性地并行执行多个不相互依赖的任务以提高速度。

通常情况下单个操作符将使用所有 CPU 或单个 GPU 上的所有计算资源。借助自动并行化框架的便利性，我们可以依靠几行 Python 代码实现相同的目标。更广泛地考虑，我们对自动并行计算的讨论主要集中在使用 CPU 和 GPU 的并行计算上，以及计算和通信的并行化内容。

请注意，我们至少需要两个GPU来运行本节中的实验。

3.1 基于GPU的并行计算

下面的 run 函数将执行 50 次“矩阵－矩阵”乘法时需要使用的数据分配到两个变量（x_gpu1和x_gpu2）中，这两个变量分别位于我们选择的不同设备上。

import torch
from d2l import torch as d2l

devices = d2l.try_all_gpus()
def run(x):
    return [x.sum(x) for _ in range(50)]

x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])

通过在测量之前预热设备（对设备执行一次传递）来确保缓存的作用不影响最终的结果。

torch.cuda.synchronize() 函数记录：

torch.cuda.synchronize() 函数将会等待一个 CUDA 设备上的所有流中的所有核心的计算完成。函数接受一个device 参数，代表是哪个设备需要同步。 如果device参数是 None（默认值），它将使用current_device() 找出的当前设备。

run(x_gpu1)
run(x_gpu2)  # 预热设备
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])

with d2l.Benchmark('GPU1 time'):
    run(x_gpu1)
    torch.cuda.synchronize(devices[0])

with d2l.Benchmark('GPU2 time'):
    run(x_gpu2)
    torch.cuda.synchronize(devices[1])

如果我们删除两个任务之间的 synchronize 语句，系统就可以在两个设备上自动实现并行计算。

with d2l.Benchmark('GPU1 & GPU2'):
    run(x_gpu1)
    run(x_gpu2)
    torch.cuda.synchronize()

3.2 并行计算与通信

在许多情况下，我们需要在不同的设备之间移动数据，比如在 CPU 和 GPU 之间，或者在不同的 GPU 之间。当我们打算执行分布式优化时，就需要移动数据来聚合多个加速卡上的梯度。让我们通过在GPU上计算，然后将结果复制回CPU来模拟这个过程。

def copy_to_cpu(x, non_blocking=False):
    return [y.to('cpu', non_blocking=non_blocking) for y in x]

with d2l.Benchmark('在GPU1上运行'):
    y = run(x_gpu1)
    torch.cuda.synchronize()

with d2l.Benchmark('复制到CPU'):
    y_cpu = copy_to_cpu(y)
    torch.cuda.synchronize()
##############
在GPU1上运行: 0.4214 sec
复制到CPU: 2.0170 sec

这种方式效率不高。注意到当列表中的其余部分还在计算时，我们可能就已经开始将 y 的部分复制到 CPU 了。例如，当我们计算一个小批量的（反传）梯度时。某些参数的梯度将比其他参数的梯度更早可用。因此，在 GPU 仍在运行时就开始使用 PCI-Express 总线带宽来移动数据对我们是有利的。在 PyTorch 中，to()和 copy_() 等函数都允许显式的 non_blocking 参数，这允许在不需要同步时调用方可以绕过同步。设置 non_blocking=True 让我们模拟这个场景。

with d2l.Benchmark('在GPU1上运行并复制到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, True)
    torch.cuda.synchronize()
##############
在GPU1上运行并复制到CPU: 1.4977 sec

两个操作所需的总时间少于它们各部分操作所需时间的总和。请注意，与并行计算的区别是通信操作使用的资源：CPU 和 GPU 之间的总线。

3.3 小结

• 现代系统拥有多种设备，如多个 GPU 和多个 CPU，还可以并行地、异步地使用它们。
• 现代系统还拥有各种通信资源，如PCI Express、存储（通常是固态硬盘或网络存储）和网络带宽，为了达到最高效率可以并行使用它们。
• 后端可以通过自动化地并行计算和通信来提高性能。

4 硬件

4.1 计算机

大多数深度学习研究者和实践者都可以使用一台具有相当数量的内存、计算资源、某种形式的加速器（如一个或者多个 GPU）的计算机。计算机由以下关键部件组成：

• 一个处理器（也被称为 CPU），它除了能够运行操作系统和许多其他功能之外，还能够执行我们给它的程序，通常由 8 个或更多个核心组成。
• 内存（RAM） 用于存储和检索计算结果，如权重向量和激活参数，以及训练数据。
• 一个或多个以太网连接，速度从 1 GB/s 到 100 GB/s 不等。在高端服务器上可能用到更高级的互连。
• 高速扩展总线（PCIe） 用于系统连接一个或多个 GPU 。服务器最多有 8 个加速卡，通常以更高级的拓扑方式连接，而桌面系统则有 1 个或 2 个加速卡，具体取决于用户的预算和电源负载的大小。
• 持久性存储设备，如磁盘驱动器、固态驱动器，在许多情况下使用高速扩展总线连接。它为系统需要的训练数据和中间检查点需要的存储提供了足够的传输速度。
  
  高速扩展总线由直接连接到CPU的多个通道组成，将 CPU 与大多数组件（网络、GPU 和存储）连接在一起。

4.2 内存

最基本的内存主要用于存储需要随时访问的数据。目前，CPU 的内存通常为 DDR4 类型，每个模块提供 20-25 Gb/s 的带宽。每个模块都有一条 64 位宽的总线。通常使用成对的内存模块来允许多个通道。CPU 有 2 到 4 个内存通道，也就是说，它们内存带宽的峰值在 40 GB/s 到 100 GB/s 之间。一般每个通道有两个物理存储体（bank）。
虽然这些数字令人印象深刻，但实际上它们只能说明了一部分故事。当我们想要从内存中读取一部分内容时，我们需要先告诉内存模块在哪里可以找到信息。也就是说，我们需要先将 地址（address）发送到 RAM 。然后我们可以选择只读取一条 64 位记录还是一长串记录。后者称为 突发读取（burst read） 。概括地说，向内存发送地址并设置传输大约需要 100 ns（细节取决于所用内存芯片的特定定时系数），每个后续传输只需要 0.2 ns。总之，第一次读取的成本是后续读取的500倍！请注意，我们每秒最多可以执行一千万次随机读取。这说明应该尽可能地避免随机内存访问，而是使用突发模式读取和写入。
当考虑到我们拥有多个物理存储体时，事情就更加复杂了。每个存储体大部分时候都可以独立地读取内存。这意味着两件事。一方面，如果随机读操作均匀分布在内存中，那么有效的随机读操作次数将高达4倍。这也意味着执行随机读取仍然不是一个好主意，因为突发读取的速度也快了4倍。另一方面，由于内存对齐是 64 位边界，因此最好将任何数据结构与相同的边界对齐。当设置了适当的标志时，编译器基本上就是自动化地执行对齐操作。
因为 GPU 的处理单元比 CPU 多得多，因此它对内存带宽的需要也更高。解决这种问题大体上有两种选择。首要方法是使内存总线变得更宽。例如：NVIDIA 的 RTX 2080 Ti 有一条 352 位宽的总线，这样就可以同时传输更多的信息。再有方法就是在 GPU 中使用特定的高性能内存。一种选择是如 NVIDIA 的消费级设备 RTX 和 Titan 系列中通常使用 GDDR6 芯片，其总带宽超过 500 GB/s。另一种选择是使用 HBM（高带宽存储器）模块。这些模块使用截然不同的接口在专用硅片上与 GPU 直接连在一起。这导致其非常昂贵，通常仅限于在高端服务器的芯片上使用，如 NVIDIA Volta V100 系列的加速卡。
GPU内存的带宽要求甚至更高，因为它们的处理单元比CPU多得多。总的来说，解决这些问题有两种选择。首先是使内存总线变得更宽。 例如，NVIDIA的RTX 2080 Ti有一条352位宽的总线。这样就可以同时传输更多的信息。其次，GPU使用特定的高性能内存。消费级设备，如NVIDIA的RTX和Titan系列，通常使用GDDR6芯片，总带宽超过500GB/s。另一种选择是使用HBM(高带宽存储器)模块。它们使用截然不同的接口，直接与专用硅片上的GPU连接。这使得它们非常昂贵，通常仅限于高端服务器芯片，如NVIDIA Volta V100系列加速卡。毫不意外的是 GPU 的内存通常比 CPU 的内存小得多，因为前者的成本更高。就目的而言，它们的性能与特征大体上是相似的，只是 GPU 的速度更快。就本书而言，我们完全可以忽略细节，因为这些技术只在调整 GPU 核心以获得高吞吐量时才起作用。

4.3 存储器

我们看到 RAM 的一些关键特性是 带宽（bandwidth）和 延迟（latency）。存储设备也是如此，只是不同设备之间的特性差异可能更大。

4.8 总结

这段和现在学习的内容关联性不是特别强，日后会慢慢学习填坑。

5 多GPU训练

到目前为止，我们讨论了如何在 CPU 和 GPU 上高效地训练模型，同时在中展示了深度学习框架如何在 CPU 和 GPU之间自动地并行化计算和通信，展示了如何使用 nvidia-smi 命令列出计算机上所有可用的 GPU。但是我们没有讨论如何真正实现深度学习训练的并行化。 是否有一种方法，以某种方式分割数据到多个设备上，并使其能够正常工作呢？

5.1 问题拆分

我们希望以一种方式对训练进行拆分，为实现良好的加速比，还能同时受益于简单且可重复的设计选择。毕竟，多个 GPU 同时增加了内存和计算能力。 简而言之，对于需要分类的小批量训练数据，我们有以下选择。
第一种方法，在多个 GPU 之间拆分网络。（网络并行） 也就是说，每个 GPU 将流入特定层的数据作为输入，跨多个后续层对数据进行处理，然后将数据发送到下一个 GPU。 与单个 GPU 所能处理的数据相比，我们可以用更大的网络处理数据。 此外，每个 GPU 占用的 显存 （memory footprint）可以得到很好的控制，虽然它只是整个网络显存的一小部分。
然而，GPU的接口之间需要的密集同步可能是很难办的，特别是层之间计算的工作负载不能正确匹配的时候， 还有层之间的接口需要大量的数据传输的时候（例如：激活值和梯度，数据量可能会超出GPU总线的带宽）。

第二种方法，拆分层内的工作。（分层并行） 例如，将问题分散到 4 个 GPU，每个 GPU 生成 16 个通道的数据，而不是在单个 GPU 上计算 64 个通道。 对于全连接的层，同样可以拆分输出单元的数量。 下图描述了这种设计，其策略用于处理显存非常小（当时为2GB）的GPU。 当通道或单元的数量不太小时，使计算性能有良好的提升。 此外，由于可用的显存呈线性扩展，多个 GPU 能够处理不断变大的网络。

然而，我们需要大量的同步或屏障操作(barrier operations)，因为每一层都依赖于所有其他层的结果。 此外，需要传输的数据量也可能比跨 GPU 拆分层时还要大。 因此，基于带宽的成本和复杂性，我们同样不推荐这种方法。

最后一种方法，跨多个 GPU 对数据进行拆分。（数据并行） 这种方式下，所有 GPU 尽管有不同的观测结果，但是执行着相同类型的工作。 在完成每个小批量数据的训练之后，梯度在 GPU 上聚合。 这种方法最简单，并可以应用于任何情况，同步只需要在每个小批量数据处理之后进行。 也就是说，当其他梯度参数仍在计算时，完成计算的梯度参数就可以开始交换。 而且，GPU 的数量越多，小批量包含的数据量就越大，从而就能提高训练效率。 但是，添加更多的 GPU 并不能让我们训练更大的模型。

总体而言，只要 GPU 的显存足够大，数据并行是最方便的。 (分布式训练分区)

5.2 数据并行性

假设一台机器有 k 个 GPU。 给定需要训练的模型，虽然每个 GPU 上的参数值都是相同且同步的，但是每个 GPU 都将独立地维护一组完整的模型参数。

一般来说， k 个GPU并行训练过程如下：

• 在任何一次训练迭代中，给定的随机的小批量样本都将被分成 k 个部分，并均匀地分配到 GPU 上。
• 每个 GPU 根据分配给它的小批量子集，计算模型参数的损失和梯度。
• 将 k 个 GPU 中的局部梯度聚合，以获得当前小批量的随机梯度。
• 聚合梯度被重新分发到每个 GPU 中。
• 每个 GPU 使用这个小批量随机梯度，来更新它所维护的完整的模型参数集。

在实践中请注意，当在 k 个 GPU 上训练时，需要扩大小批量的大小为 k 的倍数，这样每个 GPU 都有相同的工作量，就像只在单个GPU上训练一样。

5.3 简单网络

import torch
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l
import matplotlib.pyplot as plt

从零开始定义它，从而详细说明参数交换和同步。

# 初始化模型参数
scale = 0.01
W1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
W2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
W3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
W4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# 定义模型
def lenet(X, params):
    h1_conv = F.conv2d(input=X, weight=params[0], bias=params[1])
    h1_activation = F.relu(h1_conv)
    h1 = F.avg_pool2d(input=h1_activation, kernel_size=(2, 2), stride=(2, 2))
    h2_conv = F.conv2d(input=h1, weight=params[2], bias=params[3])
    h2_activation = F.relu(h2_conv)
    h2 = F.avg_pool2d(input=h2_activation, kernel_size=(2, 2), stride=(2, 2))
    h2 = h2.reshape(h2.shape[0], -1)
    h3_linear = torch.mm(h2, params[4]) + params[5]
    h3 = F.relu(h3_linear)
    y_hat = torch.mm(h3, params[6]) + params[7]
    return y_hat

# 交叉熵损失函数
loss = nn.CrossEntropyLoss(reduction='none')

5.4 数据同步

对于高效的多 GPU 训练，我们需要两个基本操作。 首先，我们需要 向多个设备分发参数 并附加梯度（get_params）。 如果没有参数，就不可能在 GPU 上评估网络。 第二，需要跨多个设备对参数求和，也就是说，需要一个 allreduce 函数。

def get_params(params, device):
    # 把params中的数据移入gpu[0]中
    new_params = [p.clone().to(device) for p in params] 
    for p in new_params:
        p.requires_grad_()  # 附加梯度
    return new_params

new_params = get_params(params, d2l.torch.try_gpu())
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)
###################### 
b1 weight: tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       device='cuda:0', requires_grad=True)
b1 grad: None

由于还没有进行任何计算，因此偏置参数的梯度仍然为零。 假设现在有一个向量分布在多个 GPU 上，下面的 allreduce 函数将所有向量相加，并将结果广播给所有 GPU。 请注意，我们需要将数据复制到累积结果的设备，才能使函数正常工作。

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i] = data[0].to(data[i].device)
data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('before allreduce:\n', data[0], '\n', data[1])
allreduce(data)
print('after allreduce:\n', data[0], '\n', data[1])

5.5 数据分发

我们需要一个简单的工具函数，将一个小批量数据均匀地分布在多个 GPU 上。

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

nn.parallel 中的函数记录

1. 复制(Replicate)：将模型拷贝到多个 GPU 上；
2. 分发(Scatter)：将输入数据根据其第一个维度(通常就是 batch 大小)划分多份，并传送到多个 GPU 上；
3. 收集(Gather)：从多个 GPU 上传送回来的数据，再次连接回一起；
4. 并行的应用(parallel_apply)：将第三步得到的分布式的输入数据应用到第一步中拷贝的多个模型上。

为了方便以后复用，我们定义了可以同时拆分数据和标签的 split_batch 函数。

def split_batch(X, y, devices):
    """将`X`和`y`拆分到多个设备上"""
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X, devices),
            nn.parallel.scatter(y, devices))

5.6 训练

现在我们可以 在一个小批量上实现多 GPU 训练。 在多个 GPU 之间同步数据将使用刚才讨论的辅助函数 allreduce 和 split_and_load。

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    # 在每个GPU上分别计算损失
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:  # 反向传播在每个GPU上分别执行
        l.backward()
    # 将每个GPU的所有梯度相加，并将其广播到所有GPU
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce([device_params[c][i].grad for c in range(len(devices))])
    # 在每个GPU上分别更新模型参数
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0]) # 在这里，我们使用全尺寸的小批量

现在，我们可以 定义训练函数。 训练函数需要分配 GPU 并将所有模型参数复制到所有设备。 显然，每个小批量都是使用 train_batch 函数来处理多个GPU。 我们只在一个 GPU 上计算模型的精确度，而让其他 GPU 保持空闲，尽管这是相对低效的，但是使用方便且代码简洁。

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # 将模型参数复制到`num_gpus`个GPU
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            # 为单个小批量执行多GPU训练
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        # 在GPU 0上评估模型
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
          f'on {str(devices)}')

# 看看 在单个GPU上运行 效果
train(num_gpus=1, batch_size=256, lr=0.2)
# 保持批量大小和学习率不变，并 增加为2个GPU
train(num_gpus=2, batch_size=256, lr=0.2)
plt.show()

6 多GPU的简洁实现

6.1 简单网络

让我们使用一个比 LeNet 更有意义的网络，它依然能够容易地和快速地训练。我们选择的是 ResNet-18。因为输入的图像很小，所以稍微修改了一下。在开始时使用了更小的卷积核、步长和填充，而且删除了最大汇聚层。

#@save
def resnet18(num_classes, in_channels=1):
    """稍加修改的 ResNet-18 模型。"""
    def resnet_block(in_channels, out_channels, num_residuals,
                     first_block=False):
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.append(d2l.Residual(in_channels, out_channels,
                                        use_1x1conv=True, strides=2))
            else:
                blk.append(d2l.Residual(out_channels, out_channels))
        return nn.Sequential(*blk)

    # 该模型使用了更小的卷积核、步长和填充，而且删除了最大汇聚层。
    net = nn.Sequential(
        nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm2d(64),
        nn.ReLU())
    net.add_module("resnet_block1", resnet_block(64, 64, 2, first_block=True))
    net.add_module("resnet_block2", resnet_block(64, 128, 2))
    net.add_module("resnet_block3", resnet_block(128, 256, 2))
    net.add_module("resnet_block4", resnet_block(256, 512, 2))
    net.add_module("global_avg_pool", nn.AdaptiveAvgPool2d((1,1)))
    net.add_module("fc", nn.Sequential(nn.Flatten(),
                                       nn.Linear(512, num_classes)))
    return net

6.2 网络初始化

net = resnet18(10)
# 获取GPU列表
devices = d2l.try_all_gpus()
# 我们将在训练代码实现中初始化网络

6.3 训练

如前所述，用于训练的代码需要执行几个基本功能才能实现高效并行：

• 需要在所有设备上初始化网络参数。
• 在数据集上迭代时，要将小批量数据分配到所有设备上。
• 跨设备并行计算损失及其梯度。
• 聚合梯度，并相应地更新参数。

最后，并行地计算精确度和发布网络的最终性能。除了需要拆分和聚合数据外，训练代码与前几章的实现非常相似。

def train(net, num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    def init_weights(m):
        if type(m) in [nn.Linear, nn.Conv2d]:
            nn.init.normal_(m.weight, std=0.01)
    net.apply(init_weights)
    # 在多个 GPU 上设置模型
    net = nn.DataParallel(net, device_ids=devices)
    trainer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        net.train()
        timer.start()
        for X, y in train_iter:
            trainer.zero_grad()
            X, y = X.to(devices[0]), y.to(devices[0])
            l = loss(net(X), y)
            l.backward()
            trainer.step()
        timer.stop()
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
    print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
          f'on {str(devices)}')

train(net, num_gpus=1, batch_size=256, lr=0.1)
plt.show()

6.4 小结

• 神经网络可以在（可找到数据的）单 GPU 上进行自动评估。
• 注意每台设备上的网络需要先初始化，然后再尝试访问该设备上的参数，否则会遇到错误。
• 优化算法在多个 GPU 上自动聚合。

7 参数服务器

当我们从一个 GPU 迁移到多个 GPU 时，以及再迁移到包含多个 GPU 的多个服务器时（可能所有服务器的分布跨越了多个机架和多个网络交换机），分布式并行训练算法也需要变得更加复杂。通过细节可以知道，一方面是不同的互连方式的带宽存在极大的区别。

7.1 数据并行训练

其中的关键是梯度的聚合需要在 GPU 0 上完成，然后再将更新后的参数广播给所有 GPU。

回顾来看，选择 GPU 0 进行聚合似乎是个很随便的决定，当然也可以选择 CPU 上聚合，事实上只要优化算法支持，在实际操作中甚至可以在某个GPU 上聚合其中一些参数，而在另一个 GPU 上聚合另一些参数。

请注意，我们还可以使用另一个工具来改善性能：在深度网络中，从顶部到底部计算所有梯度需要一些时间，因此即使还在忙着为某些参数计算梯度时，就可以开始为准备好的参数同步梯度了。

7.2 环同步（Ring Synchronization）

知识盲区，以后学习到相关内容填坑。

7.3. 多机训练

新的挑战出现在多台机器上进行分布式训练：我们需要服务器之间相互通信，而这些服务器又只通过相对较低的带宽结构连接，在某些情况下这种连接的速度可能会慢一个数量级，因此跨设备同步是个棘手的问题。毕竟，在不同机器上运行训练代码的速度会有细微的差别，因此如果想使用分布式优化的同步算法就需要 同步（synchronize）这些机器。

不会的点太多了，日后填坑。