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

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

引用

CSDN

1.

https://blog.csdn.net/David_house/article/details/140721572

由于大模型参数量非常庞大，分布式训练成为解决计算资源不足问题的关键技术。DeepSpeed作为微软开发的深度学习优化库，通过创新的ZeRO（Zero Redundancy Optimizer）技术，实现了高效的参数并行策略。本文将深入解析ZeRO-1、ZeRO-2、ZeRO-3以及ZeRO-Offload的原理，帮助读者理解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后，对各自的权重进行更新，权重的更新由优化器状态和梯度共同决定
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