3、DeepSpeed环境搭建：安装配置、依赖管理、快速启动示例
DeepSpeed 的安装并不复杂，但它对运行环境有比较明确的依赖边界：CUDA 版本、PyTorch 版本、Python 版本、编译器与 MPI 之间需要彼此匹配。很多初学者第一次失败并不是因为 DeepSpeed 本身难用，而是因为环境链条中某一环版本错位。本节从依赖关系出发，给出可复现的安装流程、依赖管理策略，并用一个最小可运行示例验证环境是否真正可用。
一、核心原理：DeepSpeed 的依赖分层
理解 DeepSpeed 的安装，关键是看清它的四层依赖结构。最底层是 GPU 驱动与 CUDA 运行时，它决定了硬件能否被调用；其上是 PyTorch，DeepSpeed 的绝大多数算子与通信都构建在 PyTorch 之上；再往上是 DeepSpeed 自身及其可选扩展，例如融合算子、CPU 优化器、通信压缩等；最上层才是你的训练脚本与配置文件。任何一层不匹配，都会在导入或初始化阶段暴露为报错。
图中自下而上的箭头表示“下层为上层提供能力”。安装时应当从下往上逐层确认：先用 nvidia-smi 确认驱动，再确认 PyTorch 能否调用 GPU，最后才安装 DeepSpeed。顺序颠倒会导致排查困难。
二、方法流程：可复现的安装步骤
推荐使用 conda 创建独立环境，避免污染系统 Python。以下流程以 CUDA 11.8 与 PyTorch 2.1 为例，实际使用时请按你的驱动版本替换。
# 1. 创建并激活独立环境
conda create -n ds python=3.10 -y
conda activate ds
# 2. 安装与 CUDA 匹配的 PyTorch
pip install torch==2.1.0 torchvision torchaudio \
--index-url https://download.pytorch.org/whl/cu118
# 3. 验证 PyTorch 能否看到 GPU
python -c "import torch; print(torch.cuda.is_available(), torch.version.cuda)"
# 4. 安装 DeepSpeed
pip install deepspeed
# 5. 查看 DeepSpeed 环境报告
ds_report
其中 ds_report 是最重要的自检命令，它会逐项列出 torch、CUDA、编译器、MPI、融合算子等是否可用。若某项显示为 NO，并不一定导致训练失败，但会影响对应优化是否生效。例如 cpu_adam 为 NO 时，CPU offload 优化器将无法使用。
pip install deepspeed 在编译阶段耗时过长或失败，可以先安装预编译版本，或使用 DS_BUILD_OPS=0 pip install deepspeed 跳过可选算子编译，先保证核心功能可用，后续再按需补装。
三、依赖管理：版本匹配清单
下表列出了常见组合的匹配关系，可作为选型参考。注意驱动版本必须大于等于 CUDA 运行时要求，而 PyTorch 的 cu 后缀必须与 CUDA 运行时一致。
| 组件 | 推荐版本 | 检查命令 | 不匹配的典型症状 |
|---|---|---|---|
| NVIDIA 驱动 | ≥ 525（CUDA 11.8） | nvidia-smi |
CUDA driver version is insufficient |
| CUDA 运行时 | 11.8 / 12.1 | nvcc --version |
PyTorch 无法加载 CUDA 扩展 |
| PyTorch | 2.0 ~ 2.2 | torch.version.cuda |
导入 deepspeed 时报符号错误 |
| Python | 3.8 ~ 3.11 | python --version |
依赖包无对应 wheel |
| DeepSpeed | 0.12 以上 | ds_report |
ZeRO 配置项不被识别 |
四、快速启动示例：最小可运行验证
环境装好后，用一个最小脚本验证分布式初始化与 ZeRO 是否生效。下面脚本使用 deepspeed.initialize 包装模型和优化器，并打印当前 rank 与 ZeRO 阶段。
# train_min.py
import torch
import deepspeed
def main():
model = torch.nn.Linear(128, 10).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
ds_config = {
"train_batch_size": 8,
"fp16": {"enabled": True},
"zero_optimization": {"stage": 2}
}
engine, opt, _, _ = deepspeed.initialize(
model=model, optimizer=optimizer, config=ds_config
)
x = torch.randn(8, 128).cuda()
loss = engine(x).sum()
engine.backward(loss)
engine.step()
print("rank", engine.local_rank, "zero stage", engine.zero_optimization_stage())
if __name__ == "__main__":
main()
使用 DeepSpeed 启动器运行，指定两张卡：
deepspeed --num_gpus=2 train_min.py
若输出中每个 rank 都打印了正确的 ZeRO stage，且没有通信超时或显存溢出，说明环境搭建成功。此时可以逐步把 stage 提升到 3，并加入 offload 参数观察显存变化。
五、误区与复盘
- 误区一：先装 DeepSpeed 再装 PyTorch。DeepSpeed 安装时会探测 PyTorch 与 CUDA，顺序颠倒会导致编译出的扩展与最终 PyTorch 不匹配。
- 误区二：忽略 MPI。多机训练依赖 MPI 或 NCCL 后端，单机多卡通常用 NCCL 即可，但跨节点必须确认 MPI 可用。
- 误区三：把
ds_report的 NO 当成致命错误。多数 NO 只影响可选优化，核心训练仍可运行，应按需补装而非盲目重装。 - 误区四：在系统 Python 中直接安装。不同项目依赖冲突时极难回退，务必使用 conda 或 venv 隔离。
ImportError: libcudart.so.11.0: cannot open shared object file，通常是 CUDA 运行时路径未加入 LD_LIBRARY_PATH，或 PyTorch 的 cu 版本与系统 CUDA 不一致。先用 ds_report 定位，再调整环境变量或重装对应版本。
复盘时建议记录三件事：驱动与 CUDA 版本、PyTorch 与 DeepSpeed 版本、ds_report 的完整输出。这三项信息足以在绝大多数环境问题中快速定位根因，也能让团队其他成员复现你的环境。