N卡环境下调用torch.sparse.spsolve遇ROCm相关错误求助

核心问题 #

错误提示中的ROCm相关内容属于文案误导，实际问题是你的PyTorch CUDA版本未编译cuDSS库——torch.sparse.spsolve在CUDA设备上运行需要依赖NVIDIA的cuDSS（CUDA Sparse Solver）组件，但并非所有PyTorch预编译包都默认包含该组件。

解决方案 #

1. 更换/重新编译支持cuDSS的PyTorch版本 #

• 推荐使用PyTorch Nightly版本：PyTorch 2.5.x正式版预编译包可能未全面启用cuDSS，Nightly开发版通常包含最新CUDA组件支持，可通过以下命令安装：

  pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu124
  

• 手动编译PyTorch：若需要稳定版，可从源码编译时启用cuDSS支持，需确保已安装CUDA Toolkit 12.4及以上版本，编译参数添加USE_CUDNN=1和USE_CUDSS=1。

2. 使用替代方法求解稀疏线性方程组 #

如果不想更换PyTorch版本，可采用以下替代方案：

方案A：借助SciPy的稀疏求解器

将PyTorch张量转换为SciPy稀疏矩阵求解，再转回PyTorch张量：

import scipy.sparse as sp
import scipy.sparse.linalg as spla

# 将CUDA稀疏张量转回CPU并转为SciPy CSR矩阵
A_scipy = sp.csr_matrix(
    (A_sparse_csr.cpu().values(), 
     A_sparse_csr.cpu().crow_indices(), 
     A_sparse_csr.cpu().col_indices()), 
    shape=A_sparse_csr.shape
)
b_np = b.cpu().numpy()

# 求解线性方程组
intercepts_np = spla.spsolve(A_scipy, b_np)

# 转回PyTorch CUDA张量
intercepts = torch.tensor(intercepts_np, dtype=b.dtype, device=b.device)

方案B：迭代法求解（适合对称正定矩阵）

若稀疏矩阵为对称正定矩阵，可使用共轭梯度（CG）迭代法实现求解：

def sparse_cg(A, b, max_iter=1000, tol=1e-6):
    x = torch.zeros_like(b)
    r = b - A @ x
    p = r.clone()
    r_dot_r = torch.dot(r, r)
    
    for _ in range(max_iter):
        Ap = A @ p
        alpha = r_dot_r / torch.dot(p, Ap)
        x += alpha * p
        r -= alpha * Ap
        new_r_dot_r = torch.dot(r, r)
        if torch.sqrt(new_r_dot_r) < tol:
            break
        beta = new_r_dot_r / r_dot_r
        p = r + beta * p
        r_dot_r = new_r_dot_r
    return x

# 调用求解（确保A是对称正定矩阵）
intercepts = sparse_cg(A_sparse_csr, b)

方案C：转为稠密矩阵求解（仅适用于小规模矩阵）

如果矩阵规模较小，可直接转换为稠密矩阵后使用torch.linalg.solve：

A_dense = A_sparse_csr.to_dense()
intercepts = torch.linalg.solve(A_dense, b)

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