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

解决 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都不行,但是突然运行这版的成功了(记得重新安装前一定要删除之前安装的版本)

相关推荐
枫叶林FYL21 小时前
【乳腺癌早期筛查(钼靶X光图像AI识别)】第一章:钼靶AI核心算法架构演进——从2D全视野到3D断层合成与视觉Transformer
人工智能·深度学习
Lethehong21 小时前
Python Selenium全栈指南:从自动化入门到企业级实战
python·selenium·测试工具·自动化
TK云大师-KK21 小时前
TikTok自动化直播遇到内容重复问题?这套技术方案了解一下
大数据·运维·人工智能·矩阵·自动化·新媒体运营·流量运营
姚青&21 小时前
大语言模型与私有部署
人工智能·语言模型·chatgpt
WeeJot嵌入式21 小时前
Meta LSP无数据训练深度解析:语言自我对弈的数学原理与实现
人工智能·机器学习·里氏替换原则
foundbug9991 天前
基于卡尔曼滤波的背景建模与车辆检测(OpenCV实现)
人工智能·opencv·计算机视觉
一个处女座的程序猿1 天前
AI之Tool:Google Stitch的简介、安装和使用方法、案例应用之详细攻略
人工智能·stitch
智算菩萨1 天前
MP3音频编码原理深度解析与Python全参数调优实战:从心理声学模型到LAME编码器精细控制
android·python·音视频
枫叶林FYL1 天前
【自然语言处理 NLP】数学与计算基础(Mathematical & Computational 完整源码实现
人工智能·深度学习·机器学习