Loading...

文章背景图

DeepSpeed:微软开源的大规模分布式训练与推理优化库|部署与使用全指南

2026-11-01
0
- 字
- 分钟
|

图1 DeepSpeed 官方 Logo(官方素材)

DeepSpeed:微软开源的大规模分布式训练与推理优化库|部署与使用全指南

项目仓库: https://github.com/microsoft/DeepSpeed

开源协议:Apache 2.0 · 交付形态:Python 库(deepspeed PyPI 包)+ 命令行入口 deepspeed / ds_report · 主要语言:Python / C++ / CUDA · 部署方式:pip 安装或源码编译 · 无独立 Web 服务端口 · 官方站点:deepspeed.ai


1 项目简介

DeepSpeed 是微软研究院开源的深度学习训练与推理优化库,核心目标是让单卡显存装不下的大模型能够在多卡、多机乃至 CPU/NVMe 异构资源上高效训练。它以 PyTorch 扩展形式发布,最有代表性的技术是 ZeRO(Zero Redundancy Optimizer),把优化器状态、梯度、模型参数按分片策略切到不同 GPU,避免每张卡都冗余保存完整副本;后续延伸出 ZeRO-Offload、ZeRO-Infinity、ZeRO++、3D-Parallelism、Ulysses 序列并行、DeepSpeed-MoE、DeepCompile 等系统能力。

非官方 / 商业产品边界声明:DeepSpeed 由 Microsoft Research 主导开发,官方仓库已迁移至 deepspeedai/DeepSpeed 组织并由 Linux Foundation 协作治理;本文按原 microsoft/DeepSpeed 仓库口径整理。NVIDIA、AMD、华为昇腾、Intel Gaudi 等硬件名称仅用于说明支持范围,相关硬件与驱动归各自厂商所有。

1.1 核心信息一览

项目项 内容
项目名称 DeepSpeed(命令行入口 deepspeed,环境自检 ds_report)
仓库地址 github.com/microsoft/DeepSpeed(现迁移至 deepspeedai/DeepSpeed)
项目定位 大规模深度学习训练与推理系统优化库,核心 ZeRO 显存分片
开源协议 Apache 2.0
交付形态 PyPI 包 deepspeed,C++/CUDA/HIP 扩展 JIT 编译
主要语言 Python、C++、CUDA、HIP
部署方式 pip install deepspeed;Windows 用 build_win.bat 构建 wheel;可在 NVIDIA PyTorch 容器内使用
独立服务端口 无独立 Web 服务,作为 Python 库与 deepspeed 启动器使用
配置文件 ZeRO 等能力通过 JSON 配置文件声明(ds_config.json)
官方文档 deepspeed.ai · API 文档

1.2 适用场景

  • 单卡显存不足以放下大模型训练任务,需要把优化器状态、梯度、参数分片到多卡或多机。
  • 想用 CPU 内存、NVMe 作为显存扩展(ZeRO-Offload / ZeRO-Infinity),在少量高端 GPU 上跑百亿到千亿参数模型。
  • 在 Hugging Face Trainer、Megatron-LM 等训练框架里直接接入 DeepSpeed,替换原生 DDP 以获得更大 batch 与更低显存占用。
  • 推理侧希望通过 DeepSpeed-Inference、MoE 并行、量化等能力降低大模型在线推理的显存与时延。

1.3 能力边界(官方口径)

  • 官方 README 自述定位为 "Extreme Speed and Scale for DL Training",已被用于训练 Megatron-Turing NLG 530B、Jurassic-1 178B、BLOOM 176B、GLM 130B 等大规模模型。
  • 官方重点测试的 NVIDIA 架构为 Pascal、Volta、Ampere、Hopper;AMD 重点测试 MI100、MI200;华为昇腾 NPU、Intel Gaudi 2、Intel Xeon、Intel Data Center GPU Max、燧原 SDAA、昆仑芯 XPU 等由贡献方验证。
  • Windows 上支持大部分训练与推理能力,但官方 README 明确异步 IO(AIO)与 GDS(GPU Direct Storage)不支持 Windows。
  • DeepSpeed 本身不附带任何模型权重、数据集或商业服务;需要自行准备 PyTorch、CUDA 工具链、模型代码与训练脚本。

2 功能特性

2.1 ZeRO 显存分片

  • ZeRO-1 / 2 / 3:分别把优化器状态、梯度、模型参数按数据并行维度分片,避免每张 GPU 冗余保存完整副本;ZeRO-3 配合 stage3_gather_16bit_weights_on_model_save 可在保存时聚合权重。
  • ZeRO-Offload:把优化器状态与部分计算卸载到 CPU 内存,用 CPU 换 GPU 显存。
  • ZeRO-Infinity:进一步把参数卸载到 NVMe,用磁盘扩展显存;配合 NVMe SSD 可在少量高端 GPU 上跑千亿参数模型。
  • ZeRO++:在 ZeRO-3 基础上优化通信量,提升多机训练吞吐。

2.2 并行与长序列

  • 3D-Parallelism:数据并行 + 张量并行 + 流水线并行组合,DeepSpeed-AutoTP 可对 Hugging Face 模型自动做张量并行。
  • Ulysses Sequence Parallelism / Ulysses-Offload:把长序列沿序列维度切分,配合 Offload 降低长上下文训练显存。
  • DeepSpeed-MoE:为专家混合模型提供专家并行、通信与显存优化。

2.3 训练与推理系统能力

  • DeepCompile:编译器级优化,把分布式训练计算图做融合与调度。
  • DeepSpeed-Inference:面向推理的并行、量化、KV-Cache 优化。
  • DeepNVMe / SDMA:面向 NVMe I/O 与 AMD GPU 卸载通信的系统优化。
  • Muon Optimizer 支持:2026 年新增 Muon 优化器集成。

2.4 启动器与可观测性

  • deepspeed 命令行启动多卡/多机训练,替代 torchrun,自动读取 ds_config.json。
  • ds_report 输出当前机器上 PyTorch、CUDA、GPU、NCCL、已编译 ops 的环境报告,用于排查编译与硬件识别问题。
  • 提供通过 JSON 配置声明的梯度裁剪、混合精度、通信 bucket、日志、张量板、checkpoint 等能力。

3 官方界面预览

DeepSpeed 是 Python 库与命令行工具,没有图形界面。官方素材以 Logo 与技术架构图为主,本文仅使用官方 Logo 作为封面与开篇图;架构图请见官方站点 deepspeed.ai。

4 官方跳转链接汇总

5 前置要求与服务器配置

5.1 软件前提

  • 操作系统:Linux(Ubuntu / CentOS 等常见发行版)为主;Windows 官方支持大部分训练/推理能力,但 AIO 与 GDS 不可用。
  • Python:与所装 PyTorch 兼容的 3.x 版本。
  • PyTorch:必须在安装 DeepSpeed 之前装好;官方推荐 PyTorch ≥ 2.0,理想情况使用最新稳定版。
  • CUDA 工具链:nvcc(NVIDIA)或 hipcc(AMD ROCm),用于 JIT 编译 C++/CUDA/HIP 扩展。
  • 编译工具:ninja、gcc/g++、Python 开发头文件;Windows 需要 VS2022 C++ x64/x86 build tools。

5.2 硬件建议

  • NVIDIA:官方重点测试 Pascal、Volta、Ampere、Hopper 架构 GPU(如 A100、H100、V100、RTX 30/40 系列)。
  • AMD:官方重点测试 MI100、MI200。
  • 其他:华为昇腾 NPU、Intel Gaudi 2、Intel Xeon CPU、Intel Data Center GPU Max、燧原 SDAA、昆仑芯 XPU 由贡献方验证。
  • 跑 ZeRO-Infinity 时建议配置高速 NVMe SSD,作为显存扩展的二级存储。

5.3 端口与磁盘

  • DeepSpeed 本身不监听固定端口;多机训练时由 NCCL / MPI 走网卡通信,需要提前规划节点间通信端口与带宽(建议 RoCE / InfiniBand 或万兆以上以太网)。
  • 磁盘空间:模型权重、检查点、数据集与激活缓存按需预留;ZeRO-3 检查点通常比单卡副本大。

6 方式一:Docker / 容器化部署(推荐)

DeepSpeed 官方不单独提供 Docker 镜像,但推荐在 NVIDIA PyTorch 官方容器内安装,以保证 CUDA、cuDNN、NCCL 版本一致。

6.1 启动 PyTorch 容器

# 拉取与本机驱动兼容的 PyTorch CUDA 容器(示例,按实际 CUDA 版本选择 tag)
docker pull pytorch/pytorch:2.3.0-cuda12.1-cudnn8-devel

# 启动容器,挂载 GPU、宿主目录与数据集
docker run -it --rm --gpus all \
  --shm-size=16g \
  -v /data:/data \
  -v /path/to/your/train:/workspace \
  pytorch/pytorch:2.3.0-cuda12.1-cudnn8-devel bash

6.2 在容器内安装 DeepSpeed

# 容器内已经带 PyTorch,直接装 deepspeed
pip install deepspeed

# 验证安装并查看环境报告
ds_report

6.3 启动训练任务

# 单节点多卡(例如 8 张 GPU)
deepspeed --num_gpus=8 train.py --deepspeed --deepspeed_config ds_config.json

# 多节点(节点 0)
deepspeed --num_nodes=2 --num_gpus=8 \
  --hostfile /job/hostfile \
  train.py --deepspeed --deepspeed_config ds_config.json

远程访问注意事项:多机训练需要节点间 SSH 互通或在 Kubernetes / Slurm 等调度平台内拉起;不要把训练任务直接暴露在公网,建议通过堡垒机或内网访问。

7 方式二:源码 / 本地运行

7.1 克隆仓库

git clone https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed

7.2 安装可编辑模式

# 先确保 PyTorch 已装好
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

# 可编辑安装,便于调试
pip install -e .

7.3 Windows 特殊构建

官方 README 给出的 Windows 步骤:

  1. 安装 PyTorch,例如 pytorch 2.3+cu121。
  2. 安装 Visual C++ build tools,例如 VS2022 C++ x64/x86 build tools。
  3. 以管理员权限打开 Cmd,确保 MSVC 工具在 PATH 中,或直接使用 "Developer Command Prompt for VS 2022"(管理员)。
  4. 运行 build_win.bat,在 dist 目录生成 wheel。

7.4 验证

ds_report

输出应列出当前识别到的 GPU、CUDA 版本、NCCL、可用 ops 列表;若某行显示 "not supported",说明对应扩展在本机未编译通过,需要按官方进阶安装说明排查。

8 配置项说明

DeepSpeed 主要通过 ds_config.json 声明训练行为。以下字段均来自官方 JSON 配置文档口径,实际使用时以 官方 config-json 文档 为准。

字段 用途 常见值 / 说明
train_batch_size 全局 batch size 通常 = 微 batch × 梯度累积 × GPU 数
train_micro_batch_size_per_gpu 单卡微 batch 根据显存调整
gradient_accumulation_steps 梯度累积步数 用于在小卡上模拟大 batch
gradient_clipping 梯度裁剪阈值 例如 1.0
fp16.enabled / bf16.enabled 半精度训练开关 Hopper/Ampere 推荐 bf16
zero.stage ZeRO 阶段 0 / 1 / 2 / 3
zero.offload_optimizer.device 优化器卸载位置 cpu(ZeRO-Offload)
zero.offload_param.device 参数卸载位置 cpu / nvme(ZeRO-Infinity)
zero.stage3_gather_16bit_weights_on_model_save 保存时聚合 16 位权重 ZeRO-3 保存 checkpoint 时常用
zero.reduce_bucket_size / stage3_prefetch_bucket_size 通信 bucket 大小 影响显存与通信效率
activation_checkpointing 激活重计算 用算力换显存
wall_clock_breakdown 输出耗时分解 用于性能分析
tensorboard / wandb 日志后端 开启后自动上报指标

最小示例(ZeRO-2 + bf16):

{
  "train_batch_size": 128,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 4,
  "bf16": {"enabled": true},
  "zero_optimization": {"stage": 2},
  "gradient_clipping": 1.0
}

9 常用运维命令

# 环境自检:查看 GPU、CUDA、已编译 ops
ds_report

# 单节点多卡启动
deepspeed --num_gpus=8 train.py --deepspeed --deepspeed_config ds_config.json

# 多节点启动
deepspeed --num_nodes=2 --num_gpus=8 --hostfile hostfile train.py

# 查看正在运行的训练进程
nvidia-smi
ps -ef | grep train.py

# 实时监控 GPU 显存与利用率
watch -n 1 nvidia-smi

# 查看训练日志(示例,按实际重定向)
tail -f /path/to/train.log

# 停止训练
kill -INT <pid>   # 优先 SIGINT,让 DeepSpeed 落 checkpoint

更新 DeepSpeed:

pip install --upgrade deepspeed
ds_report

备份:ZeRO-3 检查点按分片保存,恢复时需要把所有分片放在同一目录并由 DeepSpeed 聚合;备份时应整体拷贝 checkpoint 目录与训练脚本、配置文件,避免丢分片。

10 常见故障排查

  • ds_report 中某 op 显示 "not supported":通常是 JIT 编译失败,确认 nvcc 版本与 PyTorch CUDA 版本一致、ninja 已安装、磁盘与 /tmp 有足够空间;按官方进阶安装说明预编译对应 ops。
  • CUDA out of memory:降低 train_micro_batch_size_per_gpu,开启 activation_checkpointing,把 ZeRO 阶段从 2 升到 3,或开启 offload_optimizer 到 CPU。
  • 多节点 NCCL 超时 / 通信失败:检查节点间 SSH、网卡、NCCL 相关环境变量(NCCL_SOCKET_IFNAME、NCCL_IB_DISABLE 等),确认防火墙放行 RoCE/IB 端口。
  • ZeRO-3 保存 checkpoint 后权重不对:确认保存时开启 stage3_gather_16bit_weights_on_model_save,或在保存脚本里调用 model.module.save_16bit_model()。
  • Windows 上构建失败:确认使用管理员权限的 Developer Command Prompt,且 VS2022 C++ build tools 已安装;AIO/GDS 在 Windows 上不可用属于官方限制,不是 bug。
  • 容器内看不到 GPU:启动容器时加 --gpus all,并在容器内执行 nvidia-smi 验证;宿主机需安装与容器 CUDA 版本兼容的 NVIDIA 驱动。

11 风险与合规说明

  • DeepSpeed 采用 Apache 2.0 协议,允许商业使用与二次开发,但需保留版权与许可声明;不得使用微软、DeepSpeed 商标做背书。
  • 训练任务涉及模型权重、数据集、checkpoint 时,相关版权与合规责任由使用者承担;第三方模型(如 LLaMA、Qwen、BLOOM 等)的许可按各自仓库条款执行。
  • 多机训练会占用大量 GPU、内存、网络与电力资源,建议在配额与审批过的集群内运行,避免在共享开发机上误跑满资源。
  • 本文所有命令、版本号、配置字段均来自官方 README 与 deepspeed.ai 文档;后续版本可能调整字段名与默认值,使用前请以官方仓库最新说明为准。
评论交流

文章目录