3. 环境准备

好，咱们进入实战前的第一道坎——环境准备。

说实话，我见过太多同学在显存压测时跑出奇怪的结果，最后发现是驱动版本不对，或者CUDA和框架没对齐。这种问题排查起来特别费劲。所以这一章，咱们把环境彻底捋清楚。

3.1 GPU驱动安装——别小看这一步

驱动是GPU和操作系统之间的翻译官。翻译官不靠谱，后面全白搭。

我个人习惯用 nvidia-smi 来确认驱动状态。装完驱动后，第一件事就是跑这个命令：

nvidia-smi

你会看到类似这样的输出：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.129.03   Driver Version: 535.129.03   CUDA Version: 12.2     |
|-------------------------------+----------------------+----------------------+
| 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 80GB    Off  | 00000000:00:04.0 Off |                    0 |
|   0%   35C    P0    45W / 300W|      0MiB / 81920MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

注意看右上角的 CUDA Version。这个数字不是说你装了CUDA Toolkit，而是当前驱动支持的最高CUDA版本。我曾经遇到过一位同事，盯着这个数字以为CUDA装好了，结果编译时各种报错——嗯，这个坑我替你们踩过了。

3.2 CUDA Toolkit配置——版本对齐是门玄学

CUDA Toolkit说白了就是GPU的软件开发包。你写代码调用的 cudaMalloc、cudaMemcpy 这些API，都来自它。

这里有个关键点：PyTorch/TensorFlow 对 CUDA 版本有硬性要求。比如PyTorch 2.0 官方支持 CUDA 11.7 和 11.8，你用 CUDA 12.0 编译可能也能跑，但我不建议在生产环境这么干。

我一般这样装：

# 下载CUDA Toolkit 11.8（以Ubuntu 22.04为例）
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

# 配置环境变量
echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

验证是否装好：

nvcc --version

你会看到：

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:41:10_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

3.3 PyTorch/TensorFlow环境搭建——选对版本，少走弯路

框架版本和CUDA版本的对应关系，我建议直接查官方文档。这里我给出我压测时常用的组合：

框架	版本	推荐CUDA	备注
PyTorch	2.0.x	11.7 / 11.8	显存压测首选，API稳定
PyTorch	2.1.x	11.8 / 12.1	支持新特性，但部分算子有变动
TensorFlow	2.12.x	11.8	TF的显存管理比较特殊，后面会细讲
TensorFlow	2.13.x	11.8 / 12.0	注意TF 2.13开始默认使用XLA

安装PyTorch，我推荐用conda，省心：

conda create -n gpu_test python=3.10
conda activate gpu_test
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118

验证安装：

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

如果输出 True 和你的GPU型号，说明环境通了。

TensorFlow的安装类似：

pip install tensorflow==2.12.0
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

3.4 环境验证——跑一个简单的显存分配测试

环境搭好了，咱们跑个简单的脚本，确认显存分配和释放没问题：

import torch

# 分配1GB显存
x = torch.empty(256, 1024, 1024, dtype=torch.float32, device='cuda')
print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")

# 释放
del x
torch.cuda.empty_cache()
print(f"After release: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")

输出应该是：

Allocated: 1.00 GB
After release: 0.00 GB

如果显存没释放干净，说明你的环境可能有问题。我在压测时遇到过框架缓存导致显存泄漏的情况，后面会专门讲怎么排查。

3.5 知识体系总览

下面这张图帮你理清环境准备的核心逻辑：

说白了，环境准备就是一条链：硬件 → 驱动 → CUDA Toolkit → 框架 → 验证。哪一环断了，压测都跑不起来。