Femtotron开发日志 #10 ZeRO-3、SAC和PP：一个简单的显存节省收益数据测试和分析

这里列几组数据，和测试它们的时候使用的参数，供大家参考，并且提供一些分析。具体的内存占用太过复杂，因此分析可能有不完全详尽之处，能对读者有所启发即可。

数据

PP=1、TP=1、DP=8 (固定num_layers=8, 变化seq_len)

参数如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
model_config = AutoConfig.for_model(
    "llama",
    hidden_size=1024,
    intermediate_size=2048,
    num_attention_heads=16,
    num_key_value_heads=4,
    num_hidden_layers=8,
    max_position_embeddings=128,
    vocab_size=1024,
    rms_norm_eps=1e-5,
    hidden_act="silu",
    tie_word_embeddings=False,
)

数据如下：

Config	seq=16	seq=32	seq=1024
baseline	1653	1653	3631
baseline + AC	1653	1653	1669
ZeRO-1	460	471	2913
ZeRO-1 + AC	460	460	901
ZeRO-2	443	444	2776
ZeRO-2 + AC	443	444	765
ZeRO-3	384	419	2840
ZeRO-3 + AC	261	261	711

PP=2、TP=2、DP=2 (固定seq=32,变化num_layers)

参数如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
model_config = AutoConfig.for_model(
    "llama",
    hidden_size=2048, 
    intermediate_size=8192,
    num_attention_heads=32, 
    num_key_value_heads=8,
    num_hidden_layers=8, 
    max_position_embeddings=128,
    vocab_size=40960, 
    rms_norm_eps=1e-5,
    hidden_act="silu", 
    tie_word_embeddings=False,
)

数据如下：

Config	N=8(4/stage)	N=32(16/stage)	N=64(32/stage)
ZeRO-0	3601	12558	24404
ZeRO-0 + AC	3601	11954	23091
ZeRO-1	2136	9534	18595
ZeRO-1 + AC	2136	6993	13490
ZeRO-2	2044	9030	17627
ZeRO-2 + AC	2044	6569	12602
ZeRO-3	2233	10808	20829
ZeRO-3 + AC	2216	6417	12018

ZeRO-3 vs ZeRO-2 + SAC 对比

设置	ZeRO-2 + AC	ZeRO-3 + AC	Δ(负 = ZeRO-3 赢)
PP=1, seq=16, N=8	443	261	−182
PP=1, seq=32, N=8	444	261	−183
PP=1, seq=1024, N=8	765	711	−54
PP=2, seq=32, N=8(4/stage)	2044	2216	+172 ✗
PP=2, seq=32, N=32(16/stage)	6569	6417	−152
PP=2, seq=32, N=64(32/stage)	12602	12018	−584

分析

从上面的数据可以得出几个浅显的结论：

1、在DP规模相对较大、TP和PP都不存在的时候，ZeRO-3的优势很明显。
2、在使用PP并行和DP规模相对较小的时候，ZeRO-3的优势几乎不再明显存在，甚至可能有反转的情况。
3、在seq_len较大的时候，ZeRO-3的显存节省会被削弱甚至被ZeRO-2反转，但配合SAC可以弥补这部分额外损失，恢复明显的优势。
4、在模型单层参数增多的时候，ZeRO-3 + SAC的优势更加微弱。在模型的单层参数不变，但是层数变深的时候，ZeRO-3 + SAC的优势更加明显。
5、ZeRO-1和ZeRO-2在各种情况下基本可以提供稳定的性能优化。
6、SAC单独使用的时候，其效果取决于seq_len等输入规模参数的大小。如果这些参数不够大，SAC可能完全没有效果。

它的根源在于ZeRO-3 + SAC的显存节省来源，和改变模型参数、改变并行配置的时候ZeRO-3 + SAC的额外显存增加来源并不相同。ZeRO-3是否真的能取得收益，取决于这样两个数值的此消彼长：ZeRO-3必须保持活跃的"瞬时块"大小，和ZeRO-3总共节省下来的内存大小。前者大，那么ZeRO-3会输给ZeRO-2，反之则超过。节省下来的内存大小理论上相对来说很好计算，就是权重中被shard掉的比例，而必须保持活跃的"瞬时块"大小则和很多因素有关系。目前能够明显观察到的有：1、PP stage的流水线拖延导致的缓冲区分配；2、seq_len拉大导致的缓冲区分配；3、单个layer的unshard过程造成的临时峰值。当然，这些分析是很粗浅的，不过基本已经足以指导我们对这些优化策略的使用与否。我非常相信这当中还有更深的奥妙有待研究，这就交给有兴趣的人进行进一步的分析了。