weblog

Appearance

自动并行 ​

一、什么是自动并行 ​

深度学习框架（比如 PyTorch、MXNet、飞桨）会在后端自动构建计算图，通过计算图可以知道各个任务之间的依赖关系，然后自动把没有依赖的任务并行执行，从而提高计算速度。

简单来说，就是框架会自动帮我们把能同时做的任务放在一起做，不用我们自己写复杂的并行代码。比如我们要初始化两个独立的变量，框架可以同时初始化这两个变量，而不是一个一个来。

二、多 GPU 的自动并行计算 ​

首先我们需要至少两个 GPU 来演示自动并行的效果。我们先定义一个测试用的任务：做 50 次矩阵乘法，分别在两个 GPU 上运行。

1. 先测试单独运行的耗时 ​

我们先在每个 GPU 上单独运行这个任务，看看各自的耗时：

python
import torch
from d2l import torch as d2l

# 获取所有可用的GPU
devices = d2l.try_all_gpus()

# 定义任务：做50次矩阵乘法
def run(x):
    return [x.mm(x) for _ in range(50)]

# 在两个GPU上分别创建随机矩阵
x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])

# 预热设备，避免缓存影响结果
run(x_gpu1)
run(x_gpu2)
# 等待GPU计算完成
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])

# 测试GPU1单独运行的耗时
with d2l.Benchmark('GPU1 time'):
    run(x_gpu1)
    torch.cuda.synchronize(devices[0])

# 测试GPU2单独运行的耗时
with d2l.Benchmark('GPU2 time'):
    run(x_gpu2)
    torch.cuda.synchronize(devices[1])

运行结果大概是：

Plain
GPU1 time: 0.4600 sec
GPU2 time: 0.4706 sec

两个 GPU 单独运行的时间加起来大概是 0.93 秒左右。

2. 测试自动并行的耗时 ​

现在我们让两个 GPU 同时运行这个任务，看看总耗时：

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

运行结果大概是：

Plain
GPU1 & GPU2: 0.4580 sec

总耗时比两个单独运行的时间加起来少很多，因为框架自动把两个 GPU 的任务并行执行了，相当于两个 GPU 同时在干活，所以总时间差不多和单个 GPU 运行的时间一样。

三、计算和通信的并行优化 ​

在很多情况下，我们需要在不同设备之间传输数据，比如把 GPU 上的计算结果复制到 CPU 上。如果我们先等 GPU 算完所有结果，再复制到 CPU，效率会很低。

1. 先测试串行的耗时 ​

我们先测试先计算再复制的耗时：

python
# 定义复制到CPU的函数
def copy_to_cpu(x):
    return [y.to('cpu') for y in x]

# 先在GPU1上计算，再复制到CPU
with d2l.Benchmark('在GPU1上运行'):
    y = run(x_gpu1)
    torch.cuda.synchronize()

with d2l.Benchmark('复制到CPU'):
    y_cpu = copy_to_cpu(y)
    torch.cuda.synchronize()

运行结果大概是：

Plain
在GPU1上运行: 0.4608 sec
复制到CPU: 2.3504 sec

总耗时大概是 2.8 秒左右。

2. 测试并行的耗时 ​

其实我们可以在 GPU 还在计算的时候，就开始把已经算好的结果复制到 CPU。比如 GPU 在算第 i 个矩阵乘法的时候，我们可以把第 i-1 个结果复制到 CPU，这样计算和通信就可以同时进行了。

在 PyTorch 里，我们可以设置non_blocking=True来实现这个效果：

python
# 定义支持非阻塞复制的函数
def copy_to_cpu(x, non_blocking=False):
    return [y.to('cpu', non_blocking=non_blocking) for y in x]

# 同时进行计算和复制
with d2l.Benchmark('在GPU1上运行并复制到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, True)
    torch.cuda.synchronize()

运行结果大概是：

Plain
在GPU1上运行并复制到CPU: 1.7703 sec

总耗时比串行的时候少了很多，因为计算和通信同时进行了，节省了时间。

四、自动并行的原理 ​

深度学习框架会在后端构建计算图，计算图里会记录每个任务的依赖关系：

  • 如果两个任务之间没有依赖（比如两个独立的矩阵乘法），框架就会把它们放在一起并行执行。

  • 如果两个任务有依赖（比如必须先算完 A 才能算 B），框架就会先执行 A，再执行 B。

比如一个两层的多层感知机，在 CPU 和两个 GPU 上运行的时候，框架会自动调度各个设备的计算和通信，不用我们自己手动安排。

五、小结 ​

  1. 自动并行就是框架自动把没有依赖的任务并行执行，提高计算速度，不用我们自己写复杂的并行代码。

  2. 多 GPU 的自动并行可以让多个 GPU 同时干活，减少总耗时。

  3. 计算和通信的并行可以让我们在计算的同时传输数据，节省时间。

  4. 框架会通过计算图来跟踪任务的依赖关系，自动优化执行顺序。

（注：文档部分内容可能由 AI 生成）

京ICP备2024093538号-1

自用笔记本