🤖 AI

torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch

本文介绍了 torch-sla，这是一个开源 PyTorch 库，通过为跨多种 CPU 和 GPU 后端的不同求解器提供统一的自动微感 API，实现了可微稀疏线性代数的统一，同时借助 O(1) 图伴随框架支持批处理操作和可扩展的分布式多 GPU 执行。

原作者： Mingyuan Chi, Shizheng Wen

发布于 2026-05-07

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Mingyuan Chi, Shizheng Wen

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你正在尝试解决一个巨大而复杂的拼图。在计算机科学和机器学习的世界里，这个拼图通常是一个“稀疏线性系统”——一个巨大的数字网格，其中大多数单元格是空的（零），但少数填充的单元格掌握着一切的关键。

很长一段时间以来，流行的编程工具PyTorch（被许多 AI 研究人员使用）存在一个问题：它非常擅长处理“密集拼图”（每个单元格都填满），但在解决这些巨大且大部分为空的拼图时，同时还能“从错误中学习”，它的表现却糟糕透顶。如果你尝试解决这样一个问题，计算机要么会崩溃，要么运行得极慢，要么会忘记如何从结果中学习。

现在，torch-sla 登场了。这是一个新的开源库，充当这些拼图的“通用翻译”和“超级高效机械师”。以下是它的工作原理，使用简单的类比来说明：

1. “智能调度员”（统一后端）

想象你是一名送货司机，手里拿着一个包裹。你需要将它送到目的地，但路况会变化：有时是小城市街道，有时是巨大的高速公路，有时你甚至身处另一个国家。

以前，你必须为每一次行程手动选择不同的车辆（小任务用自行车，大任务用卡车，水上任务用船）。如果你选错了，就会浪费时间或陷入困境。

torch-sla 就像一个智能调度员。你只需说：“我需要解决这个拼图”，它会自动检查：

你在哪里？（你的计算机使用的是 CPU 还是强大的图形处理器/GPU？）
拼图有多大？（是 100 块的小拼图，还是 1 亿块的大拼图？）
它是什么形状？（是否完全对称？）

基于这些线索，它会立即从车库中的五种不同选项（包括来自 SciPy、NVIDIA 和 PyTorch 本身的工具）中挑选出最佳的“车辆”（求解器）。你无需知道该选哪一个；库会为你完成这一切。

2. “魔法捷径”（伴随微分）

这是该库最神奇的技巧。

旧方法（朴素反向传播）：
想象你在黑暗的迷宫中行走，走了 1,000 步才找到出口。为了从这段旅程中学习，你决定记下你走的每一步、脚的确切位置以及每一刻的气压。当你到达终点时，你有一本写满 1,000 页笔记的笔记本。如果你想回头看看在哪里犯了错，你必须翻阅所有 1,000 页。如果迷宫变大，你的笔记本也会变得巨大，而你的计算机甚至会因为试图容纳这些笔记而耗尽内存（RAM）。

torch-sla 方法（伴随求解器）：
torch-sla 不使用记录每一步的方法，而是利用一个称为隐函数定理的数学原理，采用“魔法捷径”。
这样想：你走一次迷宫。当你到达出口时，你不需要原路返回，而是在脑海中瞬间计算出“反向地图”。你只需要记住起点、终点和迷宫的规则。

结果： 无论迷宫走了 10 步还是 10,000 步，你的“笔记本”都保持同样小的尺寸。
优势： 论文表明，与旧方法相比，这节省了195 倍的内存。它允许你在单台计算机上解决那些以前会导致系统崩溃的巨大拼图。

3. “团队集会”（分布式求解器）

如果拼图大到一台计算机无法容纳怎么办？你需要一台计算机团队（就像一群朋友共同创作一幅巨型壁画）。

通常，当朋友们在壁画的各个部分工作时，他们必须来回喊话以确保边缘匹配。在计算机术语中，这被称为“光环交换”。

问题： 当你试图从最终图像中学习时，旧工具不知道如何正确地将修正信息“喊回”给正确的人。
torch-sla 解决方案： 它教导计算机如何反向传递“修正笔记”。如果计算机 A 向计算机 B 发送了一部分数据，torch-sla 确保当学习发生时，计算机 B 会将修正信息完美地回传给计算机 A。这使得该库能够在三块强大的 GPU 上将规模扩展到4 亿个未知数。

4. 它实际上能做什么？

论文证明，该库可以处理：

直接求解器： 为中等大小的拼图即时破解代码。
迭代求解器： 为巨大的拼图进行逐步猜测。
非线性与特征值求解器： 处理规则随求解过程变化的拼图，或在数据中寻找特定的“振动”（特征值）。
批处理： 同时解决成千上万个此类拼图（就像同时解决 1,000 个不同的迷宫）。

现实世界测试

为了证明其有效性，作者进行了一项名为“逆系数学习”的测试。
想象你看到了阴影的最终形状（解），并想找出是什么物体（输入）投射了它。

他们使用 torch-sla 通过观察物体投射的阴影来“学习”隐藏物体的形状。
该库在不到一分钟内成功以**99.77%**的准确率找出了隐藏物体。
关键的是，研究人员无需编写任何复杂的数学代码来实现学习；他们只需调用一个简单的命令：A.solve(f)。

总结

torch-sla 填补了 PyTorch 生态系统中的巨大空白。它将解决巨大、稀疏数学拼图这一困难且消耗内存的任务变得：

自动化： 自动为工作挑选最佳工具。
高效： 利用“魔法捷径”节省海量内存。
可扩展： 让计算机团队无缝协作。
可学习： 允许 AI 模型直接从这些复杂的数学问题中学习而不会崩溃。

它本质上是缺失的一环，让 PyTorch 用户能够解决那些以前只能在其他更专业的编程环境中才能处理的科学和工程问题。

技术摘要：torch-sla

问题陈述

可微稀疏线性代数是科学机器学习（包括图神经网络（GNN）、非结构化网格上的神经算子以及可微仿真流水线）的基础需求。然而，PyTorch 生态系统缺乏支持这些工作负载的统一库。原生 PyTorch 支持是碎片化的：torch.sparse 提供低级内核以及一个不可微的、仅支持 CPU 的 spsolve（SuperLU 包装器），而 torch.linalg 仅限于稠密矩阵。

现有的变通方法无法满足大规模或复杂训练场景的需求：

中断计算图：调用外部求解器（如 cuDSS、PETSc）会中断自动求导图，迫使用户放弃梯度或为每个问题编写自定义的 torch.autograd.Function 包装器。
低效的微分：朴素地通过迭代求解器的 $k$ 次迭代进行反向传播，会创建一个 $O(k)$ 节点的 computational graph。对于一个具有 $10^6$ 个自由度（DOF）且迭代 1,000 次的问题，存储中间向量大约需要 80 GB 内存，这使得大规模训练不可行。
可扩展性限制：具有 $10^7$ – $10^9$ 个未知数的问题需要分布式域分解。现有的 PyTorch 分布式原语并不原生支持稀疏算子伴随（adjoint）在反向传播中所需的转置 halo 交换。
后端碎片化：直接求解器（如 cuDSS）对于小规模问题（ $<10^5$ DOF）速度很快，但对于更大的系统，由于 $O(n^{1.5})$ 的填充（fill-in）和内存耗尽而表现不佳。迭代求解器具有更好的可扩展性，但在小规模问题上会产生内核启动开销。目前不存在统一的接口，能够根据设备和问题大小在这些异构后端之间进行调度。

方法论

作者提出了 torch-sla，这是一个开源库，通过三个核心组件填补了这一空白：

1. 统一后端抽象

torch-sla 暴露了一个单一的、感知自动求导的 API（.solve、.matvec、.eigsh、.det），该系统根据设备和问题大小在五个可互换的后端之间进行调度：

CPU：SciPy（SuperLU、UMFPACK、CG、BiCGStab）和 Eigen（CG、BiCGStab）。
GPU：NVIDIA cuDSS（LU、Cholesky、LDLT）、CuPy（迭代和直接）以及一个 PyTorch 原生的迭代求解器。
自动调度：系统自动选择后端：对于小到中等规模的问题（例如 GPU 上 $<2 \times 10^6$ DOF）使用直接求解器，对于更大规模的问题使用迭代求解器以避免内存溢出。系统会检测对称性和正定性以优化分解（例如将 LU 升级为 Cholesky）。

该库支持四种按局部性和批处理组织的稀疏张量类型：SparseTensor（单个矩阵/共享模式）、SparseTensorList（每个批次不同的模式）以及它们的分布式变体 DSparseTensor 和 DSparseTensorList。

2. 伴随微分框架

为了避免朴素反向传播带来的 $O(k)$ 内存和图复杂度，torch-sla 采用了隐函数定理方法。

线性系统：对于 $Ax=b $，反向传播求解单个伴随系统$ A^\top \lambda = \partial L / \partial x $。关于矩阵的梯度仅在稀疏模式上计算为$ -\lambda_i x_j$。
非线性系统：对于 $F(u, \theta)=0$ ，反向传播需要一个伴随线性求解（ $J^\top \lambda = \partial L / \partial u$ ）和一个向量 - 雅可比乘积，这与非线性迭代的次数无关。
特征值问题：特征值的梯度利用 Hellmann–Feynman 定理（ $\partial \lambda / \partial A_{ij} = v_i v_j$ ），仅需在稀疏模式上进行外积运算。
复杂度：该框架将自动求导图减少为 $O(1)$ 个节点，内存减少为 $O(n + \text{nnz})$ ，与求解器迭代次数（ $k$ ）无关。

3. 带有 Halo 交换的分布式层

对于多 GPU 执行，torch-sla 实现了域分解，其中每个进程拥有一个行块和 halo 元数据。

前向传播：进程在本地稀疏矩阵 - 向量（SpMV）操作之前交换 halo 值。
反向传播：该库实现了转置 halo 交换。虽然前向传播将拥有的值散射到相邻的 halo 中，但反向传播将传入的 halo 梯度求和并累加回拥有进程的边界条目中。
实现：前向和反向通信均使用 torch.distributed 封装为 torch.autograd.Function 原语，确保分布式求解的端到端可微性。

主要贡献

统一 API：一个单一的库，为直接、迭代、非线性和特征值求解器提供感知自动求导的接口，涵盖五个后端（SciPy、Eigen、cuDSS、CuPy、PyTorch 原生）。
可扩展的伴随框架：一种 $O(1)$ 图的微分方法，消除了迭代求解器反向传播的内存瓶颈，使得在具有数百万未知数的问题上进行训练成为可能。
分布式自动求导：PyTorch 中首个分布式稀疏求解器的实现，通过转置 halo 交换原生支持端到端微分。
可扩展性：一种后端抽象，允许通过实现通用接口来集成外部库（如 PETSc、Trilinos、hypre）或学习到的预条件子。

实验结果

作者在 NVIDIA H200 GPU 上对 2D 泊松基准测试评估了 torch-sla：

可扩展性：该库在单张 GPU 上可扩展至 1.69 亿 DOF，在三张 GPU 上可扩展至 4 亿 DOF。
性能：对于 100 万 DOF 的问题，torch-sla CG 后端（190 毫秒）的性能优于 cuDSS（7.3 秒）38 倍。在 200 万 DOF 以上，直接求解器会耗尽内存，而迭代求解器保持近线性的扩展性。
内存效率：在针对 1,000 次 CG 迭代的朴素反向传播与伴随方法的比较中，朴素方法需要约 64 GB 内存（在 2,000 次迭代时内存溢出），而伴随方法保持恒定的 ~328 MB 使用量，实现了 195 倍 的内存减少。
梯度验证：针对非线性和特征值问题，将梯度与有限差分进行了验证，显示相对误差低于 $10^{-5}$ 。
端到端应用：一个逆系数学习任务（恢复泊松方程中空间变化的电导率）在标准 PyTorch Adam 循环下，在一分钟内成功收敛至 0.23% 的相对 L2 误差。

意义

该论文声称，torch-sla 填补了 PyTorch 生态系统中长期存在的空白，提供了类似于 JAX 的 jax.lax.custom_linear_solve 和 jax.scipy.sparse.linalg 的功能，但位于统一的 PyTorch 库内。它使研究人员能够：

训练涉及稀疏线性求解的模型，而无需放弃梯度或将代码移植到 JAX。
将科学机器学习工作负载扩展到分布式 GPU 集群上的数亿个自由度。
将稀疏求解器视为可微层，促进学习预条件子和物理仿真端到端优化的研究。

作者指出了局限性，包括目前 PyTorch 原生后端依赖雅可比（Jacobi）预条件子（这限制了非常大系统的收敛速度），以及专注于 2D 泊松基准测试，其中 3D 偏微分方程（PDE）和不定系统被确定为未来的工作。