3、DeepSpeed核心概念：ZeRO优化器、ZeRO-Offload、ZeRO-Infinity、3D并行策略
要理解 DeepSpeed 为什么能把千亿甚至万亿参数模型塞进有限的显存里，必须抓住四条主线：ZeRO 优化器解决“训练状态冗余”，ZeRO-Offload解决“单卡显存不够”，ZeRO-Infinity解决“连 CPU 内存也不够”，而3D 并行解决“单层都放不下”的极端规模。它们不是互斥选项，而是可以叠加组合的分层显存与并行体系。
一、核心原理：训练显存到底被谁吃掉了
在标准数据并行中，每张 GPU 都保存一份完整的模型副本，包括参数（Parameters）、梯度（Gradients）和优化器状态（Optimizer States）。以 Adam 混合精度训练为例，一个 fp16 参数背后往往还跟着 fp16 梯度、fp32 主权重、fp32 一阶动量、fp32 二阶动量，显存放大倍数可达 16 倍以上。ZeRO 的核心思想就是：既然每张卡都存一份是浪费，那就把这三类状态切分到不同 GPU 上，需要时再通过通信临时聚合。
ZeRO 分为三个阶段，逐级释放显存：
- ZeRO-1（优化器状态切分）：只切分优化器状态，显存节省约 4 倍，通信量与数据并行基本持平，性价比最高。
- ZeRO-2（梯度 + 优化器状态切分）：进一步切分梯度，显存节省约 8 倍，适合大多数中等规模训练。
- ZeRO-3（参数 + 梯度 + 优化器状态切分）：连参数也切分，显存节省与并行度成正比，但前向/反向需要 All-Gather 参数，通信量显著上升。
二、ZeRO-Offload 与 ZeRO-Infinity：把显存压力转移到 CPU 和 NVMe
当 ZeRO-3 仍然放不下时，DeepSpeed 提供了两级“外援”。ZeRO-Offload 把优化器状态和梯度计算卸载到 CPU 内存，利用 CPU 的多核与大容量内存承担 Adam 更新，GPU 只负责前向和反向。它适合单卡显存小、但主机内存充足的场景，代价是 PCIe 通信和 CPU 计算延迟。
ZeRO-Infinity 则更进一步，把参数、梯度、优化器状态扩展到 NVMe SSD，形成“GPU 显存 → CPU 内存 → NVMe”的三级存储层次。它通过带宽感知的预取和分块调度，让超大模型在有限硬件上也能训练。需要注意的是，Infinity 对存储带宽和延迟敏感，适合读多写少、可预取的训练模式。
三、3D 并行策略：三种并行如何协同
当模型大到单层都无法放入一张卡时，仅靠 ZeRO 已经不够，需要引入模型并行。张量并行（TP）把单个矩阵乘法切分到多卡，通信频繁但延迟低，适合节点内 NVLink 高带宽环境；流水线并行（PP）按层切分，通信量小但存在流水线气泡，适合跨节点；数据并行（DP）复制模型、切分数据，是吞吐量的主要来源。三者组合即 3D 并行。
DeepSpeed 的典型配置中，TP 通常限制在单节点 8 卡以内，PP 跨节点扩展，DP 负责剩余维度。ZeRO 可以与 3D 并行叠加：在 TP+PP 基础上，对 DP 维度应用 ZeRO-1/2/3，进一步降低每卡状态冗余。实际调参时，应优先保证 TP 通信在 NVLink 域内，PP 的 micro-batch 数量足够掩盖气泡，再用 ZeRO 压缩 DP 组内的显存。
四、案例与配置：一个可运行的组合示例
下面是一个同时启用 ZeRO-2、CPU Offload 与 3D 并行的 DeepSpeed 配置片段，展示了各概念在 JSON 中的落点：
{
"train_batch_size": 1024,
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"allgather_partitions": true,
"overlap_comm": true
},
"tensor_parallel": { "tp_size": 8 },
"pipeline": { "pp_size": 4, "micro_batch_size": 8 },
"gradient_accumulation_steps": 4
}该配置的含义是：在 8 卡张量并行、4 段流水线的基础上，对数据并行组启用 ZeRO-2，并把优化器状态卸载到 CPU。它适合单节点显存紧张、但主机内存充裕的中大规模训练。
五、常见误区与复盘
| 误区 | 实际影响 | 正确做法 |
|---|---|---|
| ZeRO 阶段越高越好 | ZeRO-3 通信量激增，吞吐可能下降 | 先试 ZeRO-1/2，显存不足再升阶段 |
| Offload 一定提速 | PCIe 与 CPU 计算成为瓶颈 | 仅在显存不足时启用，并开启 pin_memory |
| TP 可以跨节点随意扩展 | 跨节点 All-Reduce 延迟高，效率骤降 | TP 限制在 NVLink 域内，跨节点用 PP |
| 忽略 micro-batch 与气泡 | PP 流水线利用率低，GPU 空转 | 增加 micro-batch 数量，匹配梯度累积 |
总结来看，ZeRO 优化器负责切分训练状态，ZeRO-Offload 与 ZeRO-Infinity 负责把状态外溢到 CPU 和 NVMe，3D 并行负责把模型本身拆开。四者组合的核心逻辑始终是：用通信换显存，用层次化存储换规模，用并行维度换单层可放置性。理解这条主线，才能在真实训练任务中做出合理的配置取舍。