PyTorch张量移至GPU触发Segmentation Fault问题求助

• 硬件：配备RTX 6000 ADA GPU的机器
• 背景：此前使用CUDA 11.x，采用镜像nvcr.io/nvidia/pytorch:21.04-py3（搭载PyTorch 1.x）；机器驱动更新后，nvidia-smi显示信息：NVIDIA-SMI 535.161.07 Driver Version: 535.161.07 CUDA Version: 12.2
• 故障：将张量移至CUDA设备时触发Segmentation fault (core dumped)错误，更换镜像nvcr.io/nvidia/pytorch:23.02-py3后问题依旧
• 测试代码：

import torch

tensor_cpu = torch.tensor([1, 2, 3, 4, 5])

print("Tensor device before moving to CUDA:", tensor_cpu.device)

if torch.cuda.is_available():
    device = torch.device("cuda")          
    tensor_cuda = tensor_cpu.to(device)    
    print("Tensor device after moving to CUDA:", tensor_cuda.device)
else:
    print("CUDA is not available. Cannot move tensor to CUDA.")

• 输出：

Tensor device before moving to CUDA: cpu
Segmentation fault (core dumped)

• 疑问：如何解决该问题？适配RTX 6000 ADA + PyTorch 1.x的正确镜像是什么？是否存在其他影响因素？

1. 核心问题定位 #

RTX 6000 ADA属于Ada Lovelace架构GPU，PyTorch 1.x全系列不支持该架构——PyTorch从2.0版本开始才正式加入对Ada Lovelace核心的CUDA指令支持。这是触发段错误的根本原因，和CUDA版本匹配关系不大，旧版PyTorch无法识别新架构GPU的硬件指令集。

2. 适配方案选择 #

方案A：升级PyTorch以适配新架构 #

放弃PyTorch 1.x，使用支持Ada架构的版本：

• 匹配当前驱动（CUDA 12.2），可选择nvcr.io/nvidia/pytorch:23.10-py3或更高版本（搭载PyTorch 2.1+，原生支持CUDA 12.2和Ada架构）
• 启动容器后先验证基础兼容性：

  print(torch.version.cuda)
  print(torch.cuda.get_device_name(0))
  

  若能正确输出CUDA版本和GPU名称，再执行张量迁移测试。

方案B：回退驱动兼容PyTorch 1.x #

若必须保留PyTorch 1.x，需回退GPU驱动到支持CUDA 11.x且兼容Ada架构的版本：

• Ada架构最低支持的CUDA 11.x版本是CUDA 11.8，对应驱动版本范围为520.61.05~535.104.05
• 驱动回退后，使用镜像nvcr.io/nvidia/pytorch:22.12-py3（搭载PyTorch 1.13.1，支持CUDA 11.8，适配Ada架构）

3. 其他排查要点 #

• 检查容器启动参数：确保启动时传递GPU权限，使用--gpus all参数（如docker run --gpus all -it nvcr.io/nvidia/pytorch:xxx-py3）
• 验证宿主机GPU状态：在宿主机执行nvidia-smi确认GPU无异常占用，nvidia-debugdump -l检查硬件健康状态
• 排查容器内库冲突：执行ldd $(which python) | grep cuda，确认加载的CUDA库版本与镜像内置版本一致

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