torch-sla - PyTorch 稀疏线性代数库

torch-sla: PyTorch 稀疏线性代数

torch-sla (Torch Sparse Linear Algebra) 是一个高效、可微分的 PyTorch 稀疏线性方程求解器库，支持多种后端。适用于科学计算、有限元、计算流体动力学和需要自动微分的稀疏矩阵操作的机器学习应用。

arXiv GitHub PyPI License: MIT

为什么选择 torch-sla？

  • 🚀 高性能: 通过 cuSOLVER 和 cuDSS 实现 CUDA 加速求解
  • 💾 内存高效: 仅存储非零元素，支持求解百万级未知数的系统
  • 🔄 可微分: 完整支持 torch.autograd 梯度计算
  • 📦 批量处理: 并行求解数千个系统
  • 🌐 分布式计算: 支持域分解和 halo 交换的大规模计算
  • 🔧 灵活配置: 多种后端和求解方法

核心特性

  • 内存高效: 仅存储非零元素 — 1M×1M 矩阵（1% 密度）仅需 ~80MB 而非 ~8TB
  • 完整梯度支持: 通过 torch.autograd 实现端到端可微分流程
  • 多后端支持: SciPy、Eigen、cuSOLVER、cuDSS
  • 批量求解: 支持相同和不同稀疏模式的批量矩阵
  • 分布式求解: 域分解与 halo 交换
  • 大规模测试: 经过 1.69亿+ DOF 测试，近线性复杂度扩展

快速开始

安装

pip install torch-sla

基本用法

import torch
from torch_sla import SparseTensor

# 从稠密矩阵创建稀疏矩阵（小矩阵更易读）
dense = torch.tensor([[4.0, -1.0,  0.0],
                      [-1.0, 4.0, -1.0],
                      [ 0.0, -1.0, 4.0]], dtype=torch.float64)

A = SparseTensor.from_dense(dense)

# 求解 Ax = b
b = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
x = A.solve(b)

CUDA 加速

# 移动到 GPU 进行 CUDA 加速求解
A_cuda = A.cuda()
b_cuda = b.cuda()
x = A_cuda.solve(b_cuda)  # 自动使用 cuDSS 或 cuSOLVER

应用场景

torch-sla 非常适合:

  • 有限元方法 (FEM): 求解来自 FEM 离散化的大型稀疏系统

  • 计算流体动力学 (CFD): Navier-Stokes 方程的高效稀疏求解器

  • 物理信息神经网络 (PINNs): 物理约束的可微分稀疏操作

  • 图神经网络 (GNN): 稀疏消息传递和拉普拉斯操作

  • 优化问题: 涉及稀疏线性系统的梯度优化

常见问题 (FAQ)

什么是 torch-sla？

torch-sla (Torch Sparse Linear Algebra) 是一个 Python 库，提供可微分的 PyTorch 稀疏线性方程求解器。它求解 Ax = b 形式的系统，其中 A 是稀疏矩阵，完整支持自动微分 (autograd) 和 CUDA GPU 加速。

如何在 PyTorch 中求解稀疏线性系统？

使用 torch-sla 的 SparseTensor 类:

from torch_sla import SparseTensor

# 从 COO 格式创建稀疏矩阵（值、行索引、列索引）
A = SparseTensor(values, row, col, shape)

# 求解 Ax = b
x = A.solve(b)

支持 CPU 和 GPU，并支持梯度计算。

torch-sla 支持哪些稀疏求解器？

torch-sla 支持多种后端:

  • CPU: SciPy (SuperLU, UMFPACK, CG, BiCGStab, GMRES), Eigen (CG, BiCGStab)

  • GPU: cuSOLVER (QR, Cholesky, LU), cuDSS (LU, Cholesky, LDLT)

库会根据硬件和矩阵属性自动选择最佳求解器。

能否通过稀疏求解计算梯度？

可以。torch-sla 完整支持 PyTorch autograd:

val = torch.tensor([...], requires_grad=True)
x = spsolve(val, row, col, shape, b)
loss = x.sum()
loss.backward()  # 计算 val 和 b 的梯度

如何求解批量稀疏系统？

torch-sla 支持相同稀疏模式矩阵的批量求解:

# 批量值: [batch_size, nnz]
A = SparseTensor(val_batch, row, col, (batch_size, M, N))
x = A.solve(b_batch)  # 并行求解所有系统

对于不同模式的矩阵，使用 SparseTensorList。参见 批量求解示例。

如何在 GPU 上使用 torch-sla？

只需将张量移动到 CUDA:

A_cuda = A.cuda()
x = A_cuda.solve(b.cuda())  # 使用 cuDSS 或 cuSOLVER

SparseTensor 和 DSparseTensor 有什么区别？

  • SparseTensor: 单个稀疏矩阵（可选批量），用于标准求解

  • DSparseTensor: 分布式稀疏张量，支持域分解和 halo 交换，用于大规模并行计算

与其他方案的对比

torch-sla vs scipy.sparse.linalg

特性

torch-sla ✅

scipy.sparse.linalg

PyTorch 集成

✅ 原生张量

❌ 需要 numpy 复制

GPU 加速

✅ CUDA (cuDSS, cuSOLVER)

❌ 仅 CPU

Autograd 梯度

✅ 完整支持（伴随法）

❌ 无梯度

批量求解

✅ 并行批量求解

❌ 需要循环

大规模 (>200万 DOF)

✅ 1.69亿 DOF 已测试

⚠️ 内存受限

分布式计算

✅ DSparseTensor

❌ 不支持

特征值/SVD

✅ 可微分

⚠️ 无梯度

非线性求解

✅ Newton/Anderson

❌ 不包含

torch-sla vs torch.linalg.solve

特性

torch-sla ✅

torch.linalg.solve

矩阵类型

✅ 稀疏 (COO/CSR)

❌ 仅稠密

内存 (1M×1M, 1% 密度)

✅ ~80 MB

❌ ~8 TB（不可能）

最大问题规模

✅ 5亿+ DOF （多卡可扩展）

❌ ~5万（GPU 内存限制）

专用求解器

✅ LU, Cholesky, CG, BiCGStab

⚠️ 仅稠密 LU

批量操作

✅ 相同/不同模式

⚠️ 仅相同形状

GPU 支持

✅ cuDSS, cuSOLVER, PyTorch

✅ 是

Autograd

✅ O(1) 图节点

✅ 是

torch-sla vs NVIDIA AmgX

特性

torch-sla ✅

NVIDIA AmgX

安装

✅ pip install torch-sla

❌ 复杂的构建过程

PyTorch 集成

✅ 原生

❌ 需要封装

Autograd 支持

✅ 完整梯度流

❌ 无梯度

Python API

✅ Pythonic

⚠️ 以 C++ 为主

多重网格 (AMG)

❌ 暂不支持

✅ 核心功能

预处理器

⚠️ Jacobi

✅ ILU, AMG 等

文档

✅ 完善

⚠️ 示例有限

torch-sla vs PETSc

特性

torch-sla ✅

PETSc

安装

✅ pip install

❌ 复杂（MPI、编译器）

学习曲线

✅ 简单 Python API

❌ 陡峭（C/Fortran 背景）

PyTorch 集成

✅ 原生张量

❌ 需要 petsc4py + 复制

Autograd

✅ 完整支持

❌ 无梯度

求解器种类

⚠️ 基础求解器

✅ 丰富 (KSP, SNES)

分布式

✅ DSparseTensor 多卡并行

✅ 完整 MPI 支持

生产规模

✅ 5亿+ DOF (多卡)

✅ 百亿级已验证

总结：何时使用 torch-sla

使用 torch-sla 当

考虑其他方案当

✅ 需要 PyTorch 集成

不使用 PyTorch

✅ 需要 梯度流 通过求解

不需要梯度

✅ 问题规模 5亿+ DOF (多卡)

百亿级问题（使用 PETSc）

✅ 想要 简单的 pip 安装

需要 AMG 预处理器（AmgX）

✅ 批量 稀疏系统

复杂预处理（PETSc）

✅ GPU 加速 且配置简单

完整 MPI 分布式（PETSc）

索引和搜索

许可证

torch-sla 采用 MIT 许可证发布。详见 LICENSE。

联系方式

作者: Walker Chi
邮箱: x@y，其中 x = walker.chi.000，y = gmail.com

引用

如果您在研究中使用了 torch-sla，请引用我们的论文:

@article{chi2026torchsla,
  title={torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch},
  author={Chi, Mingyuan},
  journal={arXiv preprint arXiv:2601.13994},
  year={2026},
  url={https://arxiv.org/abs/2601.13994}
}

论文: arXiv:2601.13994 - Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch