Nvidia-smi与PyTorch GPU内存占用报告不一致问题咨询

当运行以下PyTorch代码时，内存快照显示仅占用2MiB，符合PyTorch最小内存块为2MiB的文档说明：

tensor = torch.randn(256, 256, device='cuda', dtype=torch.float16) 
snapshot = torch.cuda.memory._snapshot()
total_size = snapshot['segments'][0]["total_size"] / 1024 / 1024
print(total_size,"MiB")

但通过nvidia-smi查询时，显示已使用391MiB内存：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.125.06   Driver Version: 525.125.06   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. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:41:00.0 Off |                  Off |
|  0%   34C    P8    14W / 450W |    391MiB / 24564MiB |      0%      Default |
|                               |                      |                  N/A |

这种差异的核心原因是两者统计的内存范围完全不同，具体包括以下几点：

• CUDA上下文初始化开销：PyTorch首次调用CUDA相关操作时，会触发CUDA Runtime和GPU驱动的上下文初始化，这部分内存用于存储设备管理数据、驱动内部结构等，属于GPU全局占用的内存，会被nvidia-smi统计，但不会被PyTorch的内存快照计入。
• PyTorch内存池预分配：为了提升内存分配效率，PyTorch的CUDA内存分配器会提前向GPU申请一块较大的内存池，后续tensor的内存分配直接从池中取用。nvidia-smi统计的是GPU被实际占用的总内存（包括整个内存池），而PyTorch的内存快照只统计池子里被tensor实际使用的部分。
• 系统与其他进程占用：如果GPU同时承担显示输出任务，桌面环境（如Xorg、Wayland）会占用部分显存；此外NVIDIA相关服务、其他后台GPU进程也会占用显存，这些都被nvidia-smi纳入统计，但PyTorch只会统计自身进程内的内存使用。

内容的提问来源于stack exchange，提问作者Tommaso Bendinelli