RTX pro 6000 black well最新架构下安装 PyTorch CUDA - 解决 sm_120 兼容性问题

解决 NVIDIA RTX 6000 black well 最新架构下的 PyTorch 版本兼容问题

显卡型号、CUDA、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都不行，但是突然运行这版的成功了（记得重新安装前一定要删除之前安装的版本）