解决 NVIDIA RTX 6000 black well 最新架构下的 PyTorch 版本兼容问题
显卡型号、CUDA、Pytorch等基本查询
先查询自己的服务器和目前安装torch的版本号,以及服务器支持的架构是什么 。
c
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device: {torch.cuda.get_device_name(0)}")
print(f"Compute capability: {torch.cuda.get_device_capability(0)}")
print(f"Arch list: {torch.cuda.get_arch_list()}")
print(torch.cuda.get_device_capability())
我的输出是:
c
PyTorch version: 2.11.0.dev20260115+cu128
CUDA available: True
Device: NVIDIA RTX PRO 6000 Blackwell Workstation Edition
Compute capability: (12, 0)
Arch list: ['sm_70', 'sm_75', 'sm_80', 'sm_86', 'sm_90', 'sm_100', 'sm_120']
(12, 0)
定位问题
运行程序后出现以下问题:
c
RuntimeError: CUDA error: no kernel image is available for execution on the device
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
网上查询原因是:
1、GPU 实际是 NVIDIA RTX PRO 6000 Blackwell Workstation Edition,属于最新的 Blackwell 架构,算力等级为 (12, 0)(sm_120)
2、当前安装的 PyTorch 2.3.1 仅支持 sm_50/sm_60/sm_70/sm_75/sm_80/sm_86/sm_90,缺少对 sm_120(Blackwell 架构)的内核支持
也就是
GPU 是最新 Blackwell 架构(sm_120),而 PyTorch 2.3.1 及以下版本均未提供对 sm_120 算力的官方内核支持,这是当前报错的核心症结(Blackwell 是 NVIDIA 最新架构,PyTorch 官方尚未完成全版本适配)
具体解决:
1、先卸载当前安装的torch
c
pip uninstall torch torchvision torchaudio -y
conda remove torch torchvision torchaudio pytorch-cuda -y --force
2、安装正确版本的torch
1)cu128 是 PyTorch 官方为 ** 高版本 CUDA 架构(尤其是 Blackwell sm_120)** 预留的 Nightly 版通道目录。
2)--pre 的作用是允许 pip 安装预发布版本(包括 Nightly 夜间版、测试版、候选版等),而 cu128 通道下的 PyTorch 包均为未正式发布的开发版,若缺少该参数,pip 会默认只搜索稳定正式版,直接忽略 cu128 目录下的包,导致安装失败。
3)该参数与 --index-url 配合,能让 pip 精准定位到 cu128 通道的预发布包,并自动解析依赖、完成安装,无需手动下载 .whl 包
注意:我的当前环境是 Python 3.10
c
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
接下来等一段时间安装,出现如下:
c
Successfully installed cuda-bindings-12.9.4 cuda-pathfinder-1.2.2 nvidia-cublas-cu12-12.8.4.1 nvidia-cuda-cupti-cu12-12.8.90 nvidia-cuda-nvrtc-cu12-12.8.93 nvidia-cuda-runtime-cu12-12.8.90 nvidia-cudnn-cu12-9.10.2.21 nvidia-cufft-cu12-11.3.3.83 nvidia-cufile-cu12-1.13.1.3 nvidia-curand-cu12-10.3.9.90 nvidia-cusolver-cu12-11.7.3.90 nvidia-cusparse-cu12-12.5.8.93 nvidia-cusparselt-cu12-0.7.1 nvidia-nccl-cu12-2.28.9 nvidia-nvjitlink-cu12-12.8.93 nvidia-nvshmem-cu12-3.4.5 nvidia-nvtx-cu12-12.8.90 sympy-1.14.0 torch-2.11.0.dev20260115+cu128 torchaudio-2.11.0.dev20260115+cu128 torchvision-0.25.0.dev20260115+cu128 triton-3.6.0+git9844da95
代码至此即可成功运行
tips:我之前也试过安装很多版本,包括最新正式版本的torch都不行,但是突然运行这版的成功了(记得重新安装前一定要删除之前安装的版本)