• 人工智能
• 深度学习
• 机器学习
• 教程

本指南基于《动手学深度学习》（d2l-zh）中 多GPU训练 章节展开。你将掌握在单机多卡环境下并行训练深度学习模型的三种拆分思路，重点吃透数据并行的完整实现链路——从参数分发、梯度聚合（allreduce）到批量切分与同步训练，并能在 MXNet、PyTorch、PaddlePaddle 三种框架下亲手跑通一个从零编写的多 GPU 训练示例。

三种并行策略：网络并行、分层并行与数据并行

在讨论实现之前，先回答一个根本问题：手头有多块 GPU，该如何把训练"拆开"？在 多GPU训练 中，作者给出了三种候选方案，它们分别从不同维度切割训练任务。

网络并行：按层切分模型

第一种方法是在多个 GPU 之间拆分网络。每个 GPU 只负责特定层（或连续若干层）的计算：数据流入该 GPU，经处理后发给下一个 GPU。这种做法的直接收益是——可以用远超单卡显存容量的网络来处理数据，且每个 GPU 的显存占用只是整个网络的一小部分，便于控制。

但代价同样明显：

• GPU 之间需要密集同步（barrier operation），尤其是当各层计算负载不匹配时，负载轻的 GPU 会持续空等；
• 层间接口需要搬运大量激活值和梯度，传输量可能超出 GPU 总线带宽；
• 计算密集型操作的顺序安排本质上是困难的组合优化问题。

结论很明确：除非框架或操作系统原生支持将多 GPU 连接成虚拟单设备，否则不建议采用网络并行。

分层并行：切分通道与输出单元

第二种方法是在层内部拆分工作：例如将原本在单个 GPU 上计算 64 个通道的任务，分散到 4 个 GPU 上各自计算 16 个通道；全连接层同样可以按输出单元数量切分。早期的 AlexNet 设计（见 模型并行示意图）就是这种策略的典型应用——当时 GPU 显存仅有 2GB，不得不把卷积通道和全连接单元拆到多卡上。

这种方案在通道或单元数量足够大时能获得良好的计算性能提升，且显存随 GPU 数量线性扩展，可以支撑不断变大的网络。但它的致命伤与网络并行类似：每一层都依赖于其他层的结果，需要大量屏障操作，且需要传输的数据量可能比跨 GPU 拆分层时更大。基于带宽成本与复杂度，同样不建议优先选择。

数据并行：切分小批量样本

第三种方法是跨多个 GPU 拆分数据：所有 GPU 执行完全相同的计算逻辑，只是各自处理不同的观测样本；每个小批量训练完成后，梯度在各 GPU 间聚合。三种并行方式的对比可参见 splitting.svg 示意图。

数据并行之所以成为实践主流，原因很朴素：

• 实现最简单，可以应用到任何模型结构；
• 同步只发生在每个小批量处理完之后，且某个参数梯度计算完成即可开始交换（无需等待全部梯度）；
• GPU 数量越多，等效小批量越大，训练吞吐随之提高；
• 唯一限制是模型必须能装进单卡显存——这也是如今显存普遍充足后该方案胜出的原因。

分布式训练分区的更多理论细节可参考 Li 等人 2014 年的综述工作（文中引用Li.Andersen.Park.ea.2014）。下文将集中实现数据并行。

数据并行的工作流程

假设一台机器有 $k$ 个 GPU，每个 GPU 维护一份完整且同步的模型参数副本。如># MXNet 版本 scale = 0.01 W1 = np.random.normal(scale=scale, size=(20, 1, 3, 3)) b1 = np.zeros(20) W2 = np.random.normal(scale=scale, size=(50, 20, 5, 5)) b2 = np.zeros(50) W3 = np.random.normal(scale=scale, size=(800, 128)) b3 = np.zeros(128) W4 = np.random.normal(scale=scale, size=(128, 10)) b4 = np.zeros(10) params = [W1, b1, W2, b2, W3, b3, W4, b4] def lenet(X, params): h1_conv = npx.convolution(data=X, weight=params[0], bias=params[1], kernel=(3, 3), num_filter=20) h1_activation = npx.relu(h1_conv) h1 = npx.pooling(data=h1_activation, pool_type='avg', kernel=(2, 2), stride=(2, 2)) h2_conv = npx.convolution(data=h1, weight=params[2], bias=params[3], kernel=(5, 5), num_filter=50) h2_activation = npx.relu(h2_conv) h2 = npx.pooling(data=h2_activation, pool_type='avg', kernel=(2, 2), stride=(2, 2)) h2 = h2.reshape(h2.shape[0], -1) h3_linear = np.dot(h2, params[4]) + params[5] h3 = npx.relu(h3_linear) y_hat = np.dot(h3, params[6]) + params[7] return y_hat loss = gluon.loss.SoftmaxCrossEntropyLoss()

# PyTorch 版本 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')

# PaddlePaddle 版本 scale = 0.01 W1 = paddle.randn(shape=[20, 1, 3, 3]) * scale b1 = paddle.zeros(shape=[20]) W2 = paddle.randn(shape=[50, 20, 5, 5]) * scale b2 = paddle.zeros(shape=[50]) W3 = paddle.randn(shape=[800, 128]) * scale b3 = paddle.zeros(shape=[128]) W4 = paddle.randn(shape=[128, 10]) * scale b4 = paddle.zeros(shape=[10]) params = [W1, b1, W2, b2, W3, b3, W4, b4] def lenet(X, params): h1_conv = F.conv2d(x=X, weight=params[0], bias=params[1]) h1_activation = F.relu(h1_conv) h1 = F.avg_pool2d(x=h1_activation, kernel_size=(2, 2), stride=(2, 2)) h2_conv = F.conv2d(x=h1, weight=params[2], bias=params[3]) h2_activation = F.relu(h2_conv) h2 = F.avg_pool2d(x=h2_activation, kernel_size=(2, 2), stride=(2, 2)) h2 = h2.reshape([h2.shape[0], -1]) h3_linear = paddle.mm(h2, params[4]) + params[5] h3 = F.relu(h3_linear) y_hat = paddle.mm(h3, params[6]) + params[7] return y_hat loss = nn.CrossEntropyLoss(reduction='none')

注意这里的损失函数在 PyTorch/Paddle 中显式使用reduction='none'，因为后续要对每个 GPU 的损失单独求和（见下文train_batch）。

数据同步：参数分发与 allreduce

多 GPU 训练的核心原语只有两个：

1. get_params：把同一组参数复制到每个目标设备，并开启梯度记录；
2. allreduce：把分布在多设备上的梯度向量相加，并将结果广播回所有设备。

get_params：向多个设备分发参数

# MXNet def get_params(params, device): new_params = [p.copyto(device) for p in params] for p in new_params: p.attach_grad() return new_params

# PyTorch def get_params(params, device): new_params = [p.to(device) for p in params] for p in new_params: p.requires_grad_() return new_params

# PaddlePaddle def get_params(params, device): new_params = [paddle.to_tensor(p, place=device) for p in params] for p in new_params: p.stop_gradient = False return new_params

三个框架的语义一一对应：MXNet 用copyto(device)+attach_grad()，PyTorch 用.to(device)+requires_grad_()，Paddle 用paddle.to_tensor(p, place=device)+stop_gradient = False。若不做这一步，GPU 上根本没有参数可评估。验证如下——刚复制完、尚未计算任何梯度时，梯度应为零：

new_params = get_params(params, d2l.try_gpu(0)) print('b1 权重:', new_params[1]) print('b1 梯度:', new_params[1].grad)

这里d2l.try_gpu(0)的语义可参考源码实现：PyTorch 版本在 d2l/torch.py 中按torch.cuda.device_count()判断返回cuda:i还是cpu；MXNet 版本在 d2l/mxnet.py 中按npx.num_gpus()判断；Paddle 版本在 d2l/paddle.py 中按paddle.device.cuda.device_count()判断返回CUDAPlace(i)或CPUPlace()。这意味着没有 GPU 的环境也能跑通本教程（自动退回 CPU），只是并行加速无从谈起。

allreduce：聚合梯度并广播

# MXNet def allreduce(data): for i in range(1, len(data)): data[0][:] += data[i].copyto(data[0].ctx) for i in range(1, len(data)): data[0].copyto(data[i])

# PyTorch 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)

# PaddlePaddle def allreduce(data): paddle.set_device("gpu:0") for i in range(1, len(data)): data[0] += paddle.to_tensor(data[i], place=data[0].place) for i in range(1, len(data)): data[i] = paddle.to_tensor(data[0], place=data[i].place)

逻辑分两段：第一段把第 1～k-1 个设备上的数据累加到第 0 个设备（累加前需先把数据复制到目标设备，否则跨设备加法无法进行）；第二段再把累加结果广播回所有设备。用两个设备上不同取值的向量验证：

# PyTorch 示例 data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)] print('allreduce之前：\n', data[0], '\n', data[1]) allreduce(data) print('allreduce之后：\n', data[0], '\n', data[1])

allreduce之前data[0]=[[1,1]]、data[1]=[[2,2]]；之后两者都变成[[3,3]]。MXNet 与 Paddle 版本用法一致（Paddle 需先构造devices = [d2l.try_gpu(i) for i in range(num_gpus)]）。需要指出：这个朴素实现存在数据搬运瓶颈，练习 3 中要求实现更高效率的allreduce（例如树状聚合），正是因为朴素实现中每个设备的数据都要经过设备 0 中转。

数据分发：把小批量切成 k 份

需要一个工具函数把一个小批量均匀切分到多个 GPU。三个框架都提供了内置或易写的切分手段：

# MXNet：gluon.utils.split_and_load data = np.arange(20).reshape(4, 5) devices = [npx.gpu(0), npx.gpu(1)] split = gluon.utils.split_and_load(data, devices)

# PyTorch：nn.parallel.scatter data = torch.arange(20).reshape(4, 5) devices = [torch.device('cuda:0'), torch.device('cuda:1')] split = nn.parallel.scatter(data, devices)

# PaddlePaddle：手写 paddlescatter def paddlescatter(XY, devices): xy = XY.shape[0]//len(devices) # 根据GPU数目计算分块大小 return [paddle.to_tensor(XY[i*xy:(i+1)*xy], place=device) for i, device in enumerate(devices)] data = paddle.arange(20).reshape([4, 5]) split = paddlescatter(data, devices)

在 $4\times5$ 矩阵、2 个 GPU 的例子中，输出是两块 $2\times5$ 的子矩阵分别落在两个设备上。为了方便复用，封装一个同时切分特征与标签的函数（该函数在本书后续章节被反复使用，故标注@save便于从d2l包中直接调用）：

# MXNet #@save def split_batch(X, y, devices): """将X和y拆分到多个设备上""" assert X.shape[0] == y.shape[0] return (gluon.utils.split_and_load(X, devices), gluon.utils.split_and_load(y, devices))

# PyTorch #@save 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))

# PaddlePaddle #@save def split_batch(X, y, devices): """将X和y拆分到多个设备上""" assert X.shape[0] == y.shape[0] return (paddlescatter(X, devices), paddlescatter(y, devices))

训练实现：一个完整的小批量多 GPU 训练

有了split_batch、allreduce与参数副本，单小批量训练train_batch就水到渠成。以 PyTorch 版本为例：

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]) # 这里使用全尺寸的小批量

其执行顺序严格对应数据并行的五个步骤：

1. split_batch切分并分发数据；
2. 列表推导式逐个 GPU 完成前向并得到各分片损失（.sum()是为了把分片内逐样本损失聚合成标量）；
3. 逐个backward()得到各 GPU 的局部梯度；
4. 逐参数allreduce把 $k$ 份梯度相加并广播——注意 PyTorch 中需用with torch.no_grad()包裹，避免在聚合阶段再追踪计算图；
5. 每个 GPU 用d2l.sgd(param, lr, X.shape[0])更新参数，X.shape[0]是切分前的全批量大小，保证梯度除以的是整个小批量样本数而非分片样本数。

MXNet 版用autograd.record()包裹前向、直接对device_params各参数梯度做allreduce；Paddle 版需在每个 GPU 上切换paddle.set_device(f"gpu:{i}")后再前向/反向/更新，语义完全一致。完整源码见 多GPU训练章节。

值得强调：我们不需要编写任何并行代码。因为计算图在小批量内部的设备之间没有任何依赖关系，框架会"自动地"并行执行各个分片的计算。sgd的更新语义可对照仓库实现——d2l/torch.py 中param -= lr * param.grad / batch_size并对梯度清零；d2l/paddle.py 中用set_value写回并clear_gradient()。

定义训练主循环

训练函数与前面章节的区别在于：需要枚举可用 GPU、把模型参数复制到所有设备，并逐小批量调用train_batch。评估只在一个 GPU（GPU 0）上进行，其他 GPU 保持空闲——虽然相对低效，但代码简洁、便于说明原理。PyTorch 版：

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() # 在GPU0上评估模型 animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu( lambda x: lenet(x, device_params[0]), test_iter, devices[0]),)) print(f'测试精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/轮，' f'在{str(devices)}')

关键点：

• d2l.try_gpu(i)逐个取出第 $i$ 块 GPU（不够则回退 CPU），设备列表devices的长度即并行度 $k$；
• device_params是 $k$ 份参数副本的列表，每个元素是完整参数列表在该 GPU 上的拷贝；
• 每处理完一个小批量调用torch.cuda.synchronize()（MXNet 为npx.waitall()，Paddle 为paddle.device.cuda.synchronize()），强制等待异步内核执行完毕，保证计时与下一轮迭代的数据依赖正确；
• 评估复用仓库中的d2l.evaluate_accuracy_gpu，其 PyTorch 实现在 d2l/torch.py，MXNet 实现在 d2l/mxnet.py，Paddle 实现在 d2l/paddle.py——只把数据搬到指定设备上计算精度，与训练的多卡逻辑无关。

单卡与双卡实验对比

先在 1 个 GPU 上运行，批量大小 256、学习率 0.2：

train(num_gpus=1, batch_size=256, lr=0.2)

再保持批量大小与学习率不变，增加到 2 个 GPU：

# MXNet / PyTorch train(num_gpus=2, batch_size=256, lr=0.2)

两个实验的测试精度基本持平——这符合预期：不同 GPU 个数在算法寻优方面是等价的，数据并行只是放大吞吐而非改变优化轨迹。但本章作者也坦率地指出：在这个例子里看不到有意义的加速。原因有两个：

• 模型（LeNet）太小，计算密度不足以掩盖通信开销；
• Fashion-MNIST 数据集太小，且朴素allreduce实现受巨大的 Python 开销影响。

因此该实验更多是教学验证（确认多卡训练的数值正确性），真正的加速收益要留待更复杂的模型（如 ResNet）与更精细的并行化方法——这正是下一章 多GPU训练简洁实现 以及 参数服务器 要解决的问题。

小结

• 深度网络的多 GPU 训练存在三种拆分方式：层之间拆分（网络并行）、层内拆分（分层并行）、跨数据拆分（数据并行）。前两者需要严格编排数据传输（屏障操作多、带宽需求高），数据并行最简便，也是当前主流；
• 数据并行本身并不复杂：通过扩大等效小批量提高训练效率；每个 GPU 独立做前向与反向，之后所有梯度聚合为一，再向所有 GPU 广播；
• 实现数据并行只需两个原语：get_params（分发参数并开启梯度）与allreduce（梯度求和后广播），再加上split_batch（数据切分分发）；
• 由于等效小批量变大，学习率需要相应提高；批量大小应是 GPU 数量的倍数；BatchNorm 需为每个 GPU 维护独立统计参数。

练习

1. 在 $k$ 个 GPU 上训练时，把批量大小从 $b$ 改为 $k \cdot b$，即按 GPU 数量线性扩展，观察精度与吞吐的变化；
2. 比较不同学习率下的模型精度，思考随 GPU 数量增加学习率应如何缩放（可结合小批量随机梯度下降 sgd 章节 的理论）；
3. 实现一个更高效的allreduce函数用于聚合不同 GPU 上的参数——例如树状（ring/tree）聚合，思考为什么它比朴素逐设备累加效率更高；
4. 实现模型在多个 GPU 下的测试精度计算——可参考仓库中 MXNet 的evaluate_accuracy_gpus实现（d2l/mxnet.py），它把测试数据同样切分到所有设备并行推理后再合并精度统计，思路可作为 PyTorch/Paddle 的移植参考。

本节的完整代码与三种框架的实现均可直接在本仓库的 多GPU训练章节 中复现运行；更高层的多 GPU 封装见 multiple-gpus-concise.md，本章在《动手学深度学习》"计算性能"章节中位于 硬件与批大小权衡 之后，是理解大规模分布式训练（参数服务器、多机并行）的必经起点。

• 人工智能
• 深度学习
• 机器学习
• 教程