DeepSpeed 学习指南

DeepSpeed 代码库学习指南适合希望深入理解 DeepSpeed 内部机制的工程师与研究者。目录项目定位与核心价值整体架构分层目录结构详解核心模块深度导览4.1 入口与初始化4.2 DeepSpeedEngine — 训练引擎4.3 ZeRO — 显存优化系列4.4 混合精度优化器4.5 流水线并行4.6 序列并行 / Ulysses4.7 MoE混合专家4.8 推理引擎4.9 通信层4.10 算子库与 CUDA 扩展4.11 Accelerator 抽象层4.12 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 # 公开 APIinitialize、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_configJSON 文件或 dict根据配置决定使用哪种 EngineDeepSpeedEngine/PipelineEngine/DeepSpeedHybridEngine包装 model、optimizer、lr_scheduler调用init_distributed()初始化进程组engine,optimizer,_,lr_schedulerdeepspeed.initialize(modelmodel,optimizeroptimizer,configds_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/ZeROZero Redundancy Optimizer是 DeepSpeed 最核心的技术创新阶段分片内容文件Stage 1优化器状态stage_1_and_2.pyStage 2 梯度stage_1_and_2.pyStage 3 模型参数stage3.pyZeRO-InfinityStage 3 NVMe offloadstage3.pyswap_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.pyfp16/ ├── 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.py4.7 MoE混合专家目录deepspeed/moe/moe/ ├── layer.py # MoE 层gate 专家路由 ├── sharded_moe.py # 分片 MoEDeepSpeed-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 distcomm/ ├── 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-optionbuild_ext预编译4.11 Accelerator 抽象层目录accelerator/顶层 deepspeed/accelerator/将 CUDA 特定 API 抽象为统一接口支持 CUDA / Intel XPU / Huawei Ascend NPU / AMD ROCm / Apple MPS 等。fromdeepspeed.acceleratorimportget_accelerator accget_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控制{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→DeepSpeedConfigZeRO 专属配置deepspeed/runtime/zero/config.py→DeepSpeedZeroConfig6. 测试体系tests/ ├── unit/ # 单元测试最常用 │ ├── runtime/ # 训练引擎相关 │ ├── zero/ # ZeRO 测试 │ ├── pipe/ # 流水线并行 │ ├── comm/ # 通信 │ ├── ops/ # 算子 │ └── ... ├── benchmarks/ # 性能基准 └── model/ # 模型集成测试运行单元测试# 运行特定模块的测试pytest tests/unit/runtime/-v# 运行 ZeRO 相关测试pytest tests/unit/zero/-v-kstage2重要提示大多数分布式测试需要 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 TeamC 算子通过op_builder/动态编译首次运行会较慢可通过DS_BUILD_*1环境变量预编译大量使用instrument_w_nvtx/see_memory_usage()做性能分析这些不影响逻辑可略过9. 关键概念速查表概念说明核心文件DeepSpeedEngine包装 PyTorch model统一 train/eval 接口runtime/engine.pyZeRO Stage 1/2分片优化器状态和梯度减少数据并行冗余runtime/zero/stage_1_and_2.pyZeRO Stage 3额外分片模型参数按需 all-gatherruntime/zero/stage3.pyZeRO-InfinityStage 3 NVMe offload突破 CPU 内存瓶颈runtime/swap_tensor/ZenFlow流水线化 offload消除 GPU 等待runtime/zenflow/PipelineEngine流水线并行将模型按层切分到多 GPUruntime/pipe/engine.pyUlysses序列维度并行支持百万 token 长序列sequence/layer.pyDeepSpeed-MoE稀疏 MoE支持万亿参数模型moe/InferenceEngine推理优化kernel 融合 量化 TPinference/engine.pyAccelerator硬件抽象层屏蔽 CUDA/XPU 差异accelerator/op_builderCUDA 扩展动态编译管理op_builder/FP16_OptimizerFP16 训练 FP32 master weightsruntime/fp16/fused_optimizer.pyLoss Scaler动态 loss scaling 防止 FP16 下溢runtime/fp16/loss_scaler.pyactivation_checkpointing重计算激活值用计算换显存runtime/activation_checkpointing/commtorch.distributed 统一封装comm/comm.pygroups各维度并行进程组管理utils/groups.py学习过程中遇到不理解的代码可结合blogs/目录下的对应技术博客以及 deepspeed.ai 官方文档共同理解。