浅谈AI - DeepSpeed - 单卡慎用！

前言

曾在游戏世界挥洒创意，也曾在前端和后端的浪潮间穿梭，如今，而立的我仰望AI的璀璨星空，心潮澎湃，步履不停！愿你我皆乘风破浪，逐梦星辰！

简介

Deepspeed 的 ZeRO（Zero Redundancy Optimizer）有三个阶段（Stage 1, Stage 2, Stage 3），每个阶段旨在通过不同方式优化显存使用和提高计算效率。如果你是单卡训练，尤其是消费级显卡（如 3090、4090），启用 ZeRO Stage 2 和 3 反而得不偿失。以下是每个阶段的总结：

---

Stage 1：优化器状态分散（Optimizer State Partitioning） 🧠

• 🎯 目标：减少优化器状态的冗余存储。

• ⚙️ 优化方式：

  • 在多卡训练中，将优化器的状态（如梯度、动量等）分散到不同的设备上，避免每个设备存储冗余数据。

  • 每个 GPU 只负责存储优化器状态的一部分，减少显存占用。

• ✅ 适用场景：适合 多卡训练，因为它减少了多卡训练中的显存冗余。

• ❗ 单卡训练表现：在单卡训练中，Stage 1 并没有显著的优化效果，反而可能会增加内存开销，导致显存浪费。

---

Stage 2：模型参数和梯度分散（Model and Gradient Partitioning） 🧩

• 🎯 目标：减少每个 GPU 存储模型参数和梯度的冗余，进一步节省显存。

• ⚙️ 优化方式：

  • 将模型参数和梯度的不同部分分配到多个 GPU 上。

  • 每个 GPU 只存储模型的一部分参数和梯度，避免冗余存储。

• ✅ 适用场景：适合 多卡训练，尤其在大规模模型训练时，能够显著降低显存使用。

• ❗ 单卡训练表现：在单卡环境下，Stage 2 引入的优化并不会有太大优势，反而会增加额外的管理开销和内存使用。

---

Stage 3：完全分散所有训练状态（Fully Partitioned Training States） 🧬

• 🎯 目标：进一步优化显存，分散模型的所有训练状态（包括模型参数、梯度、优化器状态等）到多个设备上。

• ⚙️ 优化方式：

  • 将模型的所有参数、梯度和优化器状态完全分散到不同的设备上，从而消除所有冗余。

  • 极大地减少每个设备的显存负担，尤其在大规模分布式训练时非常有效。

• ✅ 适用场景：适用于 超大规模多卡训练，显著提高效率并节省显存。

• ❗ 单卡训练表现：在单卡训练中，Stage 3 的优化过于复杂，不仅不会带来明显好处，还可能导致性能下降和显存增加。

---

总结 📌

• Stage 1 ✅：适用于多卡训练，减少优化器状态的冗余存储。
  ❗ 单卡环境下并不推荐，可能带来额外显存浪费。

• Stage 2 ✅：适用于大模型的多卡训练，进一步降低显存需求。
  ❗ 单卡使用会有额外开销，不建议启用。

• Stage 3 ✅：用于极大模型的分布式训练，效果最显著。
  ❗ 单卡使用严重不推荐，可能拖慢训练并占用更多显存。

---

💡结论：Deepspeed 的不同阶段是为分布式训练设计的，如果你是单卡训练，尤其是消费级显卡（如 3090、4090），启用 ZeRO Stage 2 和 3 反而得不偿失，建议关闭 ZeRO 或仅使用混合精度（如 FP16/BF16）等其他方式进行显存优化。