3. 环境准备：GPU驱动安装、CUDA Toolkit配置、PyTorch/TensorFlow环境搭建
好，咱们进入实战前的第一道坎——环境准备。
说实话，我见过太多同学在显存压测时跑出奇怪的结果，最后发现是驱动版本不对，或者CUDA和框架没对齐。这种问题排查起来特别费劲。所以这一章，咱们把环境彻底捋清楚。
3.1 GPU驱动安装——别小看这一步
驱动是GPU和操作系统之间的翻译官。翻译官不靠谱，后面全白搭。
核心原则：驱动版本 ≥ CUDA Toolkit要求的驱动版本。
我个人习惯用 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装好了，结果编译时各种报错——嗯，这个坑我替你们踩过了。
避坑指南：如果你用的是数据中心卡（比如A100、H100），建议用 nvidia-smi -pm 1 开启持久化模式。我在压测时遇到过驱动因为空闲自动降频，导致显存带宽忽高忽低，开了这个就稳了。
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
注意：千万别把 nvidia-smi 显示的CUDA版本和 nvcc --version 搞混了。前者是驱动支持的版本上限，后者是你实际安装的Toolkit版本。我见过有人对着 nvidia-smi 的12.2版本，硬要装CUDA 12.2的Toolkit，结果驱动版本不够新——装不上。
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'))"
我的经验：压测时建议用虚拟环境隔离。我习惯每个项目建一个独立的conda环境，避免依赖冲突。有一次我在一个环境里同时装了PyTorch和TF的不同版本，结果两个框架的CUDA库打架，显存分配直接崩了——从那以后我再也不偷懒了。
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 → 框架 → 验证。哪一环断了，压测都跑不起来。
总结一下：
- 驱动版本要够新，用
nvidia-smi确认 - CUDA Toolkit版本要和框架对齐，别混用
- 用conda虚拟环境隔离项目，避免依赖冲突
- 装完后一定跑显存分配测试，确认环境正常
环境准备好了，咱们就可以开始真正的压测了。下一章我会带你写第一个显存分配脚本，看看你的GPU到底能吞下多大的数据块。