4、DeepSpeed环境搭建：安装配置、依赖管理、硬件要求与快速入门示例

DeepSpeed 是微软开源的一套深度学习优化库，核心价值在于用极小的代码改动，把 ZeRO、混合精度、梯度累积、Offload 等分布式训练能力接入 PyTorch 模型。要让它稳定工作，环境搭建不是“装一个包”那么简单，而是需要同时对齐 Python 版本、PyTorch 版本、CUDA 驱动、NCCL 通信库、编译工具链 这几层依赖。本节从硬件要求、安装流程、依赖管理到最小可运行示例，给出一条可复现的搭建路径。

4.1 硬件与软件前置要求

DeepSpeed 本身是纯 Python + 少量 CUDA 扩展，真正决定能否跑起来的是底层硬件与驱动。单卡可以用于功能验证，但 ZeRO 的 Stage 2/3 和 Offload 只有在多卡或显存受限时才有明显收益。下表给出常见场景的最低与推荐配置。

4.2 安装流程与依赖管理

推荐使用 conda 或 venv 创建独立环境，避免与系统 Python 冲突。安装顺序应为：先装 PyTorch，再装 DeepSpeed，因为 DeepSpeed 在安装时会探测当前 PyTorch 的 CUDA 版本并编译对应扩展。

# 1. 创建并激活环境
conda create -n ds_train python=3.10 -y
conda activate ds_train

# 2. 安装与 CUDA 匹配的 PyTorch（以 cu121 为例）
pip install torch==2.3.0 torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/cu121

# 3. 安装 DeepSpeed
pip install deepspeed

# 4. 验证安装
ds_report

ds_report 是 DeepSpeed 自带的诊断命令，会输出 PyTorch、CUDA、NCCL、编译扩展等状态。若看到 DeepSpeed C++/CUDA extension op is installed 且 NCCL 可用，说明基础环境就绪。

4.3 依赖关系与版本匹配

DeepSpeed 与 PyTorch 的版本并非任意组合都可用。下表列出常见稳定组合，可作为选型参考。

依赖管理建议使用 requirements.txt 或 environment.yml 锁定版本，并在 CI 中加入 ds_report 检查。若使用 Docker，可直接基于 nvidia/cuda:12.1.0-devel-ubuntu22.04 构建，避免宿主机驱动与容器运行时不一致。

4.4 快速入门示例：单卡与多卡启动

下面是一个最小可运行的 DeepSpeed 训练脚本，展示如何用 deepspeed.initialize 包装模型和优化器。该示例可在单卡直接运行，也可通过 deepspeed 命令启动多卡。

import torch
import deepspeed

def main():
    model = torch.nn.Linear(1024, 1024).cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    ds_config = {
        "train_batch_size": 8,
        "fp16": {"enabled": True},
        "zero_optimization": {"stage": 2}
    }

    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        optimizer=optimizer,
        config=ds_config
    )

    for step in range(10):
        x = torch.randn(8, 1024).cuda()
        y = model_engine(x)
        loss = y.sum()
        model_engine.backward(loss)
        model_engine.step()
        print(f"step {step}, loss {loss.item():.4f}")

if __name__ == "__main__":
    main()

单卡运行：python train.py。多卡运行需先准备 hostfile，再执行：

deepspeed --num_gpus=4 train.py

若使用多机，可写 hostfile：

worker-0 slots=4
worker-1 slots=4

然后执行 deepspeed --hostfile=hostfile train.py。启动后可通过 nvidia-smi 观察显存占用，通过 ds_report 确认通信后端。

4.5 常见误区与复盘

• 误区一：先装 DeepSpeed 再装 PyTorch。 这会导致 DeepSpeed 编译时找不到正确的 CUDA 头文件，出现 nvcc not found 或算子编译失败。正确顺序是先 PyTorch 后 DeepSpeed。
• 误区二：忽略驱动与运行时版本差。 容器内 CUDA 12.1 但宿主机驱动只支持 11.8，会直接报驱动不足。用 nvidia-smi 查看驱动支持的最高 CUDA 版本。
• 误区三：多卡启动不写 hostfile。 单机多卡可用 --num_gpus，但多机必须提供 hostfile 并配置 SSH 免密，否则会卡在 Waiting for workers。
• 误区四：显存不足时只调 batch size。 更有效的是开启 ZeRO Stage 2/3 或 CPU Offload，把优化器状态和梯度分片到多卡或主机内存。