DeepSpeed：微软分布式大模型训练优化库安装与使用指南
官方信息
| 项目名称 | DeepSpeed |
| 仓库地址 | https://github.com/deepspeedai/DeepSpeed |
| 官方网站 | https://www.deepspeed.ai |
| 最新更新时间 | 2026-07-27 |
| Star 数 | 34k+ |
| 项目定位 | 微软开源的大模型分布式训练与推理优化框架，ZeRO 技术发明者 |
项目背景
DeepSpeed 是微软研究院开源的深度学习训练优化库，它解决了一个核心问题：当模型参数从数十亿增长到数千亿时，单张显卡甚至一台 8 卡服务器都装不下完整的模型权重、梯度和优化器状态。DeepSpeed 提出的 ZeRO（Zero Redundancy Optimizer）技术通过将优化器状态、梯度、模型参数在多卡之间分片存储，把训练大模型的显存需求降低了一个数量级。
从 GPT-3 级别的模型训练到如今的 Llama-3、Qwen-72B，几乎所有开源大模型的预训练和微调都离不开 DeepSpeed。它不仅支持训练，还提供了 DeepSpeed-Inference 推理引擎、MoE（混合专家）模型支持、FP8 量化训练等高级特性。HuggingFace Transformers、LLaMA-Factory、Swift 等主流微调框架都内置了 DeepSpeed 集成。
对于普通开发者而言，DeepSpeed 的价值在于：不需要自己写复杂的分布式通信代码，只需要一个 JSON 配置文件，就能让单卡训练脚本在多机多卡上高效运行。
核心特性
- ZeRO 内存优化：ZeRO-1/2/3 三个阶段，分别分片优化器状态、梯度、模型参数，最大可将显存占用降低到 1/N（N 为 GPU 数）。
- ZeRO-Offload：把优化器和参数卸载到 CPU 内存甚至 NVMe SSD，让单张消费级显卡也能微调 13B-70B 模型。
- 混合精度训练：原生支持 FP16/BF16/FP8，自动 loss scaling，训练稳定性大幅提升。
- 激活检查点：选择性重计算前向激活，以计算换显存，支持任意深度的模型。
- DeepSpeed-Inference：面向推理场景的模型并行、KV 缓存优化、MoE 专家并行。
- 通信优化：支持 NVLink、InfiniBand、RoCE，自动选择最佳集合通信算法。
- 与 PyTorch 无缝集成：只需几行代码改造，兼容 HuggingFace Trainer。
环境要求
| 项目 | 最低要求 | 推荐配置 |
|---|---|---|
| 操作系统 | Linux（Ubuntu 20.04+） | Ubuntu 22.04 LTS |
| Python | 3.8 | 3.10 / 3.11 |
| GPU | NVIDIA GPU, Volta 架构以上 | A100 / H100 / 4090 |
| CUDA | 11.6 | 12.1 / 12.4 |
| PyTorch | 1.12+ | 2.1+ |
| 内存 | 32GB（ZeRO-Offload 需要更大） | 128GB+ |
详细部署步骤
第一步：环境准备
# 创建虚拟环境
conda create -n deepspeed python=3.10 -y
conda activate deepspeed
# 安装 PyTorch
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# 验证 GPU
nvidia-smi
第二步：安装 DeepSpeed
方式一：pip 安装（推荐）
# 基础安装
pip install deepspeed
# 验证安装，会自动检查环境
ds_report
方式二：源码编译安装（需要自定义 CUDA 算子时）
git clone https://github.com/deepspeedai/DeepSpeed.git
cd DeepSpeed
# 设置编译选项，开启所有加速算子
export DS_BUILD_OPS=1
export DS_BUILD_CPU_ADAM=1
export DS_BUILD_FUSED_ADAM=1
export DS_BUILD_TRANSFORMER=1
export DS_BUILD_STOCHASTIC_TRANSFORMER=1
pip install -e .[dev]
方式三：Docker 镜像
docker pull deepspeed/deepspeed:latest
docker run -it --gpus all --shm-size=10g --network=host \
-v /your/data:/workspace \
deepspeed/deepspeed:latest bash
第三步：编写 ZeRO 配置文件
创建 ds_config_zero3.json，这是最常用的 ZeRO-3 配置：
{
"train_micro_batch_size_per_gpu": 2,
"gradient_accumulation_steps": 8,
"gradient_clipping": 1.0,
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"initial_scale_power": 16,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"bf16": {
"enabled": "auto"
},
"steps_per_print": 100,
"wall_clock_breakdown": false
}
配置解释：
stage: 3：ZeRO-3，分片所有参数+梯度+优化器状态，显存最省但通信量最大。offload_optimizer.device: cpu：把优化器状态卸载到 CPU 内存。offload_param.device: cpu：把模型参数也卸载到 CPU（适合显存极小的场景）。bf16.enabled: auto：自动检测是否使用 BF16（A100/4090 推荐）。
如果你的显卡比较充足（如 8×A100），可以用更轻量的 ZeRO-2：
{
"zero_optimization": {
"stage": 2,
"offload_optimizer": {"device": "cpu", "pin_memory": true},
"contiguous_gradients": true,
"overlap_comm": true
},
"bf16": {"enabled": true}
}
第四步：在 PyTorch 脚本中集成 DeepSpeed
改造你的训练脚本（以 PyTorch 原生为例）：
import deepspeed
import torch
from torch.utils.data import DataLoader
# 初始化
model = YourModel()
train_dataset = YourDataset()
# 关键：用 deepspeed.initialize 包装
model_engine, optimizer, trainloader, lr_scheduler = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
training_data=train_dataset,
config="ds_config_zero3.json"
)
# 训练循环和普通 PyTorch 几乎一样
for batch in trainloader:
outputs = model_engine(batch)
loss = outputs.loss
# 注意：用 model_engine.backward(loss) 代替 loss.backward()
model_engine.backward(loss)
# 用 model_engine.step() 代替 optimizer.step()
model_engine.step()
第五步：用 HuggingFace Trainer 一键启用
如果你用 HuggingFace Transformers，只需要在 TrainingArguments 里加一行：
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(
output_dir="./output",
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
learning_rate=1e-4,
num_train_epochs=3,
bf16=True,
# 关键：指定 DeepSpeed 配置
deepspeed="ds_config_zero3.json",
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
)
trainer.train()
第六步：启动多卡训练
# 单机多卡（2 张 GPU）
deepspeed --num_gpus=2 train.py
# 单机多卡（使用所有 GPU）
deepspeed --num_gpus=`nvidia-smi -L | wc -l` train.py
# 多机多卡（节点 0）
deepspeed --hostfile /job/hostfile --num_nodes=2 --node_rank=0 \
--master_addr=192.168.1.100 --master_port=29500 train.py
# 节点 1
deepspeed --hostfile /job/hostfile --num_nodes=2 --node_rank=1 \
--master_addr=192.168.1.100 --master_port=29500 train.py
hostfile 格式示例：
# /job/hostfile
worker-0 slots=8
worker-1 slots=8
第七步：验证训练和查看日志
训练启动后，你会看到类似输出：
Setting ds_accelerator to cuda (cuda)
[2026-07-27 10:00:00] [INFO] [config.py:1000:print] DeepSpeedEngine config:
{
"train_batch_size": 32,
"train_micro_batch_size_per_gpu": 2,
"gradient_accumulation_steps": 8,
"zero_optimization": {"stage": 3, ...}
}
用 watch -n1 nvidia-smi 实时查看显存占用，ZeRO-3 + Offload 下每张卡通常只用 10-20GB。
常见坑与解决
- NCCL timeout：多机训练时网络慢会导致超时，设置
export NCCL_P2P_DISABLE=0，并确保节点间能互通 29500 端口。 - 编译 C++ 算子失败：确保
nvcc --version和 PyTorch 自带的 CUDA 版本一致；实在不行用DS_BUILD_OPS=0 pip install deepspeed跳过编译。 - ZeRO-3 保存的模型无法直接加载：保存时设置
stage3_gather_16bit_weights_on_model_save: true，或用zero_to_fp32.py脚本合并权重。 - 显存仍然不够：开启
activation_checkpointing，减小train_batch_size_per_gpu，或把offload_param也开到 cpu。
使用方法示例
训练完成后，检查点目录结构：
output/
├── global_step100/
│ ├── mp_rank_00_model_states.pt
│ ├── zero_pp_rank_0_optim_states.pt
│ ├── zero_pp_rank_1_optim_states.pt
│ └── ...
合并 ZeRO-3 权重为标准 HuggingFace 格式：
python -c \
"from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint; \
load_state_dict_from_zero_checkpoint('output/global_step100', 'merged_model')"
合并后的 merged_model/ 目录可直接被 vLLM、Transformers 加载。
常见问题 FAQ
Q1：ZeRO-1/2/3 怎么选？
A：8 卡 A100 训练 7B-13B 用 ZeRO-2 足够；70B 以上或显存紧张用 ZeRO-3；只有 1-2 张消费级显卡用 ZeRO-3 + Offload。
Q2：DeepSpeed 和 FSDP 哪个好？
A：两者功能接近。DeepSpeed 生态更成熟、文档更全、社区案例多；FSDP 是 PyTorch 原生，未来可能更主流。新手建议先学 DeepSpeed。
Q3：Offload 到 CPU 会慢多少？
A：通常慢 20%-50%，但显存占用降低 60% 以上。如果内存够大（256GB+），这是单卡微调大模型的唯一选择。
Q4：支持 Windows 吗？
A：官方不支持 Windows 原生，建议用 WSL2 或 Docker。
Q5：DeepSpeed 能用于推理吗？
A：可以，DeepSpeed-Inference 支持模型并行推理，但生产环境更推荐用 vLLM 或 TensorRT-LLM。
跳转链接
- GitHub 仓库：https://github.com/deepspeedai/DeepSpeed
- 官方文档：https://www.deepspeed.ai/getting-started/
- ZeRO 配置示例：configs 目录
- HuggingFace DeepSpeed 集成：官方教程