Megatron 分布式优化器

问题分析

数据并行（DP）场景下， 梯度all-reduce之前，不同的设备并行执行不同的工作。然而all-reduce以后， 所有设备都利用同样的梯度、参数以及优化器状态做相同的梯度更新操作，因此存在计算和存储冗余。

解决思路

将重复的内存存储和计算步骤拆分给不同的设备，通过通信进行联通，最终在同时节省显存和计算时间的条件下，达到和原始DP相同的效果。

使用场景

主要用于训练场景，当DP不为1时，将优化器状态拆分到所有DP组间，在对梯度进行reduce-scatter后，每个DP组分别进行部分权重更新，最后再用一个all-gather操作收集模型权重。

使用方法

脚本中添加--use-distributed-optimizer开启分布式优化器。

使用影响

降低显存开销。