【Pytorch Distributed】DP初步
Pytorch Distributed

【Pytorch Distributed】DP初步

一、DP 是什么

PyTorch 中的 DP，指的是 nn.DataParallel。它是一种最基础的数据并行方式，核心思想很简单：

  • 把同一个模型复制到多张 GPU 上

  • 把一个 batch 的输入数据沿着 batch 维切成几份

  • 每张 GPU 用自己的模型副本处理自己那一份数据

  • 最后把输出收集起来，反向传播时再把梯度汇总

所以，DP 的本质可以概括成一句话：

模型是复制的，数据是切分的。

这句话很重要，因为很多初学者一开始会误以为 DP 是“把模型拆成几块放到不同显卡上”。其实不是。
DP 并没有拆模型，它只是让每张卡都拥有一整份相同的模型，然后让它们分别处理不同的数据子批次。

二、先分清两个概念：参数和激活值

参数（Parameters）
参数就是模型本身要学习的内容，比如线性层里的 weight 和 bias。这些值不是某一次输入临时产生的，而是模型长期保存、通过训练不断更新的。

激活值（Activations）
激活值是forward过程中，由input经过网络计算后产生的中间结果。它们不是固定存在的，而是“这一批输入来了之后，现场算出来的”。

比如，假设模型要经历一层Linear，那输入x 在中间计算得到的Wx 就属于激活值

  • 参数不能按样本切分，因为每张卡想处理自己的样本，都需要用到完整的一套模型参数。

  • 激活值会随着输入一起变小，因为它是由输入生成的。输入切成了几份，每张卡就只会产生自己那一份输入对应的激活值。

三、DP 的代码例子和详解

下面是一个最常见的 DP 写法：

import torch
import torch.nn as nn

# 先把原始模型放到 cuda:0
# 在 DP 中，原始模型通常要放在 0 号卡上。后面多卡并行时，0 号卡不仅参与计算，还常常承担“主卡”的角色。
model = MyModel().to('cuda:0')

if torch.cuda.device_count() > 1:
    # 将模型用nn.DataParallel包起来
    # 告诉Pytorch，这个模型要在0和1两张卡上DP
    model = nn.DataParallel(model, device_ids=[0, 1])

# 将x放到gpu0上
x = x.to('cuda:0')
# forward
out = model(x)

下面是步骤详解：

假设现在：

  • 输入 x.shape = [30, 5]

  • 有 2 张 GPU

  • 模型是一层 Linear(5, 2)

那么 out = model(x) 这一步内部会发生什么？

1. 切分输入

原始输入：x.shape = [30, 5]

DP 会沿着第 0 维，也就是 batch 维，把它切成两份：

  • GPU0 拿到 x0.shape = [15, 5]

  • GPU1 拿到 x1.shape = [15, 5]

也就是说，原来 30 个样本，现在变成了两张卡各处理 15 个样本。

2. 每张卡开始并行 forward

切分完输入之后：

  • GPU0 上有一份完整模型副本

  • GPU1 上也有一份完整模型副本

然后两张卡同时做前向传播：

  • GPU0：[15, 5] -> [15, 2]

  • GPU1：[15, 5] -> [15, 2]

这两个 forward 是并行发生的。

3. 激活值计算

例如：

  • 单卡时输入是 [30, 5]

  • 经过某一层后，中间结果可能是 [30, 100]

如果改成双卡 DP：

  • GPU0 输入是 [15, 5]

  • GPU0 产生的中间结果就是 [15, 100]

  • GPU1 输入是 [15, 5]

  • GPU1 产生的中间结果就是 [15, 100]

四、单卡内 15 个样本是如何并行的

单卡内 15 个样本并不是 Python 层面一个个样本 for 循环计算。
它们会作为一个小 batch，一起组成张量，参与矩阵乘法、卷积等张量运算。GPU 会利用大量计算核心，并行完成这些运算。

代码里并不会写：

for sample in samples:
    out = model(sample)

而是直接把一个小 batch 丢进去：

out = model(x_chunk)   # x_chunk.shape = [15, 5]

这 15 个样本是作为一个整体一起参与计算的。

五、0 号卡的特点

0 号卡，也就是通常的 cuda:0，在 DP 中很特殊。它不只是普通参与计算的一张卡，往往还承担着“主卡”的职责。

它通常有这些特点：

1. 原始模型在 0 号卡上

DP 包装前，模型通常先被放到 cuda:0。
这意味着 0 号卡上会驻留原始模型。

2. 输出默认汇总到 0 号卡

两张卡各自算完输出后，DP 会把结果 gather 回某个设备，默认通常就是 0 号卡。
所以最终你在外面看到的整体输出，往往先汇总到这里。

3. 梯度最终也会回到原始模型

反向传播时，各卡根据自己的子 batch 计算出梯度贡献，最后这些梯度会汇总回原始模型所在的位置，也就是通常的 0 号卡。

4. 因此 0 号卡往往更容易占更多显存

这也是很多人第一次使用 DP 时会观察到的现象：

  • GPU0 显存更高

  • GPU1、GPU2 显存稍低一些

原因就在于：

  • 0 号卡既参与自己的那部分 forward

  • 又要承担原始模型、输出汇总、梯度汇总等职责

所以 DP 的显存使用通常不是完全平均的。
0 号卡常常是压力最大的那张卡。

七、把整个过程串起来

现在可以把一次 DP 前向传播理解成下面这条流程：

假设

  • 输入：[30, 5]

  • 2 张 GPU

  • 模型：Linear(5, 2)

流程

  1. 原始模型在 GPU0

  2. DP 在 GPU0 和 GPU1 上准备模型副本

  3. 输入 [30, 5] 沿 batch 维切成两份

    • GPU0：[15, 5]

    • GPU1：[15, 5]

  4. 两张卡并行 forward

    • GPU0：输出 [15, 2]

    • GPU1：输出 [15, 2]

  5. 把两张卡的结果收集回来

  6. 得到总输出 [30, 2]

  7. backward 时，各卡梯度再汇总回原始模型