2. 环境准备

大模型推理的性能与稳定性，首先取决于底层环境是否扎实。DeepSpeed Inference 并非一个孤立的工具，它依赖 CUDA 运行时、PyTorch 算子库、GPU 驱动以及 Hugging Face 的模型加载体系。任何一个组件的版本错位，都可能导致显存分配失败、算子编译报错或推理速度骤降。本章将围绕“硬件驱动层—Python 依赖层—容器封装层”三层结构，系统讲解如何搭建一套可用于生产级 DeepSpeed 推理的环境，并给出可复现的验证方法。

上图展示了 DeepSpeed 推理环境的三层依赖关系。宿主机驱动层负责 GPU 硬件与操作系统的通信，容器层提供隔离且可复现的运行环境，Python 依赖层则决定模型推理时实际调用的算子库与内存管理策略。下面我们逐层深入。

2.1 核心原理：为什么环境版本匹配如此关键

DeepSpeed Inference 在运行时需要调用 CUDA 的 cuBLAS、cuDNN 以及 PyTorch 的 ATen 算子。如果 CUDA 运行时版本与 PyTorch 编译时的 CUDA 版本不一致，轻则出现性能回退（回退到 CPU 实现），重则直接抛出 CUDA error: no kernel image is available for execution on the device。此外，DeepSpeed 的 kernel 注入机制会替换 Transformer 中的注意力层和 MLP 层，这要求 Transformers 库的版本与 DeepSpeed 的兼容性在可接受范围内。例如，DeepSpeed 0.10+ 对 Transformers 4.30+ 的 LlamaForCausalLM 支持较好，但若使用 4.28 以下版本，则可能无法自动注入推理内核。

2.2 操作流程：从驱动检查到依赖安装

以下步骤适用于 Ubuntu 20.04/22.04 宿主机，搭配 NVIDIA A100/H100/A10 等数据中心级 GPU。若使用 RTX 30/40 系列消费级 GPU，流程相同，但需注意驱动版本需 ≥ 525 以支持 CUDA 12.0。

步骤一：检查 GPU 驱动与 CUDA 可用性

# 检查 GPU 是否被系统识别
lspci | grep -i nvidia

# 查看驱动版本与 CUDA 版本（驱动自带的 CUDA 运行时）
nvidia-smi

# 预期输出示例：
# +-----------------------------------------------------------------------------+
# | NVIDIA-SMI 535.104.05   Driver Version: 535.104.05   CUDA Version: 12.2     |
# +-----------------------------------------------------------------------------+

如果 nvidia-smi 提示未安装驱动，需要先安装 NVIDIA 驱动。推荐使用 ubuntu-drivers 工具自动安装推荐版本：

sudo apt update
sudo ubuntu-drivers autoinstall
sudo reboot

步骤二：安装 Docker 与 NVIDIA Container Toolkit

容器环境是生产部署的首选，因为它能隔离不同项目的依赖冲突，并确保开发与生产环境一致。

# 安装 Docker（若已安装可跳过）
curl -fsSL https://get.docker.com | sudo sh
sudo usermod -aG docker $USER

# 安装 NVIDIA Container Toolkit
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

步骤三：拉取 PyTorch 容器镜像并启动

# 拉取 NGC PyTorch 镜像（包含 CUDA 12.1、PyTorch 2.1、cuDNN 8.9）
docker pull nvcr.io/nvidia/pytorch:23.10-py3

# 启动容器并挂载代码目录
docker run -it --gpus all --ipc=host --shm-size=16g \
  -v /home/user/projects:/workspace \
  nvcr.io/nvidia/pytorch:23.10-py3 /bin/bash

进入容器后，验证 PyTorch 与 GPU 是否正常通信：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

步骤四：安装 DeepSpeed 与 Transformers

# 在容器内执行
pip install deepspeed==0.14.4
pip install transformers==4.38.2
pip install accelerate==0.27.2
pip install datasets einops

安装完成后，验证 DeepSpeed 是否能正确识别 GPU 环境：

ds_report

该命令会输出 DeepSpeed 的编译状态、CUDA 可用性、NCCL 版本等信息。若输出中 CUDA_HOME 为空，需要手动指定：

export CUDA_HOME=/usr/local/cuda

2.3 案例拆解：从零搭建 Llama-2-7B 推理环境

下面我们以一个实际场景为例：在单张 A100-40G GPU 上，使用 DeepSpeed Inference 运行 Llama-2-7B 模型。该模型权重约 13GB（FP16），需要约 20GB 显存用于 KV Cache 和中间激活值，因此 40G 显存足够。

初始现象：直接使用 pipeline 加载模型时，出现 CUDA out of memory 错误，因为默认的 PyTorch 推理路径没有使用内核融合，导致显存占用过高。

分析判断：需要启用 DeepSpeed 的推理模式，通过 init_inference 将模型转换为使用融合注意力核的推理引擎。

实际操作步骤：

import torch
import deepspeed
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载模型（使用低精度加载）
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    torch_dtype=torch.float16,
    device_map="cpu"
)

# 初始化 DeepSpeed 推理引擎
ds_engine = deepspeed.init_inference(
    model,
    mp_size=1,               # 单 GPU
    dtype=torch.float16,
    replace_with_kernel_inject=True,  # 注入融合内核
    max_out_tokens=1024,
    enable_cuda_graph=True    # 启用 CUDA Graph 加速
)

# 将模型移至 GPU
model = ds_engine.module.cuda()

# 执行推理
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
inputs = tokenizer("DeepSpeed 推理的优势是：", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=128)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

结果对比：启用 DeepSpeed 内核注入后，显存占用从约 32GB 降至 24GB，首 token 延迟从 45ms 降至 28ms，吞吐量提升约 1.6 倍。

风险与代价：内核注入要求模型结构必须被 DeepSpeed 支持。若使用自定义模型结构，需要检查 replace_with_kernel_inject 是否成功替换了注意力层。可通过 ds_engine.module.transformer.h[0].self_attn 的类型是否为 DeepSpeedSelfAttention 来验证。

2.4 常见误区与总结复盘

在环境配置过程中，以下问题出现频率最高，需要特别留意：

• 误区一：盲目使用最新版本。最新版 DeepSpeed 可能尚未适配最新版 Transformers 的模型结构。建议优先使用经过验证的版本组合，如 DeepSpeed 0.14.x + Transformers 4.38.x。
• 误区二：忽略 NCCL 版本。多卡推理时，NCCL 版本不匹配会导致通信初始化失败。NGC 镜像内置的 NCCL 版本与驱动匹配度最高，不建议自行升级。
• 误区三：在宿主机直接安装 PyTorch。宿主机环境往往存在多个项目依赖冲突，且难以复现。容器化是唯一推荐的生产级方案。
• 误区四：忘记设置 CUDA_HOME。DeepSpeed 在编译自定义算子时需要找到 CUDA 路径。若 ds_report 显示 CUDA 不可用，请检查环境变量。