Ascend O2 BF16 Optimizer

背景

大模型训练过程中，显存资源紧缺是常见问题。其中优化器和梯度占据大量显存，但难以做显著的优化。

解决思路

使用半精度的优化器参数和梯度，大幅降低显存占用。

使用方法和理论收益

分别使用两个参数，即可分别使能特性。

记模型参数量为N，记开启分布式并行优化器场景下数据并行度为DP，则参数的优化效果如下表所示：

特性名称	特性参数	优化效果(Byte)
半精度优化器	--o2-optimizer	4N/DP
半精度梯度	--o2-gradient	2N

若不开启分布式并行优化器，则DP=1。

使用场景

1. 应当在相同的训练场景中，以特性对标的目的开启。相比Megatron原生混精场景，理论上O2特性是一定会影响精度的。
2. 目前O2特性仅支持参数类型为BF16的训练场景。

使用说明

1. 半精度优化器、半精度梯度特性之间没有依赖关系，可以分别开启或组合开启。
2. 半精度梯度不能与matmul add梯度累积特性一同使用，原因是算子不支持。 即o2-gradient参数必须与no-gradient-accumulation-fusion参数一起使用。
3. O2特性对性能影响：理论上不会影响，实验上根据llama2-7B 4K预训练数据，性能影响 < 1%。