项目背景与定位
1.1 一句话定位
DeepSpeed 是微软研究院开源的、用"内存优化"做差异化竞争的分布式深度学习框架。它的招牌技术 ZeRO（Zero Redundancy Optimizer） 通过把优化器状态、梯度、参数依次切分到不同 GPU 上，让"原本必须 8 张 A100 才能塞下的模型，现在 2 张就够"成为可能。
如果你近几年用过 HuggingFace Transformers 训过模型，那十有八九你已经在间接用 DeepSpeed：HF Trainer 的 --deepspeed 启动参数走的就是它。
1.2 为什么会有 DeepSpeed
2019-2020 年那个时间点，业界面对的问题非常具体：
- GPT-3 出现，让"百亿到千亿参数"突然成了主流而不是实验室项目；
- 工业界很多团队有钱买卡，但没多到能买 1024 张的程度；
- PyTorch DDP 是"复制模型 + 同步梯度"的范式，显存使用上完全没做优化：N 张卡训一个 10B 模型，每张都需要 10B 的优化器状态、10B 的梯度、10B 的参数。
微软的 Samyam Rajbhandari 等人在 ZeRO 论文（2019）里给出的洞察非常朴素：这种"全卡都存一份"的模式，里面有 $N-1$ 份是重复的，可以切。Adam 优化器一份参数需要保存的额外状态约 4 倍参数量（fp32 master weights + 一阶矩 $m$ + 二阶矩 $v$），所以单卡显存占用（DP 设备数 $N$）大致为：
$$M_\text{DDP} = \underbrace{2\Psi}_{\text{参数}} + \underbrace{2\Psi}_{\text{梯度}} + \underbrace{12\Psi}_{\text{Adam 状态(fp32)}} = 16\Psi$$
其中 $\Psi$ 是参数量。ZeRO 把这 $16\Psi$ 中越来越多的部分切到 $N$ 张卡：
$$M_{\text{ZeRO-1}} = 4\Psi + \frac{12\Psi}{N}, \quad M_{\text{ZeRO-2}} = 2\Psi + \frac{14\Psi}{N}, \quad M_{\text{ZeRO-3}} = \frac{16\Psi}{N}$$
ZeRO-3 是真正的"线性可扩展"——卡越多每张占用越小。
| 阶段 | 切分对象 | 说明 |
|---|---|---|
| ZeRO-1 | 优化器状态 | Adam 的 $m/v$ 切到 $N$ 张卡 |
| ZeRO-2 | + 梯度 | 反向 reduce-scatter 而不是 all-reduce |
| ZeRO-3 | + 参数本身 | 每次 forward / backward 再 all-gather |
DeepSpeed 就是把这套"渐进式分片"工程化的产物。后来项目继续向外延伸：DeepSpeed-Inference、DeepSpeed-MoE、DeepSpeed-Chat、ZeRO-Offload / ZeRO-Infinity 等。
1.3 它和 PyTorch FSDP、Megatron-LM 是什么关系
| 维度 | DeepSpeed | PyTorch FSDP | Megatron-LM |
|---|---|---|---|
| 主战场 | ZeRO 分片 + 一切扩展 | PyTorch 原生 ZeRO-3 等价物 | TP/PP/SP 模型并行 |
| 切分对象 | 优化器/梯度/参数 | 参数 + 优化器 | 张量本身（按权重矩阵切） |
| 集成方式 | 包一个 DeepSpeedEngine | 包一个 FullyShardedDataParallel | 重写模型代码 |
| 配置形式 | JSON | Python API + auto_wrap | argparse / yaml |
| 易用性 | 中（JSON 字段多） | 高（PyTorch 原生） | 低（强工程化） |
经验法则：
- 想"基本不改业务代码、显存够紧"用 DeepSpeed ZeRO-2/3；
- 想"完全 PyTorch 原生、不引第三方"用 FSDP；
- 想"训千亿 + 极致性能 + 接受高复杂度"用 Megatron-LM。
它们不互斥——Megatron-DeepSpeed 就是把 DeepSpeed 的 ZeRO 引擎嵌进 Megatron 的并行框架里。
1.4 设计哲学：JSON 即一切
DeepSpeed 最让初学者又爱又恨的设计是：绝大多数关键参数都通过一个 JSON 文件配置，而不是 Python API。比如这样的一份典型 config：
{
"train_batch_size": 256,
"gradient_accumulation_steps": 4,
"bf16": { "enabled": true },
"zero_optimization": {
"stage": 3,
"offload_optimizer": { "device": "cpu" },
"overlap_comm": true
}
}好处是非常直观；坏处是这份 JSON 一旦超过 30 行，字段之间的相互约束就开始让人头疼——zero stage 3 不能和某些 fused optimizer 同用，offload 不能和某些 dtype 同用等等。第 7 章会给一份完整的"字段地图"。
1.5 它适合用来做什么、不适合做什么
| 判断 | 典型场景 |
|---|---|
| ✅ 适合 | 显存不够、想训得动比硬件理论上能塞下的更大的模型 |
| 已经在用 HF Trainer，希望最小代价获得分布式能力 | |
| 大规模 SFT、继续预训练、RLHF（DeepSpeed-Chat） | |
| 想在普通服务器 + NVMe 上挑战训练超大模型（ZeRO-Infinity） | |
| ❌ 不适合 | 极致 throughput 优先、有完整 NVIDIA H100 集群的预训练 —— Megatron-LM / NeMo 通常更快 |
| 想用消费级 4090 / 单卡微调 —— Unsloth / LoRA 更轻 | |
| 不想碰 JSON，想全 Python API 控制 —— FSDP / Accelerate 更顺手 |
1.6 本书将怎么带你走完这趟旅程
- 第 2-3 章先把环境搭起来，跑通 HF Trainer + DeepSpeed 的最小例子；
- 第 4 章回答"
deepspeed.initialize()之后到底发生了什么"； - 第 5-7 章深入 ZeRO 三个阶段，配上时序图说明 all-gather / reduce-scatter 何时发生，并把 config 字段彻底过一遍；
- 第 8-9 章讲它怎么和 HF 集成、怎么和 Megatron 拼出 3D 并行、怎么训 MoE；
- 第 10-12 章是工程师每天会用上的：推理、调优、踩坑；
- 第 13 章给一张源码导读图。
1.7 DeepSpeed 在 2024–2026 引入了什么
DeepSpeed 不是一个停下来的项目。打开 repos/deepspeed/README.md 的 "Latest News" 一段，几乎每个季度都有新东西落地。下面这张倒序时间线是本书基线 commit 510ebe58e4e4（2026-05-25）当时已经合入 master 的特性 —— 全部能在源码里找到，本书后续章节会逐一对应到 file:line。
| 月份 | 特性 | 是什么 | 本书章节 |
|---|---|---|---|
| 2026-05 | Muon 优化器 + SDMA for ZeRO-3 | 新接入的二阶 LLM 优化器；SDMA 直传通道用于 stage 3 通信加速 | ch05.14 / ch07.15 / ch12.13 |
| 2026-03 | SuperOffload（ASPLOS Honorable Mention） | 不止 offload 内存，把"优化器计算"也搬到 CPU worker 进程，重叠率更高 | ch05.15 / ch06.10 |
| 2025-12 | Core API 现代化 | 支持 PyTorch 风格 backward、low-precision master weight | ch04.12 |
| 2025-10 | ZenFlow 选择性梯度更新 | top-k 重要梯度立刻 GPU 更新，其余在 CPU 异步累积，stall-free offload | ch05.14 / ch07.15 / ch12.12 |
| 2025-06 | Arctic Long Sequence Training + DeepNVMe | 百万级 token 上下文训练；NVMe 直读直写优化 | ch09.14 / ch06.11 |
| 2025-04 | DeepCompile | 把 torch.compile 和 DS 的 ZeRO / offload / activation 管理打通 | ch11.12 |
| 2025-03 | AutoTP for HuggingFace | 不需要用户写 policy，自动扫 HF 模型并做张量并行切分 | ch09.13 |
| 2024-12 | Ulysses-Offload | Ulysses 序列并行 + offload 组合 | ch09.14 |
| 2024-11 | Domino 异步 TP | 列/行并行 Linear 的 allreduce 推迟到 backward，藏通信 | ch09.15 |
读完这张表，应该建立一个心智模型：DS 的"经典 ZeRO"还在演进，但很多 2024 后的新工作不是替换 ZeRO，而是在 ZeRO 的 hook / offload / 通信路径上加挂新功能。本书 ch05、ch06、ch09 的"进阶"小节就是逐个拆开这些挂件。