跳到主要内容

数据截至 (上游 commit 32e301ffaf5a)

DeepSpeed — 架构与原理

30 秒导读： DeepSpeed 是微软开源（现由 deepspeedai 社区维护）的大模型训练优化库。它的心脏是 ZeRO（Zero Redundancy Optimizer，零冗余优化器）：普通数据并行里每张卡都存一份完整的优化器状态、梯度、参数，ZeRO 把这三样按数据并行 rank 切成 N 份，每张卡只留 1/N，需要时再通过集合通信临时凑齐。再配上 CPU/NVMe offload、混合精度、流水并行、MoE 和推理内核注入，它回答的问题是——「给定手头这批 GPU，到底能训多大的模型」。

1. 这是什么（零基础也能懂）​

一句话定义​

DeepSpeed 是一个包在 PyTorch 模型外面的训练运行时：你把 nn.Module 和一份 JSON 配置交给 deepspeed.initialize()，拿回一个 DeepSpeedEngine，之后 forward / backward / step 的显存占用、通信、混合精度、offload 全部由它接管。

它要解决谁的什么问题​

假设你有 8 张 80GB 的 GPU，想训一个 13B 的模型：

  • 用混合精度 + Adam 训一个 Ψ 参数的模型，每卡要放：fp16 参数 2Ψ 字节 + fp16 梯度 2Ψ + fp32 主权重 4Ψ + Adam 的动量 4Ψ + 方差 4Ψ，合计 16Ψ 字节（ZeRO 论文 arXiv:1910.02054 的经典账）。
  • 13B 参数就是 208GB——一张卡连「模型状态」都放不下，更别提激活值。
  • 普通数据并行（DDP）让每卡都存完整的一份，这 16Ψ 在每卡上重复了 N 遍，纯粹是浪费。

DeepSpeed 的 ZeRO 就是冲这份浪费来的：把三样状态按 rank 切片，每张卡只拥有 1/N，算的时候临时 all-gather / reduce-scatter 凑齐。这份账怎么算、代码怎么落，是第 1 章的主题。

它能做什么​

能力具体支持代码位置
ZeRO 数据并行Stage 1/2/3 三级状态分片（zero_optimization.stage）deepspeed/runtime/zero/stage_1_and_2.py、stage3.py
显存外移优化器状态/梯度/参数 offload 到 CPU 或 NVMe（ZeRO-Offload / ZeRO-Infinity）deepspeed/runtime/zero/offload_config.py、deepspeed/runtime/swap_tensor/
混合精度fp16（动态 loss scale）、bf16 优化器包装deepspeed/runtime/fp16/、bf16_optimizer.py
流水并行PipelineModule + PipelineEngine 的 1F1B 调度deepspeed/runtime/pipe/
MoE专家并行 + all-to-all 路由deepspeed/moe/
推理内核注入（kernel injection）替换 transformer 层、AutoTP 张量并行推理deepspeed/inference/、deepspeed/module_inject/
长序列Ulysses 序列并行、激活重计算deepspeed/runtime/sequence_parallel/、activation_checkpointing/

用起来什么样​

最小形态：一个 JSON 配置 + 三行训练循环改造。

# 示意，非源码（摘自官方 README 用法的通用形态）
import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,                       # 带 --deepspeed_config ds_config.json
    model=model,
    model_parameters=model.parameters())

for batch in data_loader:
    loss = model_engine(batch)       # forward：引擎代管
    model_engine.backward(loss)      # backward：梯度被桶化、reduce-scatter
    model_engine.step()              # step：只更新本地分片，再 all-gather 补齐

配套的最小 ZeRO-2 配置（ds_config.json）：

{
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 4,
  "fp16": {"enabled": true},
  "zero_optimization": {"stage": 2},
  "optimizer": {"type": "AdamW", "params": {"lr": 1e-4}}
}

入口是 deepspeed/__init__.py:93 的 initialize()：初始化分布式、解析配置、按模型类型造出 DeepSpeedEngine（或 PipelineEngine），返回四元组 (engine, optimizer, dataloader, lr_scheduler)。

一句话直觉​

把 ZeRO 想成「合租」。 以前每个室友（GPU）都在自己屋里放一整套家具（参数+梯度+优化器状态）；ZeRO 之后大家合伙买一套，每人保管 1/N，谁要用哪件就喊一声（all-gather）临时搬过来，用完立刻还回去（release / partition）。屋子（显存）省了，代价是楼道（网络）上多了搬来搬去的脚步声。

2. 顶层全景（它大概怎么转）​

2.1 部件拓扑​

用户脚本
  │  deepspeed.initialize(model, config=ds_config.json)     deepspeed/__init__.py:93
  ▼
┌────────────────────────────────────────────────────────────┐
│ DeepSpeedEngine (runtime/engine.py:252)                    │
│  ├─ DeepSpeedConfig：JSON → 对象 (runtime/config.py:692)   │
│  └─ 优化器包装层 _configure_optimizer (engine.py:2028)     │
│       │  按 zero stage 分流 (engine.py:2366)               │
│       ├─ Stage 1/2 → DeepSpeedZeroOptimizer                │
│       │     (zero/stage_1_and_2.py:135)：flatten→切片→     │
│       │     梯度桶 reduce-scatter→本地 step→all-gather     │
│       └─ Stage 3   → DeepSpeedZeroOptimizer_Stage3         │
│           (zero/stage3.py:149) + DeepSpeedZeRoOffload      │
│           (zero/parameter_offload.py:130)：hook 驱动       │
│           参数 all-gather 预取/释放                        │
└────────────────────────────────────────────────────────────┘
        │ offload 配置 (zero/offload_config.py)
        ▼
  CPU Adam (ops/adam/cpu_adam.py)  ←→  NVMe swap (runtime/swap_tensor/)

怎么读这张图： 竖着看是「一次初始化的装配顺序」，横着看 DeepSpeedEngine 内部是「配置驱动优化器选择」。stage 1/2 与 stage 3 是两条几乎独立的实现路线——前者动「优化器看到的参数」，后者动「模型参数本身的存在形态」。

2.2 部件一句话职责​

部件干什么在哪个文件
initialize()入口：初始化分布式、建 engine、返回四元组deepspeed/__init__.py:93
DeepSpeedConfig把 JSON 解析成带校验的对象（pydantic 子配置 + 手写 getter 混合）deepspeed/runtime/config.py:692
DeepSpeedEngine包住用户模型的 nn.Module：forward/backward/step/梯度累积边界deepspeed/runtime/engine.py:334
DeepSpeedZeroOptimizerZeRO-1/2 优化器：分片优化器状态（+梯度），桶化通信deepspeed/runtime/zero/stage_1_and_2.py:136
DeepSpeedZeroOptimizer_Stage3ZeRO-3 优化器：参数也分片，step 按 sub-group 轮换deepspeed/runtime/zero/stage3.py:150
DeepSpeedZeRoOffloadZeRO-3 的「参数调度器」：给每个模块挂 hookdeepspeed/runtime/zero/parameter_offload.py:130
PartitionedParameterCoordinatorZeRO-3 的调度中枢：trace 记录、预取、释放deepspeed/runtime/zero/partitioned_param_coordinator.py:73
Init（zero.Init）上下文管理器：让参数一出生即分片，省掉整模型落地deepspeed/runtime/zero/partition_parameters.py:940
AsyncPartitionedParameterSwapper参数级 NVMe 换入换出（AIO/GDS）deepspeed/runtime/swap_tensor/partitioned_param_swapper.py:37
DeepSpeedCPUAdam向量化 CPU 版 Adam/AdamW，offload 时计算在 CPU 上跑deepspeed/ops/adam/cpu_adam.py:13
PipelineEngine流水并行引擎（DeepSpeedEngine 子类）deepspeed/runtime/pipe/engine.py:60
InferenceEngine推理引擎：dtype 转换、内核注入、TP 组deepspeed/inference/engine.py:43

2.3 主线走一遍（ZeRO-2 的一个训练步）​

以最常见的 ZeRO-2 为例，高层动作序列是：

① engine(batch) 前向
   和普通 PyTorch 一样；参数全员在位（stage ≤2 不分参数）
        │
② engine.backward(loss)
   每个参数的 grad 一就绪，hook 就把它拷进连续梯度桶
   桶满 → reduce-scatter：每 rank 只收「归自己更新」的 1/N 段
   （其余分片的梯度当场释放——这就是 stage 2 省下的显存）
        │
③ engine.step() 到达累积边界时
   本 rank 只对自己那段 fp32 分片跑 Adam（优化器状态也只有 1/N）
   更新结果写回 fp16 分片 → all-gather 把新参数凑齐到所有 rank
        │
   下一步前向，又是一套完整参数

这条线最值得记住的一点： stage 1/2 里「参数」从头到尾是完整的，只有梯度通信的落点变了——从「每卡收全量」变成「每卡只收自己负责的那段」。通信总量没变，显存占用降了。stage 3 更进一步，把「临时凑齐」推到了参数上，那是第 3 章的故事。

3. 阅读地图（建议顺序）​

五章由浅入深。时间有限的话，读 01 → 03 就能抓住 DeepSpeed 的全部要害。

顺序章节讲什么适合谁
101-zero-stages-1-2.md显存账（16Ψ→16Ψ/N）、摊平切片、梯度桶、step 的全部动作所有人必读，ZeRO 的地基
202-engine-and-config.mdinitialize 装配链、配置体系、forward/backward/step 骨架、梯度累积要接入/调试 DeepSpeed 的人
303-zero-stage-3.mdzero.Init、hook 机制、trace 驱动的预取/释放、stage3 的 step想懂「参数也分片」的人
404-offload.mdCPU/NVMe offload：pin memory、AIO swapper、流水线式优化器状态交换显存实在不够、想买时间换空间的人
505-pipeline-moe-inference.md流水并行调度、MoE all-to-all、推理内核注入、ZeRO++ 等速览需要这些特性时再来查

4. 巧妙之处（可借鉴的技术）​

  1. 摊平再切，一切通信按 bucket 走。 每个参数组先被 _flatten_dense_tensors 摊成一条连续 fp16 向量，再按 DP rank 数等分（deepspeed/runtime/zero/stage_1_and_2.py:1976 get_data_parallel_partitions）。参数边界因此不再重要，通信和内存分配都以固定大小的桶为单位。

  2. 一行代码换掉优化器的整个世界。 初始化时执行 param_group['params'] = [self.single_partition_of_fp32_groups[i]]（deepspeed/runtime/zero/stage_1_and_2.py:518）——真优化器（FusedAdam 等）从此只看到本 rank 那 1/N 分片，优化器状态自动变成 1/N，Adam 代码一行不用改。这是 stage 1 的全部机关。

  3. 梯度边产生边走，不等 backward 结束。 每个参数注册 grad hook，梯度一就绪就拷进 reduce_bucket_size 大小的连续桶，桶满即 reduce-scatter（deepspeed/runtime/zero/stage_1_and_2.py:1162 create_gradient_handling_hooks + :1692 reduce_ipg_grads）。通信被摊进整个反向过程，而不是末尾一次大同步。

  4. ZeRO-3 用「第一圈记录、以后照演」解决预取。 第一轮 forward/backward 记录模块执行顺序（trace），之后每轮按这个队列提前 all-gather 下一步要用的参数（deepspeed/runtime/zero/partitioned_param_coordinator.py:251 reset_step）。不用静态分析计算图，也不用用户标注，却拿到了全局的取用顺序。

  5. 参数一出生即分片。 zero.Init 用猴子补丁挂钩所有 nn.Module 的子类构造，每个参数刚被 __init__ 造出来就被就地切成 1/N（deepspeed/runtime/zero/partition_parameters.py:334 InsertPostInitMethodToModuleSubClasses）。造 1T 模型时每个进程只需要装 1/N 的 CPU 内存，而不是整份。

  6. 释放也要看「复用距离」。 ZeRO-3 释放参数前向后看 max_reuse_distance 个元素：马上要复用的不释放，省一次 all-gather（deepspeed/runtime/zero/partitioned_param_coordinator.py:634 __params_to_release）。

5. 边界与局限​

  • ZeRO-2/3 与流水并行互斥。 PipelineEngine.__init__ 里硬断言 zero_optimization_stage < ZeroStageEnum.gradients（deepspeed/runtime/pipe/engine.py:77）。要 3D 并行请走 Megatron 路线，或 ZeRO stage 1 + pipeline。
  • 配置面巨大，多数问题是配置问题。 源卡片原话：「Configuration is a large JSON surface; most problems are config problems, not code problems.」zero_optimization 下有几十个旋钮（deepspeed/runtime/zero/config.py:90 DeepSpeedZeroConfig），默认值之间还存在联动（如 overlap_comm 的默认值随 stage 变化，:383 overlap_comm_valid）。
  • ZeRO-3 要求模块执行序跨 rank 一致且逐步稳定。 trace 发现实际模块顺序和缓存不符就整体失效重录（deepspeed/runtime/zero/partitioned_param_coordinator.py:202 trace_prologue）；带条件分支的动态图会让预取反复失效，性能掉回无预取水平。
  • 造模型本身可能先 OOM。 若 fp16 权重单卡都放不下，必须用 zero.Init 包着构造（Init docstring 明确写了这条，deepspeed/runtime/zero/partition_parameters.py:940 起）。
  • Muon 优化器与若干 ZeRO 特性互斥。 reduce_scatter + offload 的组合会直接 ValueError（deepspeed/runtime/zero/stage_1_and_2.py:233、deepspeed/runtime/zero/stage3.py:353）。
  • offload 是吞吐换显存。 梯度/优化器状态走 PCIe（NVMe 还要再走 SSD），官方定位就是「用通信和 I/O 带宽换模型规模」；具体掉多少吞吐取决于硬件，代码里看不出固定数字（这是诚实的留白）。

6. 横向对比​

同书架上还有四个「训练侧」兄弟，取舍各不相同：

项目与 DeepSpeed 的关系核心取舍
Megatron-LM同为大规模训练库，可叠加（Megatron-DeepSpeed）Megatron 押注张量并行+流水并行的 3D 切分；DeepSpeed 押注数据并行内的状态分片（ZeRO）
PyTorchDeepSpeed 跑在 PyTorch 之上；FSDP 是 ZeRO-3 思路的原生实现FSDP 与 PyTorch 运行时深度整合；DeepSpeed 提供更全的一揽子（offload、MoE、推理）
AccelerateHF 的薄封装，可以把 DeepSpeed 当后端托管Accelerate 管「启动与设备摆放」，重活仍交给 DeepSpeed/FSDP
verlRL 后训练框架，训练后端用 FSDP/Megatron 而不用 ZeROverl 解决「训推同集群编排」，ZeRO 解决「单模型显存」；两个问题正交

一句话：要纯数据并行的显存杠杆、又不想碰张量并行改造模型，DeepSpeed 的 ZeRO 是这个书架上最直接的答案。

7. 代码地图（入口级）​

每章末尾有细粒度地图，这里只列从零开始读源码的五个入口：

主题文件路径符号名
入口与装配deepspeed/__init__.pyinitialize
引擎本体deepspeed/runtime/engine.pyDeepSpeedEngine、forward、backward、step
配置解析deepspeed/runtime/config.py、deepspeed/runtime/zero/config.pyDeepSpeedConfig、DeepSpeedZeroConfig、ZeroStageEnum
ZeRO-1/2deepspeed/runtime/zero/stage_1_and_2.pyDeepSpeedZeroOptimizer、step、average_tensor
ZeRO-3deepspeed/runtime/zero/stage3.py、partitioned_param_coordinator.pyDeepSpeedZeroOptimizer_Stage3、PartitionedParameterCoordinator