浅谈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） 🧩

• 🎯 目标