2、环境准备：安装nvidia-ml-py、pynvml、torch等工具库
好，咱们正式开始动手前的第一步——搭环境。
说实话，排查显存碎片这件事，最怕的不是问题本身，而是你连监控工具都没装好，就开始瞎猜。我见过不少同学，一上来就盯着nvidia-smi看，看了半天也看不出个所以然。为什么？因为nvidia-smi只告诉你显存用了多少，它不告诉你碎片率。
所以，咱们得装几样趁手的兵器。
2.1 核心工具清单
先列个单子，看看今天要装什么：
| 工具库 | 用途 | 我为什么推荐它 |
|---|---|---|
| nvidia-ml-py | Python封装的NVML接口 | 轻量，直接调底层API |
| pynvml | 另一个NVML封装 | 社区活跃，文档全 |
| torch | 深度学习框架 | 咱们主要用它来模拟显存分配 |
| cuda-python | CUDA驱动级接口 | 查底层信息时用得上 |
嗯，这里要注意：nvidia-ml-py和pynvml其实是一回事，底层都调的是NVML库。我个人习惯用pynvml，因为它的异常处理更友好一些。但如果你喜欢原汁原味的，nvidia-ml-py也没问题。
2.2 安装步骤
安装其实很简单，一行命令搞定。但别急着复制粘贴，先看看你的环境。
推荐做法：先创建一个干净的虚拟环境，别把系统Python搞乱了。
# 创建虚拟环境（Python 3.8+ 都行）
conda create -n gpu_mem_debug python=3.9
conda activate gpu_mem_debug
# 安装核心工具
pip install nvidia-ml-py==11.525.84
pip install pynvml==11.5.0
pip install torch==2.1.0 # 根据你的CUDA版本选
pip install cuda-python==12.1.0我曾经踩过一个坑：直接pip install nvidia-ml-py，结果装了个老版本，连显存温度都读不出来。所以，版本号最好指定一下。上面列的是我测试过的稳定版本。
小技巧：如果你用的是公司内网或者代理环境，记得加 --trusted-host 参数。不然下载到一半断了，你想想看，多闹心。
2.3 验证安装
装完了，怎么知道好不好使？写几行代码测一下。
import pynvml
import torch
# 测试pynvml
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")
# 测试torch
if torch.cuda.is_available():
print(f"CUDA版本: {torch.version.cuda}")
print(f"GPU型号: {torch.cuda.get_device_name(0)}")
else:
print("CUDA不可用，检查驱动！")跑完之后，如果能看到显存信息和GPU型号，说明环境没问题。如果报错，别慌，八成是驱动版本不匹配。
注意：pynvml需要管理员权限才能初始化。如果你在容器里跑，记得加 --privileged 参数。我刚开始用Docker时，折腾了半天才发现是权限问题。
2.4 知识体系一览
为了让你对整个环境有个全局认识，我画了张图。说白了，咱们的排查工作就是围绕这三层展开的：
你看，最上层是咱们的模型代码，中间是监控工具，最下面是驱动。排查显存碎片，其实就是用中间这层工具，去观察最上层代码在底层驱动上的行为。
2.5 避坑指南
最后，分享几个我实际遇到过的坑：
- 版本不匹配：pynvml 11.x 需要 CUDA 11.x 以上。如果你还在用 CUDA 10，老老实实装 pynvml 9.x。
- 多卡环境：如果你有4张卡，记得遍历 index 0,1,2,3。我一开始只查了 index 0，结果另外三张卡碎片成啥样都不知道。
- 虚拟环境隔离：别在 base 环境装。我曾经图省事，结果项目之间依赖打架，浪费了一下午。
好了，环境就绪。工具都装好了，咱们下一节就可以开始真正的排查了。
想要更多资料，联系可可