分布式训练

DP

最早的pytorch分布训练框架叫做dp(data parallel)

它的运行方式是硬盘读取数据然后给cpu分发给不同的显卡,每个gpu独立的计算出网络的梯度。然后所有的gpu都将自己计算的梯度传到一个gpu上。然后gpu0用全局平均梯度来更新网络参数。然后将更新后的参数广播到其他GPU上。在进行分布式训练时,最关键的任务是如何减少多卡的计算量。分布式训练中一半的时间可能都被通讯所占用。

总的GPU越多通信量就越大。

对于其他GPU梯度为X,传入参数为$\Psi$,传出梯度为$\Psi$

对于GPU0,传入为$(n-1)\Psi$输出参数为$(n-1)\Psi$

但是它由俩个问题。

  1. 但进程,多线程,python的GIL只能利用一个CPU核
  2. CPU1负责收集梯度,更新参数,同步参数。通信,计算压力大。

DDP

因此pytorch有一个替代DP的框架DDP:Distributied Data parallel

它把多个节点连接成一个环来进行通讯。

比如我们有GPU:0-2。我们希望通信结束后,让GPU都有,a0,a1,a2

首先

接着继续进行累加

接下来

这样每个gpu就有了所有的参数和。同时发送和接受,可以利用每个显卡的上下形带宽。

具体的步骤。它把神经网络按照参数定义反序排列,输出层最前面,输入层最后面。这是因为反向传播输出层的梯度是先被计算出来的,这样能在计算梯度的时候就将数据发送出去。每一层会放到不同的桶里面。

通讯量分析

参数量为$\Psi$,进程数为N

阶段传入传出$(n-1)\frac{\Psi}{N}\approx\Psi $

allGather阶段传入传出:$(n-1)\frac{\Psi}{N}\approx\Psi $

总传入传出$2\Psi$与集群大小无关。

deepspeed zero-1

上面的DDP每个神经网络需要存储完整的神经网络优化器。这个zero的意思是0荣誉优化器。

每个gpu在运行神经网络的时候需要存储

  • FP16参数
  • FP16梯度
  • fp32梯度
  • fp32一阶动量
  • fp32二阶动量
  • fp32参数

那么能不能让每个GPU存储一部分优化器的状态吗。比如我们让三个GPU每个存储三分之一的状态。

这样只需要在真正反向传播更新参数时的GPU有着一部分动量就可以完成更新了。

传播完成后,每个GPU拿到自己对应部分的梯度均值。然后将梯度转换为fp32,进行缩放,然后更新一阶二阶动量,最后优化器更新参数。最后更新优化器对应的那部分fp16的网络参数。在将更新后的部分广播给其他GPU。在发送梯度给其他gpu的时候当前gpu仍然可以继续计算。

zero2-3

zero2在deepSpeed zero的基础上进行再次优化,它把fp16梯度也按gpu进行了划分。

因为既然每个GPU只更新一部分的参数,那么其他的参数没有必要保存。

zero2的通信量与zero1和ddp的通讯量都是相同的。

zero3

zero3进一步的划分了参数,在前向传播的时候遇到自己没有的参数就靠其他gpu来广播给它。收到广播后计算完之后对该层进行丢弃。

阶段传入传出$(n-1)\frac{\Psi}{N}\approx\Psi $(不变)

allGather阶段传入传出:$2\times (n-1)\frac{\Psi}{N}\approx\Psi $

因此传输量是ddp的1.5倍

从上到下以此是ddp,zero0,zero1,zero2。

Last modification:October 9, 2024
如果觉得我的文章对你有用,请随意赞赏