DeepSpeed三个ZeRO优化技术对比

1. ZeRO-1（拆分优化器状态）
   先理解什么是“优化器状态”：
   在训练神经网络时，我们使用优化器（例如 Adam、SGD）来更新模型参数。优化器除了模型参数外，还会存储一些额外的信息来帮助加速训练过程，比如：
   • 动量（Momentum）：记录之前梯度的方向，以更稳定地更新参数。
   • 平方梯度（在 Adam 优化器中）：用于调整学习率。
   这些额外的存储信息称为“优化器状态”，它会占用相当大的内存，特别是在大模型中。
   ZeRO-1 的核心原理：
   ZeRO-1 通过将优化器状态在多个 GPU 之间分配来减少单个 GPU 的内存负担。每个 GPU 不再存储整个模型的优化器状态，而是只存储它所负责的那部分。
   • 举例：假设有 4 个 GPU 和一个非常大的模型。如果没有 ZeRO，所有的 GPU 都需要保存完整模型的所有优化器状态。而在 ZeRO-1 中，每个 GPU 只需要存储优化器状态的一部分。比如，GPU 1 只存储动量的一部分，GPU 2 存储另一个部分，依此类推。
   • 效果：节省了内存，每个 GPU 只处理一部分优化器状态，这样就可以用相同的硬件训练更大的模型。
2. ZeRO-2（拆分优化器状态和梯度）
   先理解“梯度”：
   在模型训练时，我们会通过反向传播计算梯度，梯度是用来更新模型参数的数值。在每次反向传播后，优化器根据这些梯度调整模型的参数。
   梯度也会占用大量内存，特别是在大模型和大批量数据的情况下。每个 GPU 都需要存储模型的所有梯度，这会导致内存不足。
   ZeRO-2 的