DeepSpeed 学习指南

DeepSpeed 代码库学习指南

适合希望深入理解 DeepSpeed 内部机制的工程师与研究者。

项目定位与核心价值
整体架构分层
目录结构详解
核心模块深度导览
- 4.1 入口与初始化
- 4.2 [DeepSpeedEngine --- 训练引擎](#DeepSpeedEngine — 训练引擎)
- 4.3 [ZeRO --- 显存优化系列](#ZeRO — 显存优化系列)
- 4.4 混合精度优化器
- 4.5 流水线并行
- 4.6 [序列并行 / Ulysses](#序列并行 / Ulysses)
- 4.7 MoE（混合专家）
- 4.8 推理引擎
- 4.9 通信层
- 4.10 [算子库与 CUDA 扩展](#算子库与 CUDA 扩展)
- 4.11 [Accelerator 抽象层](#Accelerator 抽象层)
- 4.12 [ZenFlow --- 无停顿 Offload](#ZenFlow — 无停顿 Offload)
配置系统
测试体系
推荐的阅读路径
阅读技巧与注意事项
关键概念速查表

1. 项目定位与核心价值

DeepSpeed 是 Microsoft 开源的大规模分布式深度学习训练与推理优化框架。其核心价值：

能力	技术手段
训练超大模型（单机放不下）	ZeRO Stage 1/2/3、ZeRO-Infinity
加速训练	融合 CUDA 算子、混合精度、流水线并行
降低显存占用	参数/梯度/优化器状态卸载至 CPU/NVMe
高效推理	量化、张量并行、内核融合
易用性	只需在已有 PyTorch 代码中调用 `deepspeed.initialize()`

2. 整体架构分层

复制代码

┌─────────────────────────────────────────────────────┐
│                   用户 Training Script               │
│         model, optimizer, dataloader, config        │
└───────────────────────┬─────────────────────────────┘
                        │ deepspeed.initialize()
                        ▼
┌─────────────────────────────────────────────────────┐
│              DeepSpeedEngine  (runtime/engine.py)   │
│  ┌──────────┐  ┌──────────┐  ┌────────────────────┐ │
│  │  ZeRO    │  │   FP16/  │  │  Pipeline Parallel │ │
│  │ Stage1-3 │  │  BF16    │  │  (PipelineEngine)  │ │
│  │ Optimizer│  │ Optimizer│  └────────────────────┘ │
│  └──────────┘  └──────────┘                         │
│  ┌──────────────────────────────────────────────┐   │
│  │         Activation Checkpointing             │   │
│  └──────────────────────────────────────────────┘   │
└────────────────────────┬────────────────────────────┘
                         │
        ┌────────────────┼──────────────────┐
        ▼                ▼                  ▼
┌──────────────┐  ┌──────────────┐  ┌────────────────┐
│ comm layer   │  │  ops / csrc  │  │  Accelerator   │
│ (deepspeed.  │  │  (CUDA/CPU   │  │  Abstraction   │
│  comm)       │  │   kernels)   │  │  (CUDA/XPU/HPU)│
└──────────────┘  └──────────────┘  └────────────────┘

3. 目录结构详解

复制代码

DeepSpeed/
├── deepspeed/               # 主包，所有 Python 逻辑
│   ├── __init__.py          # 公开 API（initialize、dist 等）
│   ├── runtime/             # ★ 训练核心：引擎、ZeRO、优化器
│   │   ├── engine.py        # DeepSpeedEngine（最核心文件，~4500行）
│   │   ├── config.py        # 配置解析
│   │   ├── zero/            # ZeRO Stage 1/2/3 实现
│   │   ├── fp16/            # FP16 混合精度优化器
│   │   ├── bf16_optimizer.py
│   │   ├── pipe/            # 流水线并行引擎
│   │   ├── tensor_parallel/ # 张量并行初始化工具
│   │   ├── sequence_parallel/
│   │   ├── zenflow/         # ZenFlow 无停顿 offload
│   │   ├── superoffload/    # SuperOffload
│   │   ├── swap_tensor/     # NVMe tensor swap
│   │   ├── activation_checkpointing/
│   │   └── model_checkpointing/
│   ├── inference/           # 推理引擎
│   ├── comm/                # 通信抽象层（封装 torch.distributed）
│   ├── ops/                 # Python 层算子入口（指向 csrc）
│   ├── moe/                 # MoE 层实现
│   ├── sequence/            # Ulysses 序列并行层
│   ├── pipe/                # PipelineModule（模型定义侧）
│   ├── module_inject/       # 推理时替换 Transformer 层
│   ├── launcher/            # deepspeed CLI 启动器
│   ├── autotuning/          # 自动超参调优
│   ├── checkpoint/          # 检查点工具（ZeRO/Universal CP）
│   ├── accelerator/         # 硬件抽象（CUDA/XPU/HPU 等）
│   ├── utils/               # 通用工具（logging、timer、groups）
│   └── linear/              # LoRA 优化线性层
├── csrc/                    # ★ C++/CUDA 算子实现
│   ├── adam/                # CPU Adam、Fused Adam
│   ├── transformer/         # Transformer 融合算子
│   ├── quantization/        # 量化 kernel
│   ├── aio/                 # 异步 NVMe I/O
│   └── ...
├── op_builder/              # PyTorch Extension 构建系统
├── tests/
│   ├── unit/                # 单元测试（按模块分目录）
│   └── ...
├── docs/                    # 文档与博客
└── blogs/                   # 技术博客（附代码分析）

4. 核心模块深度导览

4.1 入口与初始化

文件：deepspeed/__init__.py

最重要的 API 是 deepspeed.initialize()，它：

解析 ds_config（JSON 文件或 dict）
根据配置决定使用哪种 Engine（DeepSpeedEngine / PipelineEngine / DeepSpeedHybridEngine）
包装 model、optimizer、lr_scheduler
调用 init_distributed() 初始化进程组

python 复制代码

engine, optimizer, _, lr_scheduler = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config="ds_config.json"
)

阅读建议 ：先看 initialize() 函数体，追踪 DeepSpeedEngine.__init__() 的调用链。

4.2 DeepSpeedEngine --- 训练引擎

文件：deepspeed/runtime/engine.py（约 4500 行，最核心）

职责：

封装 torch.nn.Module，提供 forward() / backward() / step()
根据配置挂载 ZeRO 优化器、FP16/BF16 优化器
管理梯度累积、梯度裁剪
管理 checkpoint 保存与加载

关键方法阅读顺序：

复制代码

__init__()
  → _configure_optimizer()       # 决定用哪个优化器
  → _configure_zero_optimizer()  # ZeRO 包装
  → _configure_fp16_optimizer()  # FP16 包装

train_batch()                    # 完整训练步骤
  → forward()
  → backward()
  → step()

4.3 ZeRO --- 显存优化系列

目录：deepspeed/runtime/zero/

ZeRO（Zero Redundancy Optimizer）是 DeepSpeed 最核心的技术创新：

阶段	分片内容	文件
Stage 1	优化器状态	`stage_1_and_2.py`
Stage 2	+ 梯度	`stage_1_and_2.py`
Stage 3	+ 模型参数	`stage3.py`
ZeRO-Infinity	Stage 3 + NVMe offload	`stage3.py` + `swap_tensor/`

关键文件：

复制代码

zero/
├── stage_1_and_2.py        # Stage1/2 优化器（~3000行）
│   └── DeepSpeedZeroOptimizer  # 核心类
├── stage3.py               # Stage3（~3500行）
│   └── DeepSpeedZeroOptimizer_Stage3
├── partition_parameters.py # 参数分区与 all-gather 逻辑
├── parameter_offload.py    # 参数 CPU offload
├── offload_config.py       # offload 配置模型
├── partitioned_param_coordinator.py  # Stage3 参数调度
└── config.py               # ZeRO 配置解析

阅读建议：

先读 zero/config.py 理解配置项
再读 stage_1_and_2.py 的 DeepSpeedZeroOptimizer.__init__() 和 step()
最后攻坚 stage3.py（需要理解 partition_parameters.py 中的参数生命周期）

4.4 混合精度优化器

目录：deepspeed/runtime/fp16/、bf16_optimizer.py

复制代码

fp16/
├── fused_optimizer.py      # FP16_Optimizer：维护 FP32 master weights
├── unfused_optimizer.py    # 非融合版本
├── loss_scaler.py          # 动态/静态 loss scaling
└── onebit/                 # 1-bit Adam/LAMB

核心逻辑：FP16 前向/反向 → FP16 梯度 → 动态 loss scale → unscale → FP32 优化器步骤 → 拷贝回 FP16 参数。

4.5 流水线并行

目录：deepspeed/runtime/pipe/ + deepspeed/pipe/

复制代码

runtime/pipe/
├── engine.py      # PipelineEngine，继承 DeepSpeedEngine
├── module.py      # 流水线调度逻辑
├── schedule.py    # 1F1B 调度策略
├── p2p.py         # 流水线间 P2P 通信
└── topology.py    # 3D 并行拓扑（数据/模型/流水线维度）

pipe/
└── PipelineModule  # 用户用来定义流水线模型的接口

阅读建议 ：先看 topology.py 理解进程组划分，再看 schedule.py 的 1F1B 调度，最后看 engine.py 如何驱动调度。

4.6 序列并行 / Ulysses

目录：deepspeed/sequence/、deepspeed/runtime/sequence_parallel/

实现 DeepSpeed-Ulysses：在序列维度切分注意力头，通过 All-to-All 通信替代 All-Reduce，大幅降低长序列训练通信量。

复制代码

sequence/
├── layer.py        # UlyssesAttention：核心 all-to-all 切分逻辑
├── fpdt_layer.py   # Flash Attention + Ulysses 融合版
└── cross_entropy.py

4.7 MoE（混合专家）

目录：deepspeed/moe/

复制代码

moe/
├── layer.py        # MoE 层（gate + 专家路由）
├── sharded_moe.py  # 分片 MoE，DeepSpeed-MoE 核心
├── experts.py      # 专家模块封装
└── mappings.py     # 专家并行通信映射

4.8 推理引擎

目录：deepspeed/inference/

复制代码

inference/
├── engine.py       # InferenceEngine
├── config.py       # 推理配置
└── v2/             # 新版推理引擎（模块化设计）

module_inject/ 负责在推理时将标准 Transformer 层替换为 DeepSpeed 融合 kernel 实现。

4.9 通信层

目录：deepspeed/comm/

DeepSpeed 对 torch.distributed 做了统一封装。规则：

永远不要直接 import torch.distributed，始终使用 import deepspeed.comm as dist

复制代码

comm/
├── comm.py     # 统一接口（all_reduce、broadcast、all_gather 等）
├── torch.py    # torch.distributed 后端适配
├── ccl.py      # Intel CCL 后端
├── backend.py  # 后端注册
└── config.py   # 通信配置（log 等）

utils/groups.py 管理数据并行、模型并行、流水线并行等各类进程组。

4.10 算子库与 CUDA 扩展

目录：deepspeed/ops/（Python 入口） + csrc/（C++/CUDA 实现） + op_builder/（构建系统）

复制代码

ops/
├── adam/           # CPU Adam、Fused Adam (CUDA)
├── transformer/    # Transformer 融合算子
├── quantizer/      # 量化算子
├── aio/            # 异步 NVMe I/O
└── lion/           # Lion 优化器

csrc/
├── adam/           # CPU/GPU Adam kernel
├── transformer/    # Attention/FF 融合
├── quantization/   # 量化 CUDA kernel
└── aio/            # libaio 异步 I/O

op_builder/         # 每个算子对应一个 Builder 类
                    # 负责动态/AOT 编译

构建模式：

JIT 编译（默认）：首次调用时自动编译
AOT 编译 ：pip install deepspeed --global-option="build_ext" 预编译

4.11 Accelerator 抽象层

目录：accelerator/（顶层） + deepspeed/accelerator/

将 CUDA 特定 API 抽象为统一接口，支持 CUDA / Intel XPU / Huawei Ascend NPU / AMD ROCm / Apple MPS 等。

python 复制代码

from deepspeed.accelerator import get_accelerator
acc = get_accelerator()
acc.device_name()    # "cuda" / "xpu" / ...
acc.synchronize()

移植新硬件只需继承 abstract_accelerator.py 中的 DeepSpeedAccelerator。

4.12 ZenFlow --- 无停顿 Offload

目录：deepspeed/runtime/zenflow/

ZenFlow 是 DeepSpeed 2025 年新引入的 Offload 引擎，通过流水线化的方式消除 CPU-GPU 数据搬运造成的 GPU 停顿。

复制代码

zenflow/
├── zenflow_stage_1_and_2.py   # ZenFlow 版 Stage1/2 优化器
├── engine_stage3.py           # Stage3 集成
├── zenflow_config.py          # 配置项
└── zenflow_utils.py           # 工具函数

5. 配置系统

DeepSpeed 的全部特性通过 JSON 配置文件（ds_config.json）控制：

json 复制代码

{
  "train_batch_size": 32,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" }
  },
  "gradient_clipping": 1.0,
  "pipeline": { "stages": 4 }
}

配置解析入口 ：deepspeed/runtime/config.py → DeepSpeedConfig

ZeRO 专属配置 ：deepspeed/runtime/zero/config.py → DeepSpeedZeroConfig

6. 测试体系

复制代码

tests/
├── unit/                   # 单元测试（最常用）
│   ├── runtime/            # 训练引擎相关
│   ├── zero/               # ZeRO 测试
│   ├── pipe/               # 流水线并行
│   ├── comm/               # 通信
│   ├── ops/                # 算子
│   └── ...
├── benchmarks/             # 性能基准
└── model/                  # 模型集成测试

运行单元测试：

bash 复制代码

# 运行特定模块的测试
pytest tests/unit/runtime/ -v

# 运行 ZeRO 相关测试
pytest tests/unit/zero/ -v -k "stage2"

重要提示 ：大多数分布式测试需要 GPU，可用 @pytest.mark.skip 在无 GPU 环境跳过。

7. 推荐的阅读路径

路径一：理解用户接口（入门级，1-2天）

复制代码

1. README.md                        # 了解项目目标和功能概述
2. blogs/                           # 选读 1-2 篇技术博客（推荐 ZeRO、Ulysses）
3. deepspeed/__init__.py            # initialize() 函数签名
4. deepspeed/runtime/config.py      # 配置项全貌
5. tests/unit/runtime/test_ds_initialize.py  # 通过测试理解用法

路径二：深入 ZeRO（核心技术，3-5天）

复制代码

1. deepspeed/runtime/zero/config.py          # 配置
2. deepspeed/runtime/zero/offload_config.py  # offload 配置
3. deepspeed/runtime/engine.py               # _configure_zero_optimizer()
4. deepspeed/runtime/zero/stage_1_and_2.py   # Stage1/2 主循环
   → __init__(), reduce_gradients(), step()
5. deepspeed/runtime/zero/partition_parameters.py  # 参数分区
6. deepspeed/runtime/zero/stage3.py          # Stage3
   → __init__(), _pre_forward_hook(), _post_backward_hook()
7. deepspeed/runtime/zero/parameter_offload.py  # CPU offload

路径三：流水线并行（模型并行，2-3天）

复制代码

1. deepspeed/pipe/__init__.py               # PipelineModule 使用方式
2. deepspeed/runtime/pipe/topology.py       # 进程组
3. deepspeed/runtime/pipe/schedule.py       # 1F1B 调度
4. deepspeed/runtime/pipe/p2p.py            # P2P 通信
5. deepspeed/runtime/pipe/engine.py         # PipelineEngine

路径四：新增硬件支持（加速器移植）

复制代码

1. accelerator/abstract_accelerator.py      # 接口定义
2. accelerator/cuda_accelerator.py          # 参考实现
3. op_builder/builder.py                    # 算子构建基类

8. 阅读技巧与注意事项

技巧

从 engine.py 出发 ：几乎所有训练特性都在 DeepSpeedEngine.__init__() 中被条件性地启用，它是最好的"地图"。
跟踪 _configure_* 系列方法 ：如 _configure_optimizer()、_configure_zero_optimizer()、_configure_lr_scheduler()，这些方法串联起各模块。
善用测试文件 ：tests/unit/ 下的测试是最直接的使用示例，比文档更新、更准确。
通过 ds_config.json 理解特性边界 ：每一个配置项背后对应一段 C/P 分支逻辑，在 config.py 中找到配置解析，再追踪使用它的代码。
关注 comm 调用 ：分布式逻辑的核心在 dist.all_reduce()、dist.all_gather() 等调用处，重点阅读这些周边逻辑。
看博客辅助理解 ：blogs/ 目录下的 Markdown 文章解释了算法动机，是理解代码 "为什么这样写" 的关键。

注意事项

禁止直接使用 torch.distributed ，代码中统一用 import deepspeed.comm as dist
新文件必须有 License 头：# SPDX-License-Identifier: Apache-2.0 + # DeepSpeed Team
C++ 算子通过 op_builder/ 动态编译，首次运行会较慢，可通过 DS_BUILD_*=1 环境变量预编译
大量使用 @instrument_w_nvtx / see_memory_usage() 做性能分析，这些不影响逻辑，可略过

9. 关键概念速查表

概念	说明	核心文件
`DeepSpeedEngine`	包装 PyTorch model，统一 train/eval 接口	`runtime/engine.py`
`ZeRO Stage 1/2`	分片优化器状态和梯度，减少数据并行冗余	`runtime/zero/stage_1_and_2.py`
`ZeRO Stage 3`	额外分片模型参数，按需 all-gather	`runtime/zero/stage3.py`
`ZeRO-Infinity`	Stage 3 + NVMe offload，突破 CPU 内存瓶颈	`runtime/swap_tensor/`
`ZenFlow`	流水线化 offload，消除 GPU 等待	`runtime/zenflow/`
`PipelineEngine`	流水线并行，将模型按层切分到多 GPU	`runtime/pipe/engine.py`
`Ulysses`	序列维度并行，支持百万 token 长序列	`sequence/layer.py`
`DeepSpeed-MoE`	稀疏 MoE，支持万亿参数模型	`moe/`
`InferenceEngine`	推理优化：kernel 融合 + 量化 + TP	`inference/engine.py`
`Accelerator`	硬件抽象层，屏蔽 CUDA/XPU 差异	`accelerator/`
`op_builder`	CUDA 扩展动态编译管理	`op_builder/`
`FP16_Optimizer`	FP16 训练 + FP32 master weights	`runtime/fp16/fused_optimizer.py`
`Loss Scaler`	动态 loss scaling 防止 FP16 下溢	`runtime/fp16/loss_scaler.py`
`activation_checkpointing`	重计算激活值，用计算换显存	`runtime/activation_checkpointing/`
`comm`	torch.distributed 统一封装	`comm/comm.py`
`groups`	各维度并行进程组管理	`utils/groups.py`

学习过程中遇到不理解的代码，可结合 blogs/ 目录下的对应技术博客，以及 deepspeed.ai 官方文档共同理解。