环境准备：NVIDIA驱动、CUDA Toolkit、Docker 环境搭建

做GPU集群压测，环境搭建是第一步，也是最容易踩坑的一步。我见过不少团队，压测脚本写得漂亮，结果环境没配好，一跑就崩。说白了，底层驱动和容器化环境就是地基，地基不稳，上层再花哨也没用。

这一章，我带你一步步把NVIDIA驱动、CUDA Toolkit和Docker环境搭起来。嗯，都是我在项目中反复验证过的流程，你照着做就行。

3.1 NVIDIA驱动安装：别小看这一步

NVIDIA驱动是GPU和操作系统之间的桥梁。没有它，你的显卡就是个摆设。我个人习惯用runfile方式安装，虽然比包管理器麻烦点，但可控性高，不容易出现版本冲突。

3.1.1 检查当前硬件和驱动

先看看你的机器上有没有NVIDIA显卡，以及当前驱动状态：

# 查看显卡型号
lspci | grep -i nvidia

# 查看当前驱动版本（如果有的话）
nvidia-smi

如果nvidia-smi报错，说明没装驱动或者驱动挂了。我曾经在一台新到的DGX机器上，发现系统自带的驱动版本太老，导致CUDA 11.8死活装不上。折腾了两小时才发现是驱动的问题。

3.1.2 卸载旧驱动（重要）

为什么要卸载旧的？因为新旧驱动混在一起，会导致nvidia-smi输出乱码，甚至内核崩溃。我建议用以下命令彻底清理：

# 卸载NVIDIA相关包
sudo apt-get purge nvidia-* -y
sudo apt-get autoremove -y

# 如果是runfile安装的，用这个
sudo nvidia-uninstall

3.1.3 安装驱动

去NVIDIA官网下载对应型号的驱动。我个人推荐525.105.17这个版本，兼容性好，我在多个集群上验证过。

# 禁用nouveau驱动
echo "blacklist nouveau" | sudo tee /etc/modprobe.d/blacklist-nvidia-nouveau.conf
sudo update-initramfs -u

# 重启后，安装驱动
sudo chmod +x NVIDIA-Linux-x86_64-525.105.17.run
sudo ./NVIDIA-Linux-x86_64-525.105.17.run --no-opengl-files

安装完成后，运行nvidia-smi，你应该能看到类似这样的输出：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  NVIDIA A100-SXM4...  On  | 00000000:00:04.0 Off |                    0 |
|   0%   35C    P0    45W / 400W|      0MiB / 40960MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

3.2 CUDA Toolkit安装：选对版本很重要

CUDA Toolkit是GPU编程的SDK，包含编译器、库和工具。做压测时，我们主要用它来编译和运行GPU算子。

3.2.1 版本选择策略

你想想看，CUDA版本选错了，后面所有依赖都得重装。我建议遵循这个原则：

• 看框架要求： PyTorch、TensorFlow都有明确的CUDA版本要求。比如PyTorch 2.0推荐CUDA 11.8。
• 看驱动版本： 驱动支持的CUDA版本必须高于你要装的Toolkit版本。用nvidia-smi顶部显示的CUDA Version作为上限。
• 看集群统一： 整个集群最好用同一个CUDA版本，否则管理起来很头疼。

3.2.2 安装CUDA 11.8

以CUDA 11.8为例，我习惯用runfile安装，因为可以指定安装路径，方便多版本共存。

# 下载CUDA 11.8 runfile
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run

# 运行安装
sudo sh cuda_11.8.0_520.61.05_linux.run --toolkit --silent --override

安装完成后，配置环境变量：

# 添加到 ~/.bashrc
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH

# 生效
source ~/.bashrc

验证安装：

nvcc --version

输出应该显示Cuda compilation tools, release 11.8, V11.8.89。

3.3 Docker环境搭建：容器化是压测的标配

为什么压测要用Docker？说白了，就是为了环境一致性。你在本地调好的压测脚本，放到集群上跑，结果因为依赖版本不同报错——这种事我遇到过太多次了。

3.3.1 安装Docker Engine

我推荐用官方脚本安装，简单粗暴：

# 安装Docker
curl -fsSL https://get.docker.com -o get-docker.sh
sudo sh get-docker.sh

# 将当前用户加入docker组（避免每次sudo）
sudo usermod -aG docker $USER

# 重启docker服务
sudo systemctl restart docker

3.3.2 安装NVIDIA Container Toolkit

这是关键一步。没有它，Docker容器里是认不出GPU的。你想想看，容器里跑nvidia-smi报错，那压测还怎么搞？

# 添加NVIDIA容器工具包仓库
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

# 重启Docker
sudo systemctl restart docker

3.3.3 验证GPU容器

跑一个测试容器，看看GPU能不能用：

docker run --rm --gpus all nvidia/cuda:11.8.0-base nvidia-smi

如果能看到GPU信息，说明环境搭好了。如果报错could not select device driver "" with capabilities: [[gpu]]，八成是nvidia-container-toolkit没装好，或者Docker没重启。

3.4 知识体系总览

下面这张图，帮你理清这三者的关系：

这张图很直观：硬件在最底下，驱动在中间，CUDA Toolkit在上层，Docker容器在最上面。每一层都依赖下一层。你装环境的时候，必须从下往上装，顺序不能乱。

3.5 环境验证脚本

环境搭好后，我习惯跑一个完整的验证脚本，确保所有组件都正常工作：

#!/bin/bash
# 环境验证脚本

echo "=== 1. 检查NVIDIA驱动 ==="
nvidia-smi || { echo "驱动未安装"; exit 1; }

echo "=== 2. 检查CUDA Toolkit ==="
nvcc --version || { echo "CUDA Toolkit未安装"; exit 1; }

echo "=== 3. 检查Docker ==="
docker --version || { echo "Docker未安装"; exit 1; }

echo "=== 4. 检查GPU容器 ==="
docker run --rm --gpus all nvidia/cuda:11.8.0-base nvidia-smi || { echo "GPU容器无法运行"; exit 1; }

echo "=== 全部通过！环境准备完成 ==="

把这个脚本保存为check_env.sh，每次新机器到手，先跑一遍。省心。

想要更多资料，联系可可