deepspeedai/DeepSpeed · 上手攻略

• 仓库：deepspeedai/DeepSpeed
• 链接：https://github.com/deepspeedai/DeepSpeed
• 分类：ai（llm-infra）
• 作者：Jay
• 更新：2026-07-11

是什么

DeepSpeed 是微软开源的深度学习优化库，专为大规模分布式训练和推理而设计。它让训练千亿、万亿参数模型不再是巨头的专利，普通研究团队也能在有限硬件上跑起来。

DeepSpeed 的核心创新是 ZeRO（Zero Redundancy Optimizer） 系列技术，通过分片优化器状态、梯度、参数，大幅降低单 GPU 显存占用。此外还支持 3D 并行（数据并行 + 张量并行 + 流水线并行）、ZeRO-Infinity（支持把参数卸载到 CPU/NVMe）、DeepSpeed-MoE（专家混合模型）、FP8 混合精度训练等特性。

在历史上，DeepSpeed 支撑了 MT-530B（当时全球最大语言模型）和 BLOOM（1760 亿参数）等大模型的训练。

解决什么问题

• 显存不够：单卡跑不动大模型，ZeRO 把优化器状态分片，70B 模型可以从需要 80GB+ 显存降低到数 GB
• 多卡通信效率低：1-bit Adam、ZeRO-DRAM 等技术减少多卡通信量
• 上手门槛高：兼容 PyTorch Lightning、Hugging Face Transformers、FastChat 等主流框架，加载 DeepSpeed 只需改几行代码
• 推理也慢：DeepSpeed-Inference 支持在 GPU 上高效推理大模型

快速安装

# 基础安装（JIT 编译，无需预编译）
pip install deepspeed

# 验证安装 & 查看支持的 ops
ds_report
# 或
python -m deepspeed.env_report

# 预编译所有 C++/CUDA 扩展（加速后续使用）
DS_BUILD_OPS=1 pip install deepspeed

# 从源码安装（不含预编译）
git clone https://github.com/deepspeedai/DeepSpeed.git
cd DeepSpeed
pip install .

⚠️ 依赖：PyTorch 必须先安装。DeepSpeed 通过 torch.utils.cpp_extension 实现 JIT 编译，需要 ninja 构建工具（JIT 模式下会自动处理）。建议 Python 3.8–3.11，PyTorch ≥1.9。

核心用法

1. 用 DeepSpeed 训练（单 GPU 示例）

创建配置文件 ds_config.json：

{
  "train_batch_size": 8,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 1e-4
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 1
  }
}

在训练脚本中加入：

import deepspeed

# 用 DeepSpeed 初始化模型
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config="ds_config.json"
)

for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

2. ZeRO Stage 选择指引

3. Hugging Face Transformers 用户

# 训练
ds_config = {
    "zero_optimization": {"stage": 2},
    "fp16": {"enabled": True}
}

trainer = transformers.Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    # DeepSpeed 插件会自动处理分布式
)
trainer.train()

4. DeepSpeed 推理

import deepspeed

# 推理引擎初始化
ds_engine = deepspeed.init_inference(
    model=model,
    tp_size=1,  # 张量并行大小
    dtype=torch.float16
)
output = ds_engine.generate(input_tokens)

5. 单行代码让训练快 30%

FP8 混合精度（需 H100/A100 8+ GPU）：

{
  "fp8": {"enabled": true},
  "gradient_clipping": 1.0
}

📌 注意：FP8 训练需要 NVIDIA H100/A100（Compute Capability ≥ 9.0）和对应 CUDA 版本。DeepSpeed 会自动检测硬件兼容性，不满足条件时降级为 BF16。

典型适用场景

• 大模型预训练：LLaMA、BLOOM、Megatron 等模型的高效分布式训练
• 模型微调：Lora、RLHF 等下游任务降低显存需求
• 推理服务：DeepSpeed-Inference 加速 vLLM 等推理框架
• 多模态训练：与 Lightning 集成后支持图像、语言等多种模型
• 科研场景：买不起 8×A100？用 ZeRO-3 在消费级 GPU 上跑 7B 模型

坑与注意

1. JIT 编译失败：部分 CUDA 版本不兼容。遇到 Unable to build DeepSpeed C++ extensions 时，建议使用 DS_BUILD_OPS=1 手动预编译，或确认 PyTorch 和 CUDA 版本匹配
2. ZeRO-3 + 流水线并行冲突：ZeRO-3 目前与流水线并行有已知兼容问题，建议单独使用
3. Windows 支持有限：训练和推理基本支持，但 AIO 和 GDS（GPUDirect Storage）不支持 Windows
4. PyTorch 版本锁定：某些旧版本 DeepSpeed 需要指定 PyTorch 版本，安装前查 PyPI 页面的版本对应表
5. 多节点训练：需要 NCCL（GPU 间通信）和 MPI 或 PyTorch 的 torchrun，配置相对复杂，新手建议从单节点多卡开始
6. DeepSpeed 与 transformers 版本耦合：Hugging Face transformers 版本过低可能无法识别 DeepSpeed ZeRO，建议 transformers ≥ 4.20

与同类对比

结论：DeepSpeed 适合已经在用 PyTorch 生态的用户——Hugging Face、Lightning 都能无缝接入。ColossalAI 在视频生成（Open-Sora）生态上有优势，DeepSpeed 在通用性和推理优化上积累更深。

一句话推荐结论

DeepSpeed 是大模型训练的「显存救星」，ZeRO 技术让 70B 模型在数张消费级 GPU 上成为可能——无论你是研究员还是工程师，都值得在训练脚本里加上它。