经验总结

  1. 首页 > 生活 > >

使用Pytorch进行多卡训练( 三 )

生活百科

以上代码包含模型在多GPU上读取权重、进行分布式训练、保存权重等过程 。细节注释如下：
1、初始化进程组 ， 由于使用GPU通信 ， 后端应该写为NCCL 。不过经过实验 ， 即使错写为gloo ， DDP内部也会自动使用NCCL作为通信模块 。
2、由于后面使用DDP包裹模型进行训练 ， 其内部会自动将所有rank的模型权重同步为rank 0的权重 ， 因此我们只需在rank 0上读取模型权重即可 。这是基于Pytorch版本1.12.1 ， 低级版本似乎没有这个特性 ， 需要在不同rank分别导入权重 ， 则load需要传入map_location ， 如下面注释的两行代码所示 。
3、这里创建model的优化器 ， 而不是创建用ddp包裹后的ddp_model的优化器 ， 是为了兼容单GPU训练 ， 读取优化器权重更方便 。
4、将优化器权重读取至该进程占用的GPU 。如果没有map_location参数 ， load会将权重读取到原本保存它时的设备 。
5、优化器获取权重 。经过实验 ， 即使权重不在优化器所在的GPU ， 权重也会迁移过去而不会报错 。当然load直接读取到相应GPU会减少数据传输 。
6、DDP包裹模型 ， 为模型复制一个副本到相应GPU中 。所有rank的模型副本会与rank 0保持一致 。注意 ， DDP并不复制模型优化器的副本 ， 因此各进程的优化器需要我们在初始化时保持一致 。权重要么不读取 ， 要么都读取 。
7、这里开始模型的训练 。数据需转移到相应的GPU设备 。
8、在backward中 ， 所有进程的模型计算梯度后 ， 会进行平均（不是相加） 。也就是说 ， DDP在backward函数添加了hook ， 所有进程的模型梯度的ring_reduce将在这里执行 。这个可以通过给各进程模型分别输入不同的数据进行验证 ， backward后这些模型有相同的梯度 ， 且验算的确是所有进程梯度的平均 。此外 ， 还可以验证backward函数会阻断(block)各进程使用梯度 ， 只有当所有进程都完成backward之后 ， 各进程才能读取和使用梯度 。这保证了所有进程在梯度上的一致性 。
9、各进程优化器使用梯度更新其模型副本权重 。由于初始化时各进程模型、优化器权重一致 ， 每次反向传播梯度也保持一致 ， 则所有进程的模型在整个训练过程中都能保持一致 。
10、由于所有进程权重保持一致 ， 我们只需通过一个进程保存即可 。
11、定义rank 0的IP和端口 ， 使用mp.spawn ， 只需在主进程中定义即可 ， 无需分别在子进程中定义 。
12、创建子进程 ， 传入：子进程调用的函数（该函数第一个参数必须是rank）、子进程函数的参数（除了rank参数外）、子进程数、是否等待所有子进程创建完毕再开始执行 。
参考： https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

经验总结扩展阅读


上一篇：泳镜泳帽在什么店里有卖

下一篇：2023高考播音主持专业好考吗 通过率高吗