Skip to content

DeepSpeed

微软的大规模模型训练优化库,通过 ZeRO 技术实现显存高效利用

Overview

DeepSpeed 是训练大语言模型的核心基础设施。辉少的面试笔记详细解释了 ZeRO(Zero Redundancy Optimizer)的三个阶段,以及如何在有限显存下训练超大规模模型。

Key Facts / Claims

显存消耗构成

混合精度训练(FP16/BF16)中: - 模型参数:FP16 格式权重 - 梯度:FP16 格式 - 优化器状态:FP32 动量、方差、主权重(占模型参数 3-4 倍)

ZeRO 三阶段

阶段 切分内容 显存节省 通信开销 适用场景
ZeRO-1 优化器状态 4x 大多数 LLM 微调首选
ZeRO-2 优化器+梯度 8x 显存稍紧缺时
ZeRO-3 优化器+梯度+参数 线性 N 倍 高(~50%) 超大模型预训练(70B+)

ZeRO-Offload

  • 将优化器状态和部分参数下放到 CPU RAM
  • 利用 CPU 进行权重更新
  • PCIe 带宽成为瓶颈,速度变慢

工程实践

  • ZeRO-1 最常用:性价比最高,通信开销几乎无增加
  • ZeRO-3 仅在必要时:70B+ 参数模型,用时间换空间
  • [[llm-rl-algorithms]] — 大模型训练中的 RL 阶段
  • [[grpo-global]] — 多 GPU 训练设置
  • [[transformer]] — 需要优化的基础架构
  • [[quantization]] — 另一项显存优化技术
  • [[flash-attention]] — 注意力计算优化

Counter-arguments & Data Gaps

  • DeepSpeed vs FSDP(PyTorch 原生)的对比
  • ZeRO-Infinity(NVMe 卸载)的适用场景
  • 通信优化:All-Reduce vs Reduce-Scatter + All-Gather

Sources