根治 PyTorch CUDA `pynvml` 弃用警告：直接修改 `torch/cuda/__init__.py` 的实践记录

根治 PyTorch CUDA `pynvml` 弃用警告：直接修改 `torch/cuda/init.py` 的实践记录

【笔记】消除 "FutureWarning: The pynvml package is deprecated"警告

1. 问题现象

在 Windows 环境下运行 ComfyUI（或任何基于 PyTorch 的 CUDA 应用）时，控制台会在 PyTorch 初始化阶段抛出一条 FutureWarning：

复制代码

H:\PythonProjects3\Win_ComfyUI\.venv\Lib\site-packages\torch\cuda\__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead.
  import pynvml  # type: ignore[import]

虽然这条警告不影响 CUDA 的正常使用，也不会阻断程序运行，但在每次启动时污染日志输出，尤其在使用 --verbose 或排查其他问题时，显得格外扎眼。

2. 根因分析

这条警告的源头并非 PyTorch 本身，而是 pynvml 这个包的设计策略。

pynvml 13.x 的本质 ：当前 PyPI 上的 pynvml 已经演变成一个兼容 shim（垫片） 。它内部完全依赖 nvidia-ml-py 提供 NVML（NVIDIA Management Library）的 Python 绑定，但会在被 import 时主动抛出 FutureWarning，提醒用户迁移。
PyTorch 的行为 ：PyTorch 在 torch.cuda 初始化时，硬编码了 import pynvml 来获取驱动版本、显存信息等。由于它先尝试导入 pynvml，自然就触发了 shim 层的警告。
环境的现状 ：通过 pipdeptree 可以确认，nvidia-ml-py 已经正确安装，且被 flashinfer-python、inference 等包直接依赖；而 pynvml 的 Required-by 为空，说明没有硬依赖它的上层包。

因此，真正提供功能的是 nvidia-ml-py，而 pynvml 只是一个多此一举的警告发射器。

python 复制代码

pipdeptree --reverse --packages pynvml

3. 解决思路：让 PyTorch 直接对接底层库

与其用环境变量或 warnings.filterwarnings 去"捂嘴"屏蔽警告，不如从源头根治：让 PyTorch 跳过 pynvml 这个 shim，直接导入 nvidia-ml-py。

由于 PyTorch 后续代码中大量使用了 pynvml.nvmlInit()、pynvml.nvmlDeviceGetCount() 等调用，我们只需要在导入时做一个别名映射 ，让 pynvml 这个命名空间直接指向 nvidia-ml-py 即可。这样无需改动 PyTorch 的任何业务逻辑，实现零侵入式修复。

4. 具体操作步骤

定位文件（请根据你的实际虚拟环境路径调整）：

复制代码

H:\PythonProjects3\Win_ComfyUI\.venv\Lib\site-packages\torch\cuda\__init__.py

找到对应代码块（大约在第 59~61 行）：

python 复制代码

    try:
        if not _version.hip:
            import pynvml  # type: ignore[import]

修改后的代码：

python 复制代码

    try:
        if not _version.hip:
#            import pynvml  # type: ignore[import]
            import nvidia_ml_py as pynvml

关键细节：

保持缩进 ：确保新添加的 import nvidia_ml_py as pynvml 与原 import pynvml 的缩进完全一致（4 个空格），否则 Python 解释器会抛出 IndentationError。
注释而非删除：将原行注释掉，而不是直接删除。这样一旦需要回滚，只需取消注释即可，无需记忆修改内容。
添加书签 ：在编辑器（如 VS Code、PyCharm、Notepad++ 等）中为该行添加书签。PyTorch 重装（pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu126 --force-reinstall ）会覆盖此文件，书签能帮你下次快速定位并重新修改。

python 复制代码

pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu126 --force-reinstall

5. 验证效果

保存文件后，重启 ComfyUI。观察日志：

修改前 ：FutureWarning: The pynvml package is deprecated... 出现在 torch/cuda/__init__.py:61
修改后：该位置警告完全消失，且 CUDA 初始化、GPU 检测、显存查询等功能一切正常

6. 注意事项与回溯策略

事项	说明
虚拟环境隔离	修改仅作用于当前 `.venv`，不会影响系统 Python 或其他项目的 PyTorch 安装
PyTorch 重装	执行 `torch` 重装后，`__init__.py` 会被覆盖，需重新执行上述修改。这也是建议添加书签的原因
版本兼容性	`nvidia-ml-py` 是 NVML 的官方 Python 绑定，API 与旧 `pynvml` 完全兼容，不存在功能风险
备份建议	如果你担心改坏，可先复制一份 `__init__.py` 为 `__init__.py.bak`

7. 与其他方案的对比

为什么推荐直接改源码，而不是其他更"温和"的手段？

方案	缺点
`set PYTHONWARNINGS=ignore::FutureWarning`	过于粗暴，会同时屏蔽项目中其他可能有价值的 `FutureWarning`
在 `sitecustomize.py` 中 `warnings.filterwarnings`	需要额外维护一个文件，且属于"掩耳盗铃"，警告的触发逻辑仍在运行
直接 `pip uninstall pynvml`	虽然 `pipdeptree` 显示无依赖，但某些 ComfyUI 自定义节点可能硬编码 `import pynvml`，卸载后可能导致其报错

相比之下，修改 torch/cuda/__init__.py 的别名导入 是风险最低、最优雅的根治方案：它保留了 pynvml 包在环境中的存在（安抚可能依赖它的第三方插件），同时让 PyTorch 直接走官方库，彻底消除警告。

8. 总结

这是一次典型的"面向源码排障"实践。通过分析依赖树的实际指向，我们发现 pynvml 已经退化为无实质功能的警告层，而 PyTorch 恰好是这个警告的触发者。利用 Python import ... as ... 的别名机制，我们在 PyTorch 内部做了一个无缝的"拨乱反正"，既保留了环境兼容性，又实现了日志的彻底清净。

下次重装 / 升级 PyTorch 后，只需凭借编辑器书签回到同一位置，重复一次注释与替换，即可继续享受零警告的启动体验。

根治 PyTorch CUDA `pynvml` 弃用警告：直接修改 `torch/cuda/__init__.py` 的实践记录