3、环境准备与安装：CUDA Toolkit版本要求、Python依赖安装、Docker镜像构建

说实话，搞显存碎片治理这事儿，环境搭不对，后面全是坑。我见过太多人花了一周调参数，最后发现是CUDA版本不兼容——那种感觉，啧，比加班到凌晨还难受。

这一章咱们就把环境彻底搞定。你跟着我走一遍，后面实操的时候就不会被环境问题绊住脚。

3.1 CUDA Toolkit版本要求

先说说CUDA版本。显存碎片治理工具，说白了就是跟GPU驱动和CUDA runtime打交道。版本不对，工具根本跑不起来。

我个人的建议是：

• 最低要求：CUDA 11.4 —— 这是底线。低于这个版本，很多内存管理API（比如 cudaMallocAsync）压根不存在。
• 推荐版本：CUDA 11.8 或 12.1 —— 这两个版本我实测过，稳定性最好。特别是11.8，我在好几个生产项目里都用它，没出过幺蛾子。
• 注意： CUDA 12.x 对显存池管理做了大改，如果你用12.0之前的版本，有些新特性（比如 cudaMemPoolSetAttribute）会报错。

为什么会这样？因为CUDA Toolkit是向上兼容的，但驱动是向下兼容的。你装个新驱动，老CUDA还能用；反过来就不行。嗯，这里要注意。

3.2 Python依赖安装

Python环境这块，我建议你用 Python 3.9 ~ 3.11。3.8太老，3.12有些库还没适配好。

先列一下核心依赖：

包名	版本要求	用途
torch	≥2.0.0	深度学习框架，显存分析的基础
nvidia-ml-py	≥11.515	获取GPU显存使用详情
pynvml	≥11.5.0	另一个NVML封装，跟上面二选一
cupy	≥12.0.0	GPU上的NumPy，做显存分配测试用
matplotlib	≥3.7	画显存碎片分布图
psutil	≥5.9	监控系统内存，辅助分析

安装命令很简单：

pip install torch nvidia-ml-py cupy matplotlib psutil

避坑指南： 我曾经在某个项目里直接 pip install torch，结果装了个CPU版本。你想想看，GPU工具装了个CPU版PyTorch，那还玩啥？所以一定要指定CUDA版本：

# 装CUDA 11.8对应的PyTorch
pip install torch --index-url https://download.pytorch.org/whl/cu118

# 或者CUDA 12.1
pip install torch --index-url https://download.pytorch.org/whl/cu121

3.3 Docker镜像构建

说实话，Docker是解决环境问题的最优解。你想想看，团队里有人用Ubuntu 20.04，有人用CentOS，还有人用Mac——光CUDA驱动版本就能吵一架。用Docker，一次构建，到处运行。

下面是我自己项目里用的Dockerfile，你直接拿去用：

# 基础镜像：选CUDA 11.8 + Ubuntu 22.04
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

# 安装Python 3.10
RUN apt-get update && apt-get install -y \
    python3.10 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

# 设置工作目录
WORKDIR /workspace

# 复制依赖文件
COPY requirements.txt .

# 安装Python依赖
RUN pip3 install --no-cache-dir -r requirements.txt

# 复制项目代码
COPY . .

# 默认命令
CMD ["python3"]

对应的 requirements.txt 内容：

torch==2.1.0+cu118
nvidia-ml-py==11.515.0
cupy-cuda11x==12.0.0
matplotlib==3.7.2
psutil==5.9.5

构建和运行命令：

# 构建镜像
docker build -t mem-frag-tool:latest .

# 运行容器（记得加 --gpus all）
docker run --gpus all -it --rm \
    -v $(pwd):/workspace \
    mem-frag-tool:latest

3.4 验证环境是否就绪

装完别急着走，跑个验证脚本确认一下：

import torch
import pynvml

# 检查PyTorch能否看到GPU
print(f"PyTorch CUDA可用: {torch.cuda.is_available()}")
print(f"GPU数量: {torch.cuda.device_count()}")
print(f"当前GPU: {torch.cuda.get_device_name(0)}")

# 检查NVML是否正常
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(f"总显存: {mem_info.total / 1024**3:.2f} GB")
print(f"已用显存: {mem_info.used / 1024**3:.2f} GB")
print(f"空闲显存: {mem_info.free / 1024**3:.2f} GB")

如果输出里能看到GPU型号和显存信息，那环境就齐活了。如果报错，多半是CUDA版本或驱动问题，回头检查一下3.1节的内容。

下面这张图帮你理清整个环境搭建的逻辑：

说白了，环境搭建就是一层叠一层：硬件 → 驱动 → CUDA → Python/Docker。哪一层出问题，工具都跑不起来。我建议你按这个顺序排查，别跳步。

好了，环境准备好了。接下来就可以真正动手治理显存碎片了。

---