Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Mingyuan Chi

TL;DR

torch-sla addresses differentiable sparse linear algebra at industrial scale by fusing adjoint-based differentiation, distributed sparse tensor parallelism, and PyTorch integration. It achieves backward passes with $O(1)$ graph nodes and $O(\text{nnz})$ memory for linear, eigenvalue, and nonlinear problems, while enabling domain-decomposition with halo exchange that scales to $4\times 10^8$ DOF on $3$ GPUs. The experiments demonstrate GPU-accelerated solvers across backends and gradient verification showing high accuracy ($<10^{-5}$ relative error), enabling end-to-end differentiable simulations in PyTorch for unstructured sparse systems. The library provides a practical pathway to scalable physics-informed learning and neural operators on large sparse meshes and graphs.

Abstract

Industrial scientific computing predominantly uses sparse matrices to represent unstructured data -- finite element meshes, graphs, point clouds. We present \torchsla{}, an open-source PyTorch library that enables GPU-accelerated, scalable, and differentiable sparse linear algebra. The library addresses three fundamental challenges: (1) GPU acceleration for sparse linear solves, nonlinear solves (Newton, Picard, Anderson), and eigenvalue computation; (2) Multi-GPU scaling via domain decomposition with halo exchange, reaching \textbf{400 million DOF linear solve on 3 GPUs}; and (3) Adjoint-based differentiation} achieving $\mathcal{O}(1)$ computational graph nodes (for autograd) and $\mathcal{O}(\text{nnz})$ memory -- independent of solver iterations. \torchsla{} supports multiple backends (SciPy, cuDSS, PyTorch-native) and seamlessly integrates with PyTorch autograd for end-to-end differentiable simulations. Code is available at https://github.com/walkerchi/torch-sla.

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

TL;DR

torch-sla addresses differentiable sparse linear algebra at industrial scale by fusing adjoint-based differentiation, distributed sparse tensor parallelism, and PyTorch integration. It achieves backward passes with graph nodes and memory for linear, eigenvalue, and nonlinear problems, while enabling domain-decomposition with halo exchange that scales to DOF on GPUs. The experiments demonstrate GPU-accelerated solvers across backends and gradient verification showing high accuracy ( relative error), enabling end-to-end differentiable simulations in PyTorch for unstructured sparse systems. The library provides a practical pathway to scalable physics-informed learning and neural operators on large sparse meshes and graphs.

Abstract

Industrial scientific computing predominantly uses sparse matrices to represent unstructured data -- finite element meshes, graphs, point clouds. We present \torchsla{}, an open-source PyTorch library that enables GPU-accelerated, scalable, and differentiable sparse linear algebra. The library addresses three fundamental challenges: (1) GPU acceleration for sparse linear solves, nonlinear solves (Newton, Picard, Anderson), and eigenvalue computation; (2) Multi-GPU scaling via domain decomposition with halo exchange, reaching \textbf{400 million DOF linear solve on 3 GPUs}; and (3) Adjoint-based differentiation} achieving computational graph nodes (for autograd) and memory -- independent of solver iterations. \torchsla{} supports multiple backends (SciPy, cuDSS, PyTorch-native) and seamlessly integrates with PyTorch autograd for end-to-end differentiable simulations. Code is available at https://github.com/walkerchi/torch-sla.
Paper Structure (23 sections, 6 equations, 3 figures, 4 tables)

This paper contains 23 sections, 6 equations, 3 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Preliminaries: Sparse Linear Systems
  5. Adjoint-Based Differentiation
  6. Linear Systems
  7. Eigenvalue Problems
  8. Nonlinear Systems
  9. Sparse Tensor Parallel Computing
  10. Domain Decomposition
  11. Halo Exchange
  12. Distributed Conjugate Gradient
  13. Experiments
  14. Single-GPU Scalability
  15. Multi-GPU Performance
  16. ...and 8 more sections

Figures (3)

  • Figure 1: Halo exchange in domain decomposition. Each process owns a subset of nodes (solid colored) and maintains halo copies of boundary neighbors (dashed). Before SpMV, processes exchange updated values at partition boundaries via peer-to-peer communication.
  • Figure 2: Single-GPU benchmark results showing performance, memory usage, and residual across solver backends (SciPy, cuDSS, PyTorch CG) on 2D Poisson equation with H200 GPU.
  • Figure 3: Multi-GPU scaling: distributed CG with NCCL backend on 3--4 H200 GPUs, showing time and memory scaling across problem sizes.