赞
踩
MegatronLM
的第二篇论文【Efficient Large-Scale Language Model Training on GPU ClustersUsing Megatron-LM】是2021年出的,同时GPT-3模型参数已经达到了175B参数,GPU显存占用越来越多,训练时间也越来越长。
在本文中,MegatronLM结合了tensor模型并行、pipeline模型并行和数据并行来提升整体的训练速度,以及解决大模型在多机上的扩展性问题。在pipeline模型并行上提出了interleaved pipelining schedule
方法,提升了10%的吞吐。
每个work都有一份模型的完整拷贝,输入数据被分成了多个shard(份)分给不同work进行计算。work会定期汇总梯度信息,保证所有worker的权重一致。在大模型训练时,一张卡放不下整个模型,只能放下部分模型,一个大模型会分成多个shard(份)放在不同卡上。
复用了之前Megatron-1上的Tensor并行策略,transformer层分为自注意力和两层MLP组成,以MLP为例,原有 Y = G e L U ( X A ) , Z = D r o p o u t ( Y B ) Y=GeLU(XA), Z=Dropout(YB) Y=GeLU(XA),Z=Dropout(YB) ; 由于GeLU是非线性的,所以这里对参数按列来切分 A = [ A 1 , A 2 ] A=[A_1, A_2] A=[A1,A2],得到 [ Y 1 , Y 2 ] = [ G e L U ( X A 1 ) , G e L U ( X A 2 ) ] \left[ Y_1, Y_2 \right] = \left[ GeLU(XA_1), GeLU(XA_2) \right] [Y1,Y2]=[GeLU(XA1),GeLU(XA2)] ; 对于B按行来切分参数, B = [ B 2 B 1 ] , Y = [ Y 1 , Y 2 ] B=\left[ {}_{B_2}^{B_1}\right], Y=\left[ Y_1, Y_2\right] B=[B2B1],Y=[Y1,Y2] 。
通过pipeline并行, 一个模型可以被拆解为多份放到不同节点上, 以transformer为例,在transformer中block多次重复,所以每个device会分别处理相同个数的layer,对于非对称结构不好切分,这里不做考虑。如下图,不同TransformerLayer做为不同的pipeline的stage,也称为partition。
pipeline并行的基本思路是把一个batch
切分为多个更小的microbatch
, 多个microbatch
同时并行训练,训练中要注意保证权重weight同步更新的语义。常见的pipeline有以下几种方式:
Gpipe中重点有两块,一块是Mini-Batches
改为了Micro-Batches
;另一块是通过recompute
方法在反向过程中降低显存占用, 主要是减少activation memory
的占用(这里的activation
指的是前向的输出结果,不是激活activation function
)。
Mini-Batches
的更新过程, 横轴表示时间,由于计算中存在上下依赖,会存在下方大面积的空白,空白部分也被称为Bubble
, 这个图中有个问题是下标没有对应更新为0,1,2,3
;Mini-Batches
被分为4个Micro-Batches
, 通过pipeline并行使用microbatch
可以大幅减少Bubble
,例如
F
1
,
2
F_{1,2}
F1,2 依赖
F
0
,
2
F_{0,2}
F0,2,
B
1
,
2
B_{1,2}
B1,2 依赖
B
2
,
2
B_{2, 2}
B2,2 和
F
1
,
2
F_{1, 2}
F1,2。在GPipe中所有的microbatch
的前向都先执行,然后再执行训练的后向。在所有前后向执行完后进行pipeline的flush操作。如图每个batch被切分为8个microbatch
。这里假设反向的计算耗时是前向的两倍。
接下来计算下pipeline的bubble耗时与理想计算时间的占比(Bubble time fraction (pipeline bubble size)
),pipeline的bubble耗时用
t
p
b
t_{pb}
tpb 表示,microbatche的大小用
m
m
m 表示,pipeline的stage个数用
p
p
p 表示,每个iter的理想时间用
t
i
d
t_{id}
tid 表示,执行一个microbatch的前向和后向分别用
t
f
t_f
tf 和
t
b
t_b
tb 表示。
以第一个microbatch为例,在前向计算时bubble耗时有 ( p − 1 ) × t f (p-1) \times t_f (p−1)×tf, 后向计算时bubble耗时有 ( p − 1 ) × t b (p-1) \times t_b (p−1)×tb, 总的bubble耗时有 t p b = ( p − 1 ) × ( t f + t b ) t_{pb} = (p-1) \times (t_f + t_b) tpb=(p−1)×(tf+tb) , 理想的计算时间 t i d = m ∗ ( t f + t b ) t_{id} = m * (t_f + t_b) tid=m∗(tf+tb)。最终理想计算耗时与bubble耗时的比例如下。为了降低占比就要增大 m m m, 使得 m ≫ p m \gg p m≫p , 这时显存占用就会上升。 计算公式如下:
B u b b l e t i m e f r a c t i o n ( p i p e l i n e b u b b l e s i z e ) = t p b t i d = p − 1 m Bubble time fraction (pipeline bubble size)=tpbtid=p−1m
通常情况下,会缓存记录下所有算子的前向计算结果,用于backward时的梯度计算,这部分结果的显存占用(activation memory
)可能会比weight参数的显存占用还大。为了减少少这部分activation memory
显存存占用,在GPipe中当进行反向计算时,第k个设备会重新计算前向的过程
F
k
F_k
Fk。通过重计算,activation memory
从
O
(
N
×
L
)
O\left( N \times L \right)
O(N×L) 降为了
O
(
N
+
L
K
×
N
M
)
O\left( N + \frac{L}{K} \times \frac{N}{M} \right)
O(N+KL×MN)。
L
L
L 表示总层数,
N
N
N 表示mini-batch的大小,
M
M
M 表示micro-batch的个数。
在GPipe的结果中, AmoebaNet
通过重计算,activation memory
从6.26G降为了3.46G;Transformer
通过重计算每个device/accelerator可以训练模型参数量较之前大了2.7倍。
论文出自PipeDream: Generalized Pipeline Parallelism for DNN Training。
之前的pipeline并行先统一做所有前向再开始做后向,后向计算中会用到前向对应的输出(activation)和前一个算子的后向输出(这里叫成state状态)。存在问题在于前后向的state在不同batch中可能会被误用,也就是没有对应的版本区分。GPipe的方式存在问题在于频繁的pipeline flush
也会降低效率。
在PipeDream中有两个重点,一个是论文提出了1F1B
算法进行并行的调度,使得硬件利用率在语义层面接近数据并行,不过跟数据并行不同的是在1F1B
使用了不同版本的weight参数进行训练,在训练中避免了GPipe中定期的pipeline flush
;另一个点是会自动进行pipeline stage的切分。
对于1F1B
算法实现来说,在前向startup
计算中,每个stage计算完前向后会异步把输出结果发送给下一个stage;最后一个stage来说,会在前向计算完成后立马开始后向的计算(worker4完成蓝色块1的前向计算后,直接开始绿色块1的后向计算);steady稳定状态下,前向和后向是交替进行的,也就是1F1B
。这样就不用在最后统一做梯度更新,每次PP(pipeline parallel)执行过程中都只传一部分梯度和输出结果给一个worker,减少了通信量。对于采用数据并行的stage,通过round-robin方法分发数据,保证一个minibatch的数据是在同一个worker上的。
对于pipeline stage的切分,PipeDream先通过profile方式在一个GPU卡上跑1000个minibatch,统计每层的前向和后向计算耗时,每层的输出大小和weight权重大小,预估all-reduce等通信量。partition算法来确定:如何划分stage到worker、哪些stage需要复制进行数据并行,以及mini-batch的大小。
在之前的pipeline并行系统中,每个stage的前向和后向用到的参数都是不同的,这样可能会导致模型收敛问题。如图4中,mini-batch 5是在1的后向算完执行,但5的后向是在2,3,4的后向算完才执行,用到的参数不一样。因此在PipeDream中采用了weight stashing
机制,前向stage每次都采用最新的参数处理minibatch,在前向算完后,计算完后保存参数提供给后向使用,保证一个stage内前向和后向用的参数是一致的, 如下图worker1(stage1)和worker3(stage3)中5的前向和后向都是用的一样的参数,不同stage的参数版本是不一样的。
同时还有一个vertical sync
可选的方法,为了消除跨stage的不一致问题,意思跟weight stashing
稍有不同是,在minibatch第一次进入pipeline stage时采用最新参数版本,固定当前的参数版本,在后续stage计算中都采用跟第一个stage相同的参数版本,相当于按列固化stage的参数版本号。
PipeDream-2BW
重点关注降低显存占用和提升吞吐,是在PipeDream的后续,文中提出了两种pipeline调度方法来降低显存,一种是Double-buffered weight updates (2BW)
,另一种是PipeDream-Flush
;同时PipeDream的Planner决定如何合理切分模型。
先回顾下GPipe和PipeDream的区别,再看和PipeDream-2BW
的比较。
1F1B
方式,会同时保留多个版本的参数;PipeDream中用的还是mini-batch,不同的mini-batch分别更新。
D
o
u
b
l
e
b
u
f
f
e
r
e
d
w
e
i
g
h
t
u
p
d
a
t
e
s
(
2
B
W
)
Double\ buffered\ weight\ updates\ (2BW)
Double buffered weight updates (2BW):在PipeDream-2BW
中的粒度是micro-batch, 对于每个micro-batch在前向和后向过程中都是采用相同版本的参数进行计算。不同的是PipeDream-2BW
最多只会保留两个版本的参数,减少显存占用。具体做法:每处理
m
m
m 个micro-batch(
m
≫
d
m \gg d
m≫d ,
d
d
d 是pipeline的stage数,也就是pipeline的深度)会更新一版参数,下图示例中
m
=
d
=
4
m=d=4
m=d=4 。新版权重不是立马使用,而是等之前用老版本参数的micro-batch的反向计算完后,新的micro-batch再采用新版的参数。跟Gpipe相比,这里是每过
m
m
m 个micro-batch进行梯度累加,假设micro-batch的大小是
b
b
b, 实际的有效的batch大小是
m
⋅
b
m \cdot b
m⋅b 。
W
e
i
g
h
t
u
p
d
a
t
e
s
w
i
t
h
F
l
u
s
h
e
s
(
P
i
p
e
D
r
e
a
m
F
l
u
s
h
)
Weight\ updates\ with\ Flushes\ (PipeDream\ Flush)
Weight updates with Flushes (PipeDream Flush):这个是另外一种节省显存的pipeline调度方法,比2BW
的显存更少,这个方法复用了1F1B
调度的思路,但不同的是只保存了一个版本的weight,PipeDream-Flush
可以看成是1F1B
方法在GPipe
中的应用。
在之前pipeline基础上很容易看到Megatron pipeline的思路,在Megatron中pipeline并行采用了Interleaved Schedule
的方式。之前做pipeline stage切分时,对于每个device会被分到连续的layer(比如device1分到layer1-4,device2分到layer5-8),一组连续的layer被称为一个model chunk
,每个device只会处理一个stage;为了减少pipeline bubble气泡的大小,在Interleaved Schedule
方法中,改为每个device会被分到两组model chunk
的内容,(比如device1分到layer1/2和layer9/10, device2分到layer3/4和layer11/12),这样每个device会来处理多个stage。
如下图,深色表示第一个model chunk的内容,浅色表示第二个model chunk的内容。每个device分到 v v v 个stage,对于每个microbatch的前向和后向时间分别为 t f / v t_f/v tf/v 和 t b / v t_b/v tb/v 。对应的bubble时间减少为 t p b i n t = ( p − 1 ) ⋅ ( t f + t b ) v t_{pb}^{int}=\frac{(p-1) \cdot (t_f + t_b)}{v} tpbint=v(p−1)⋅(tf+tb) 。
B u b b l e t i m e f r a c t i o n ( p i p e l i n e b u b b l e s i z e ) = t p b i n t t i d = 1 v ⋅ p − 1 m Bubble time fraction (pipeline bubble size)=tintpbtid=1v⋅p−1m
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。