CodeWalk

DeepSpeed的核心特性与引擎架构

作者:小字辈 · 2026-05-30 12:55

DeepSpeed是微软开源的大模型训练优化库。请介绍其ZeRO优化器、DeepSpeed Engine的架构设计,以及在ChatGPT类模型训练中的典型配置(ZeRO Stage/Offload/混合精度)。

回答

小字辈

DeepSpeed是Microsoft DeepSpeed团队为PyTorch构建的分布式训练优化库,核心围绕三个模块。

1. ZeRO优化器

  • Stage 1/2/3 逐步消除冗余存储(如前所述)
  • ZeRO-Offload:利用CPU内存扩展显存
  • ZeRO-Infinity:扩展到NVMe

2. DeepSpeed Engine架构: 用户模型 → DeepSpeed Engine注入 → 自动优化:自动ZeRO包装、自动混合精度设置、自动梯度累积、自动通信优化、自动Checkpointing

使用方式即替换优化器和Trainer:model_engine, optimizer, _, _ = deepspeed.initialize(model, optimizer, ...)

3. 典型ChatGPT(SFT/RLHF)配置: ZeRO Stage 3 + Offload + FP16混合精度 + 梯度累积

4. 其他特性

  • Sparse Attention(减少长序列计算)
  • 1-bit Adam(梯度压缩通信)
  • Progressive Layer Dropping(随机跳过部分层加速)
  • ZeRO++(通信量化)

DeepSpeed现被广泛用于LLaMA-Factory、HuggingFace Trainer中。