分布式训练
DP
最早的pytorch分布训练框架叫做dp(data parallel)
它的运行方式是硬盘读取数据然后给cpu分发给不同的显卡,每个gpu独立的计算出网络的梯度。然后所有的gpu都将自己计算的梯度传到一个gpu上。然后gpu0用全局平均梯度来更新网络参数。然后将更新后的参数广播到其他GPU上。在进行分布式训练时,最关键的任务是如何减少多卡的计算量。分布式训练中一半的时间可能都被通讯所占用。
总的GPU越多通信量就越大。
对于其他GPU梯度为X,传入参数为$\Psi$,传出梯度为$\Psi$
对于GPU0,传入为$(n-1)\Psi$输出参数为$(n-1)\Psi$
但是它由俩个问题。
- 但进程,多线程,python的GIL只能利用一个CPU核
- 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。