源码导读：从 deepspeed.initialize() 开始

13.1 顶层目录地图

DeepSpeed/
├── deepspeed/                       ← Python 主代码（本章重点）
│   ├── __init__.py                  入口，导出 initialize / init_inference / tp_model_init
│   ├── runtime/                     训练 runtime
│   │   ├── engine.py                DeepSpeedEngine（forward / backward / step / save）
│   │   ├── config.py                顶层 ds_config 解析
│   │   ├── bf16_optimizer.py        BF16 wrapper（loss scaling、master fp32）
│   │   ├── fp16/                    FP16 dynamic loss scaler
│   │   ├── zero/                    ZeRO 全套（stage_1_and_2.py、stage3.py、partitioned_param_coordinator.py、parameter_offload.py、offload_config.py、muon/）
│   │   ├── pipe/                    Pipeline 训练引擎（engine.py、schedule.py、topology.py、p2p.py）
│   │   ├── tensor_parallel/         DS 原生 TP（与 Megatron 互操作）
│   │   ├── sequence_parallel/       序列并行（Ulysses）
│   │   ├── comm/                    通信原语包装
│   │   ├── superoffload/            2026-03 SuperOffload
│   │   ├── zenflow/                 2025-10 选择性梯度更新
│   │   ├── domino/                  2024-11 异步 TP
│   │   ├── data_pipeline/           data efficiency / curriculum
│   │   ├── hybrid_engine.py         DS-Chat / RLHF 用
│   │   └── activation_checkpointing/
│   ├── inference/                   推理 runtime（v1 engine.py、v2/）
│   ├── ops/                         Python 封装的 fused op（adam / lamb / lion / muon / quantizer）
│   ├── moe/                         layer.py / sharded_moe.py / experts.py
│   ├── nvme/                        NVMe IO
│   ├── pipe/                        PipelineModule / LayerSpec 公开入口
│   ├── launcher/                    deepspeed CLI 启动器
│   ├── module_inject/               推理 kernel 替换（auto_tp.py、replace_policy.py）
│   ├── monitor/                     TensorBoard / WandB / CSV / Comet
│   ├── compression/                 量化 / 剪枝 / 蒸馏
│   ├── elasticity/                  弹性扩缩容
│   ├── compile/                     2025-04 DeepCompile（torch.compile 集成）
│   ├── sequence/                    Ulysses + ALST（fpdt_layer.py、auto_sp.py）
│   ├── checkpoint/                  shard 合并 / universal checkpoint
│   ├── linear/                      TP 化的 Linear
│   ├── autotuning/                  config 自动搜索
│   ├── utils/                       工具（含 zero_to_fp32.py 转换脚本）
│   └── ...
├── op_builder/                      C++/CUDA op 编译规则（每个 op 一个 .py）
├── csrc/                            CUDA / C++ 源码
├── tests/                           单元测试（unit/runtime/zero/ 是读懂 ZeRO 的好入口）
├── docs/                            API doc 源码
├── benchmarks/                      性能基准
├── blogs/                           官方博客（很多技术细节藏在这里）
├── docker/                          Dockerfile
├── requirements/                    各组依赖
└── setup.py                         package + ds_report 入口

13.2 三条主线代码

13.3 训练调用栈（rank 0 视角）

13.4 runtime/engine.py：DeepSpeedEngine 的 5 块

DeepSpeedEngine 这个类大约 3000+ 行，分块阅读：

13.5 runtime/zero/stage3.py：ZeRO-3 状态机

这是整个 DS 最难读的文件之一（~2500 行）。建议这样读：

1. 先读 PartitionedParamCoordinator：参数状态机 NOT_AVAILABLE / INFLIGHT / AVAILABLE；
2. 再读 DeepSpeedZeRoOffload：参数 offload 到 CPU / NVMe 的实现；
3. 再读 DeepSpeedZeroOptimizer_Stage3 的 backward / step：怎么触发 reduce-scatter；
4. 最后看 parameter_offload.py：hook 怎么挂的，何时 release 别人那 (N-1)/N 的参数。

13.6 launcher/：deepspeed CLI 的实现

13.7 op_builder/ + csrc/：怎么编出 fused kernel

想加一个自己的 fused op？模板：在 csrc/myop/ 写 .cu 和 binding，op_builder/myop.py 写编译规则，op_builder/__init__.py 注册到 ALL_OPS。

13.8 inference/：两套推理引擎

13.9 checkpoint/：从 ZeRO shard 到 HF 模型

13.10 阅读路线总结

"读完 DeepSpeed"是个庞大目标，建议这样推进：

1. 跑通最小例子（本书 ch03）：建立感觉；
2. 读 __init__.py + runtime/config.py：理解配置驱动；
3. 顺着 DeepSpeedEngine.__init__ 看一遍：理解初始化；
4. 挑 stage_1_and_2.py 读 backward / step：先看简单 stage；
5. 再啃 stage3.py + parameter_offload.py：ZeRO-3 状态机；
6. 顺手看一下 launcher/runner.py：理解多机；
7. 有 fused op 兴趣再看 op_builder/ + csrc/；
8. 动手提一个 PR（哪怕修 doc），跑通 contribution 流程。

13.11 想提 PR？从这里入手

13.12 延伸阅读

13.13 Stage 3 状态机的"读源码顺序"

ch05 已经讲清了 PartitionedParameterCoordinator 的语义。这里给一个最少行数的读源码路径——按顺序读，能完整理解 stage 3 一个 step 内 GPU 上的参数怎么动：

1. runtime/zero/partition_parameters.py:40–72 — ZeROOrderedDict.__getitem__，看入口拦截；
2. runtime/zero/parameter_offload.py:114 — DeepSpeedZeRoOffload.__init__，注入这个 dict、构造 Coordinator；
3. runtime/zero/partitioned_param_coordinator.py:44–50 — ZeRoTraceMode 枚举；
4. runtime/zero/partitioned_param_coordinator.py:64–149 — __init__ 与所有内部字段；
5. runtime/zero/partitioned_param_coordinator.py:159–196 — trace_prologue、状态切换逻辑；
6. runtime/zero/stage3.py:136–250 — DeepSpeedZeroOptimizer_Stage3.__init__，组装；
7. runtime/zero/stage3.py 内的 backward / step 实现 — 找 reduce_partitioned_gradients、independent_gradient_partition_epilogue；
8. 最后回头看 runtime/engine.py:2351 / 2615 / 2790 的 forward / backward / step，理解上层包装。

读完这 8 步，你对 stage 3 的脑中模型应该和源码 1:1。

13.14 op_builder ↔ csrc 对照

每个 op 的 Python 编译脚本在 op_builder/，实际 C++/CUDA 源码在 csrc/。两者一一对应：

注册中心：op_builder/all_ops.py:30 的 ALL_OPS = {op.name: op for op in __op_builders__}。新增 op 时往 __op_builders__ 里加自己的 builder 就会自动出现在 ds_report。

"自己加一个 op" 的步骤：

1. 在 csrc/myop/ 写 .cu + pybind11 binding；
2. 在 op_builder/myop.py 写一个 MyOpBuilder(OpBuilder) 子类，实现 sources()、include_paths()、absolute_name()；
3. 在 op_builder/__init__.py 把 MyOpBuilder 加进 __op_builders__；
4. 在 deepspeed/ops/myop/__init__.py 写 Python wrapper；
5. 跑 DS_BUILD_MYOP=1 pip install -e . 编一次。ds_report 里应该能看到 myop ............... [YES] ...... [OKAY]。

13.15 2024–2026 新子系统源码地图

这些目录是新的，老书 / 老 blog 里不会提，但都已经在 master：

13.16 跟着测试读 API

不知道一个新 feature 怎么用最快的方法：grep 它的测试。下面这几份测试堪称"API 教科书"：

读测试的好处：每个测试自带最小 config 与最小模型，不用先读上下游就能跑起来。怀疑某个 config 字段的真实行为时，去 tests/ 里 grep 字段名比看 doc 快。

13.17 zero_to_fp32.py 完整 CLI

这个脚本几乎是每个 DS 用户最终都会用一次的：把 ZeRO checkpoint 的多个 shard 合并成一份 PyTorch / HuggingFace 可加载的 fp32 / safetensors 权重。位置 deepspeed/utils/zero_to_fp32.py。完整签名：

python -m deepspeed.utils.zero_to_fp32 \
    <checkpoint_dir> \                  # ZeRO 保存的目录，里面有 global_step100/ 这种子目录
    <output_dir> \                       # 输出位置；得到 pytorch_model.bin 或 model.safetensors
    [--max_shard_size 5GB] \             # HF shard 切分大小（保存超大模型时用）
    [--safe_serialization] \             # 输出 safetensors 格式，不是 .bin
    [--tag global_step100] \             # 选用哪个 step 的 checkpoint（默认最后一个）
    [--exclude_frozen_parameters] \      # PEFT 场景下不导出冻结参数（只导 LoRA）
    [-d] [--debug]                       # 多打日志

关键函数（同文件内）：

• convert_zero_checkpoint_to_fp32_state_dict() — 公开 API；返回标准 PyTorch state_dict；
• _get_fp32_state_dict_from_zero_checkpoint() — 主体逻辑；
• parse_optim_states() — 从 shard 里读 fp32 master weights。

下游加载：

# HF 模型加载
model = AutoModelForCausalLM.from_pretrained("output_dir/")

# 或者裸 PyTorch
state = torch.load("output_dir/pytorch_model.bin")
model.load_state_dict(state)