大模型分布式训练之DeepSpeed优化器并行(ZeRO)原理

由于大模型参数量非常庞大，所以我们常常需要用到分布式训练来解决训练过程中计算资源不足的问题，现在也出现了很多大模型相关的分布式训练框架，但是使用的比较多的还是deepspeed的数据并行，那么deepspeed是怎么实现数据并行的呢

DeepSpeed

DeepSpeed是一个开源的深度学习优化库，它由微软开发并维护，旨在提高大规模模型训练的效率和可扩展性。通过创新的算法和技术，DeepSpeed能够降低训练超大规模模型的复杂性和资源需求，让深度学习训练变得更快、更高效。

DeepSpeed特点和优势：

• 高效的并行化策略：DeepSpeed支持多种并行化方法，包括数据并行、模型并行和流水线并行。这些方法可以灵活组合，以适应不同规模和复杂度的深度学习模型。通过并行化，DeepSpeed能够显著提高训练速度和可扩展性。
• 内存优化技术：为了降低内存占用和提高训练效率，DeepSpeed引入了ZeRO（Zero Redundancy Optimizer）技术。ZeRO通过将优化器的状态、梯度和参数在分布式环境中进行分割，从而减少了冗余的内存占用。这使得在有限的内存资源下训练更大的模型成为可能。
• 混合精度训练支持：DeepSpeed支持混合精度训练，即同时使用单精度和半精度浮点数进行训练。这种方法可以在保持模型性能的同时，减少内存占用和计算时间，降低能耗。
• 易用性和兼容性：DeepSpeed与PyTorch等主流深度学习框架紧密集成，提供了易用的API和丰富的文档支持。这使得用户能够轻松地将DeepSpeed集成到他们的项目中，并充分利用其提供的优化功能。此外，DeepSpeed还提供了高度优化的数据加载和网络通信工具，以减少通信量并提高多GPU和多节点环境下的训练效率。

DeepSpeed实现数据并行主要是通过ZeRO（Zero Redundancy Optimizer）技术和混合进度训练，混合进度训练这里不再过多介绍，ZeRO原理到底是什么呢？小编花了一个晚上终于搞懂了！
优化器数据并行有三种方式，即ZeRO-1/23

ZeRO-1原理

只对optimizer状态进行切分
具体步骤：

1. 把batch分成N份，每张卡（GPU）一份

2. 执行一步前向和反向传播计算后，每个GPU各得一份梯度

3. 对梯度执行all-reduce操作，得到完整梯度，这里all-reduce操作是将每个节点上的部分梯度累加起来，并将结果广播到所有节点，这样所有节点都拥有一份完整的梯度

4. 每个 GPU 得到完整的梯度 G后，对各自的权重进行更新，权重的更新由优化器状态和梯度共同决定
   在这里，优化器的状态也被切分成了N份，每个GPU只需要存储和更新总优化器状态的 1/N ，并更新1/N 的权重；

5. 每个GPU维护各自优化器里面更新的权重，最后执行all-gather，使得每个GPU都有更新后的权重，all-gather 操作将更新后的权重同步到所有节点
   

ZeRO-2原理

在optimizer状态划分的基础上，再对梯度也进行划分，每个GPU各自维护一块自己的梯度
具体步骤：

1. 把batch分成N份，每张卡（GPU）一份
2. 执行一步前向和反向传播计算后，每个GPU各得一份梯度
3. 对梯度执行all-reduce，保证每个GPU所维护的梯度是聚合梯度（即各节点汇总之后的梯度），聚合之后对梯度进行切分
   具体是怎么切分梯度的呢？ 举个栗子——
   eg：① 比如GPU1只负责维护梯度G1，其他GPU只需要把对应位置梯度发给GPU1即可
   ② 汇总完毕后，其他不是GPU1维护的梯度会从GPU1中移除，即更新后马上释放
4. 每个GPU用所维护的优化器和梯度更新相应的权重，即每块GPU维护独立的权重
5. 最后对权重执行all-gather，将其他GPU的权重同步一份完整的到自己节点上来
   

ZeRO-3原理

ZeRO-3在ZeRO-1和ZeRO-2的基础上，对权重进行各自的维护，即不再维护一整份权重了
具体步骤：

1. 把batch分成N份，每张卡（GPU）一份
2. 这个时候模型的权重参数（张量）也被分成N份
3. 在进行前向计算之前，对权重执行all-gather操作取回分布在各GPU上的权重，组成完整的参数进行前向计算，计算完成后，把不属于自身维护的权重抛弃
4. 在进行反向传播计算之前，对权重执行all-gather操作取回分布在各GPU上的权重，组成完整的参数进行反向传播计算，计算完成后，把不属于自身维护的权重抛弃
5. backward之后得到各自的梯度，对梯度执行all-reduce，得到聚合的梯度之后更新其自身维护的权重，然后立刻把不是自己维护的梯度抛弃
6. 由于每个GPU只保存其自身维护的权重参数，因为无需对权重进行all-reduce
   
   注：

• ZeRO-3其实是增加了通信开销，来减少每张GPU的显存占用，以通信换显存
• 可能有人会问，ZeRO-3对权重参数进行了切分，应该算张量并行吧？？？其实不是，因为ZeRO-3在前向和反向的时候，还是用的完整的权重来计算的（张量并行在前向和反向的时候也只用一部分权重）

ZeRO-Offload

把占用显存多的部分卸载到cpu上，计算和激活值部分放到gpu上，这样比起跨机，更能节省内存，也能减少跨机跨通信域的通信压力
原理：

• 高计算：前向传播和反向传播计算量高，相关的权重参数计算和激活值计算仍然在gpu
• 低计算：权重更新部分计算量低，以通信为主，且需要的显存较大，放入到cpu中

deepspeed的使用

官方教程
感觉官网说的挺明白的，这里不再过多赘述~~~
有任何说的不对的地方，欢迎大家指正！！！

参考

• https://blog.csdn.net/myTomorrow_better/article/details/138945584
• https://www.bilibili.com/video/BV1fb421t7KN/?spm_id_from=333.788&vd_source=9b5bcbc75b2a02599faa4bac703679a1