跳转至

3.2 数据并行DP

数据并行原理


学习目标

  • 理解数据并行DP的原理.
  • 理解分布式数据并行DDP的原理.

数据并行

  • 所谓数据并行, 就是由于训练数据集太大, 因此将数据集分为N份, 每一份分别装载到N个GPU节点中. 同时, 每个GPU节点持有一个完整的模型副本, 分别基于每个GPU中的数据去进行梯度求导. 然后, 在GPU0上对每个GPU中的梯度进行累加. 最后, 再将GPU0聚合后的结果广播到其他GPU节点, 如下图所示:


  • 当然, 也可以将参数服务器分布在所有GPU节点上面, 每个GPU只更新其中一部分梯度, 如下图所示:



数据并行Pytorch DP

  • 数据并行的流程: 数据并行(torch.nn.DataParallel), 这是Pytorch最早提供的一种数据并行方式, 它基于单进程多线程进行实现的, 它使用一个进程来计算模型权重, 在每个批处理期间将数据分发到每个GPU:
    • 1: 将inputs从主GPU分发到所有GPU上.
    • 2: 将model从主GPU分发到所有GPU上.
    • 3: 每个GPU分别独立进行前向传播, 得到outputs.
    • 4: 将每个GPU的outputs发回主GPU.
    • 5: 在主GPU上, 通过loss function计算出loss, 对loss function求导, 求出损失梯度.
    • 6: 计算得到的梯度分发到所有GPU上.
    • 7: 反向传播计算参数梯度.
    • 8: 将所有梯度回传到主GPU, 通过梯度更新模型权重.
    • 9: 不断重复上面的过程.


  • 数据并行使用起来非常简单:
# 将数据和模型分布式放在0号GPU, 1号GPU, 2号GPU上.
net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

# input_var can be on any device, including CPU
output = net(input_var)

  • Pytorch DP的缺点:
    • 单进程多线程带来的问题: DataParallel使用单进程多线程进行实现的, 方便了信息的交换, 但受困于GIL, 会带来性能开销, 速度很慢. 而且, 只能在单台服务器(单机多卡)上使用(不支持分布式). 同时, 不能使用Apex进行混合精度训练.
    • 效率问题: 主卡性能和通信开销容易成为瓶颈, GPU利用率通常很低, 数据集需要先拷贝到主进程, 然后再分片(split)到每个设备上; 权重参数只在主卡(GPU0)上更新, 需要每次迭代前向所有设备做一次同步; 每次迭代的网络输出需要聚集到主卡(GPU0)上. 所以通信很快成为一个瓶颈. 除此之外, 这将导致主卡和其他卡之间, GPU利用率严重不均衡(比如主卡使用了10G显存, 而其他卡只使用了2G显存, batch size稍微设置大一点主卡的显存就OOM了).
    • 不支持模型并行, 由于其本身的局限性, 没办法与模型并行组合使用.

  • 注意: 目前PyTorch官方建议使用DistributedDataParallel, 而不是DataParallel类来进行多GPU训练, 即使在单机多卡的情况下!!!


分布式数据并行Pytorch DDP

  • 分布式数据并行(torch.nn.DistributedDataParallel), 基于多进程进行实现的, 每个进程都有独立的优化器, 执行自己的更新过程. 每个进程都执行相同的任务, 并且每个进程都与所有其他进程通信. 进程(GPU)之间只传递梯度, 这样网络通信就不再是瓶颈.


  • Pytorch DDP具体流程:
    • 1: 首先将rank=0进程中的模型参数广播到进程组中的其他进程.
    • 2: 每个DDP进程都会创建一个local Reducer来负责梯度同步.
    • 3: 在训练过程中, 每个进程从磁盘加载batch数据, 并将它们传递到其GPU. 每个GPU都有自己的前向过程, 完成前向传播后, 梯度在各个GPUs间进行All-Reduce, 每个GPU都收到其他GPU的梯度, 从而可以独自进行反向传播和参数更新.
    • 4: 同时, 每一层的梯度不依赖于前一层, 所以梯度的All Reduce和后向过程同时计算, 以进一步缓解网络瓶颈.
    • 5: 在后向过程的最后, 每个节点都得到了平均梯度, 这样各个GPU中的模型参数保持同步.


  • 注意: DataParallel是将梯度reduce到主卡, 在主卡上更新参数, 再将参数broadcast给其他GPU, 这样无论是主卡的负载还是通信开销都比DDP大很多). 相比于DataParallel, DistributedDataParallel方式可以更好地进行多机多卡运算, 更好的进行负载均衡, 运行效率也更高, 虽然使用起来较为麻烦, 但对于追求性能来讲是一个更好的选择.

  • 下面给出一个DDP的代码示例:
# 使用torch.nn.Linear作为本地模型, 用DDP对其进行包装, 
# 然后在DDP模型上运行一次前向传播, 一次反向传播和更新优化器参数步骤. 
# 之后, 本地模型上的参数将被更新, 并且不同进程上的所有模型完全相同.

import torch
import dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # 创建进程组
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # 实例化本地模型
    model = nn.Linear(10, 10).to(rank)

    # 实例化DDP模型对象
    ddp_model = DDP(model, device_ids=[rank])

    # 定义损失函数和优化器
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # 前向计算
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)

    # 反向传播
    loss_fn(outputs, labels).backward()

    # 参数更新
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example, args=(world_size,), nprocs=world_size, join=True)

if __name__ == '__main__':
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()

DP与DDP的区别

  • DP与DDP的主要区别有如下几点:
    • DP是基于单进程多线程的实现, 只用于单机情况, 而DDP是多进程实现的, 每个GPU对应一个进程, 适用于单机和多机情况, 是真正实现分布式训练. 并且因为每个进程都是独立的Python解释器, DDP避免了GIL带来的性能开销.
    • 参数更新的方式不同, DDP在各进程梯度计算完成之后, 各进程需要将梯度进行汇总平均, 然后再由rank=0的进程, 将其广播到所有进程后, 各进程用该梯度来独立的更新参数(而DP是梯度汇总到GPU0, 反向传播更新参数, 再广播参数给其他剩余的GPU). 由于DDP各进程中的模型, 初始参数一致(初始时刻进行一次广播), 而每次用于更新参数的梯度也一致; 因此, 各进程的模型参数始终保持一致(而在DP中, 全程维护一个optimizer, 对各个GPU上梯度进行求平均, 而在主卡进行参数更新, 之后再将模型参数广播到其他GPU). 相较于DP, DDP传输的数据量更少, 训练更高效, 不存在DP中负载不均衡的问题. 目前, 基本上DP已经被弃用.
    • DDP支持模型并行, 而DP并不支持, 这意味如果模型太大单卡显存不足时, 只能使用DDP.

  • DP数据传输过程:
    • 1: 前向传播得到的输出结果gather到主cuda计算loss.
    • 2: scatter上述loss到各个cuda.
    • 3: 各个cuda反向传播计算得到梯度后gather到主cuda后, 主cuda的模型参数被更新.
    • 4: 主cuda将模型参数broadcast到其它cuda设备上, 至此, 完成权重参数值的同步.

  • DDP数据传输过程:
    • 1: 前向传播的输出和loss的计算都是在每个cuda独立计算的.
    • 2: 梯度all-reduce到所有的CUDA(传输梯度).
    • 3: 这样初始参数相同, para.grad也相同, 反向传播后参数就还是保持一致的.