🤖 AI

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

Dit artikel introduceert torch-sla, een open-source PyTorch-bibliotheek die differentieerbare sparse lineaire algebra verenigt door een enkele autograd-bewuste API te bieden voor diverse oplossers op meerdere CPU- en GPU-backends, terwijl het batchbewerkingen en schaalbare gedistribueerde multi-GPU-uitvoering ondersteunt via een O(1)-graafadjoint-framework.

Oorspronkelijke auteurs: Mingyuan Chi, Shizheng Wen

Gepubliceerd 2026-05-07

📖 6 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Mingyuan Chi, Shizheng Wen

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je probeert een enorm, complex puzzel op te lossen. In de wereld van de informatica en machine learning is deze puzzel vaak een "spaarzaam lineair stelsel" – een gigantisch raster van getallen waarbij de meeste cellen leeg (nul) zijn, maar de weinige ingevulde cellen de sleutel tot alles bevatten.

Lange tijd had het populaire programmeergereedschap PyTorch (dat door veel AI-onderzoekers wordt gebruikt) een probleem: het was uitstekend in het hanteren van dichte puzzels (waarbij elke cel is ingevuld), maar het was vreselijk in het oplossen van deze gigantische, grotendeels lege puzzels terwijl het ook uit de fouten leerde. Als je probeerde er één op te lossen, zou de computer ofwel crashen, ongelofelijk traag werken, of vergeten hoe het uit het resultaat moest leren.

Dan komt torch-sla in beeld, een nieuwe open-source bibliotheek die fungeert als een universele vertaler en een super-efficiënte monteur voor deze puzzels. Hier is hoe het werkt, met behulp van eenvoudige analogieën:

1. De "Slimme Dispatcher" (Gecombineerde Backend)

Stel je voor dat je een bezorger bent met een pakket. Je moet het naar een bestemming brengen, maar de verkeersomstandigheden veranderen: soms is het een kleine stadsstraat, soms een enorme snelweg, en soms bevind je je helemaal in een ander land.

Voorheen moest je voor elke enkele reis handmatig een ander voertuig kiezen (een fiets voor kleine klussen, een vrachtwagen voor grote, een boot voor water). Als je de verkeerde koos, waste je tijd of bleef je steken.

torch-sla is als een slimme dispatcher. Je zegt gewoon: "Ik moet deze puzzel oplossen", en het controleert automatisch:

Waar ben je? (Gebruikt je computer een CPU of een krachtige grafische kaart/GPU?)
Hoe groot is de puzzel? (Is het een klein puzzeltje van 100 stukjes of een van 100 miljoen stukjes?)
Welke vorm heeft het? (Is het perfect symmetrisch?)

Op basis van deze aanwijzingen kiest het direct de beste "voertuig" (oplosser) uit een garage met vijf verschillende opties (waaronder tools van SciPy, NVIDIA en PyTorch zelf). Je hoeft niet te weten welke je moet kiezen; de bibliotheek doet dit voor jou.

2. De "Magische Kortsluiting" (Adjoint Differentiatie)

Dit is de meest magische truc van de bibliotheek.

De Oude Manier (Naïeve Backpropagatie):
Stel je voor dat je door een donker doolhof loopt, waarbij je 1.000 stappen doet om de uitgang te vinden. Om te leren van je reis, besluit je elke enkele stap die je hebt genomen, de exacte positie van je voeten en de luchtdruk op elk moment op te schrijven. Als je aan het einde bent, heb je een notitieboek met 1.000 pagina's notities. Als je terug wilt gaan om te zien waar je een fout hebt gemaakt, moet je door al die 1.000 pagina's bladeren. Als het doolhof groter wordt, wordt je notitieboek enorm, en raakt je computer de geheugenruimte (RAM) kwijt terwijl hij alleen al probeert de notities vast te houden.

De torch-sla Manier (Adjoint Oplosser):
In plaats van elke stap op te schrijven, gebruikt torch-sla een "magische kortsluiting" gebaseerd op een wiskundig principe genaamd de Stelling van de Impliciete Functie.
Denk er zo over: Je loopt het doolhof één keer door. Als je de uitgang bereikt, bereken je in plaats van je stappen terug te doen, direct de "omgekeerde kaart" in je hoofd. Je hoeft alleen maar het begin, het einde en de regels van het doolhof te onthouden.

Resultaat: Of het doolhof nu 10 stappen of 10.000 stappen duurde, je "notitieboek" blijft even groot.
Voordeel: Het artikel toont aan dat dit 195 keer meer geheugen bespaart dan de oude manier. Het stelt je in staat om enorme puzzels op één computer op te lossen die eerder het systeem zouden hebben laten crashen.

3. De "Teamhuddle" (Gedistribueerde Oplossers)

Wat als de puzzel zo groot is dat één computer het niet kan vasthouden? Je hebt een team van computers nodig (zoals een groep vrienden die werken aan een gigantisch muurschildering).

Meestal moeten vrienden die aan aparte delen van een muurschildering werken, heen en weer schreeuwen om ervoor te zorgen dat de randen overeenkomen. In computertaal heet dit "halo-uitwisseling".

Het Probleem: Als je probeert te leren van het eindbeeld, wisten de oude tools niet hoe ze de correcties correct naar de juiste personen moesten "terugschreeuwen".
De torch-sla Oplossing: Het leert de computers hoe ze de "correctienotities" in omgekeerde richting moeten doorgeven. Als Computer A een stuk data naar Computer B stuurt, zorgt torch-sla ervoor dat, wanneer het leren plaatsvindt, Computer B de correctie terug naar Computer A stuurt, perfect. Dit stelt de bibliotheek in staat om op te schalen tot 400 miljoen onbekenden verspreid over drie krachtige GPU's.

4. Wat Kan Het Eigenlijk?

Het artikel demonstreert dat deze bibliotheek het volgende aankan:

Directe Oplossers: De code direct kraken voor middelgrote puzzels.
Iteratieve Oplossers: Stap-voor-stap gissingen doen voor enorme puzzels.
Niet-lineaire & Eigenwaarde Oplossers: Puzzels aanpakken waarbij de regels veranderen terwijl je ze oplost, of specifieke "trillingen" (eigenwaarden) in de data vinden.
Batchverwerking: Duizenden van deze puzzels tegelijkertijd oplossen (zoals het oplossen van 1.000 verschillende doolhoven tegelijk).

De Realiteitstest

Om te bewijzen dat het werkt, draaiden de auteurs een test genaamd "Inverse Coëfficiëntenleren".
Stel je voor dat je de uiteindelijke vorm van een schaduw ziet (de oplossing) en je wilt uitzoeken welk object (de invoer) deze heeft geworpen.

Ze gebruikten torch-sla om de vorm van een verborgen object te "leren" door te kijken naar de schaduw die het wierp.
De bibliotheek slaagde erin het verborgen object te achterhalen met 99,77% nauwkeurigheid in minder dan een minuut.
Cruciaal: de onderzoeker hoefde geen complexe wiskundecode te schrijven om het leren te laten gebeuren; ze riep gewoon een eenvoudige opdracht aan: A.solve(f).

Samenvatting

torch-sla vult een enorm gat in het PyTorch-ecosysteem. Het neemt de moeilijke, geheugenverslindende taak van het oplossen van gigantische, spaarzame wiskundepuzzels en maakt het:

Automatisch: Het kiest het beste gereedschap voor de klus.
Efficiënt: Het gebruikt een "magische kortsluiting" om enorme hoeveelheden geheugen te besparen.
Schaalbaar: Het stelt teams van computers in staat naadloos samen te werken.
Leerbaar: Het stelt AI-modellen in staat om direct uit deze complexe wiskundeproblemen te leren zonder te breken.

Het is in wezen het ontbrekende schakel dat PyTorch-gebruikers in staat stelt om dezelfde soorten wetenschappelijke en technische problemen aan te pakken die voorheen alleen mogelijk waren in andere, meer gespecialiseerde programmeeromgevingen.

Technische Samenvatting: torch-sla

Probleemstelling

Differentieerbare lineaire algebra voor dunne matrices is een fundamentele vereiste voor wetenschappelijk machine learning, waaronder grafische neurale netwerken (GNNs), neurale operatoren op ongestructureerde meshes en differentieerbare simulatiepijplijnen. Het PyTorch-ecosysteem mist echter een geünificeerde bibliotheek om deze workloads te ondersteunen. De native PyTorch-ondersteuning is gefragmenteerd: torch.sparse biedt low-level kernels en een niet-differentieerbare, alleen-CPU-gebaseerde spsolve (een SuperLU-wrapper), terwijl torch.linalg beperkt is tot dichte matrices.

Bestaande oplossingen zijn ontoereikend voor grootschalige of complexe trainingsscenario's:

Onderbreken van de Graaf: Het aanroepen van externe oplosprogramma's (bijv. cuDSS, PETSc) onderbreekt de autograd-graaf, waardoor gebruikers genoodzaakt zijn om gradients te laten vallen of aangepaste torch.autograd.Function-wrappers te schrijven voor elk probleem.
Inefficiënte Differentiatie: Het naief terugpropageren door $k$ iteraties van een iteratief oplosprogramma creëert een computatiegraaf met $O(k)$ knopen. Voor een probleem met $10^6$ vrijheidsgraden (DOF) en 1.000 iteraties vereist het opslaan van tussentijdse vectoren ongeveer 80 GB geheugen, waardoor grootschalige training onhaalbaar wordt.
Schalingsgrenzen: Problemen met $10^7$ – $10^9$ onbekenden vereisen gedistribueerde domeindecompositie. Bestaande gedistribueerde primitieven in PyTorch ondersteunen niet op native wijze de getransponeerde halo-uitwisseling die nodig is voor geadjungeerde operatoren van dunne matrices in de backward-pass.
Fragmentatie van Backends: Directe oplosprogramma's (bijv. cuDSS) zijn snel voor kleine problemen ( $<10^5$ DOF), maar lijden onder $O(n^{1.5})$ opvulling en geheugenuitputting voor grotere systemen. Iteratieve oplosprogramma's schalen beter, maar brengen kernel-startkosten met zich mee bij kleine problemen. Er bestaat geen geünificeerde interface om te dispatchen over deze heterogene backends op basis van apparaat en probleemgrootte.

Methodologie

De auteurs presenteren torch-sla, een open-source bibliotheek die deze kloof dicht via drie kerncomponenten:

1. Geünificeerde Backend-abstractie

torch-sla blootst een enkele autograd-bewuste API (.solve, .matvec, .eigsh, .det) die dispatcht over vijf uitwisselbare backends op basis van apparaat en probleemgrootte:

CPU: SciPy (SuperLU, UMFPACK, CG, BiCGStab) en Eigen (CG, BiCGStab).
GPU: NVIDIA cuDSS (LU, Cholesky, LDLT), CuPy (iteratief en direct) en een PyTorch-native iteratief oplosprogramma.
Auto-Dispatch: Het systeem selecteert automatisch de backend: directe oplosprogramma's voor kleine tot middelgrote problemen (bijv. $<2 \times 10^6$ DOF op GPU) en iteratieve oplosprogramma's voor grotere problemen om geheugenuitputting te voorkomen. Symmetrie en positief-definite aard worden gedetecteerd om factorisatie te optimaliseren (bijv. het upgraden van LU naar Cholesky).

De bibliotheek ondersteunt vier soorten dunne tensoren, georganiseerd op basis van localiteit en batching: SparseTensor (één matrix/gedeeld patroon), SparseTensorList (verschillende patronen per batch) en hun gedistribueerde varianten DSparseTensor en DSparseTensorList.

2. Geadjungeerd Differentiatieframework

Om de $O(k)$ geheugen- en graafcomplexiteit van naieve backpropagation te vermijden, hanteert torch-sla een aanpak gebaseerd op de stelling van de impliciete functie.

Lineaire Systemen: Voor $Ax=b$ lost de backward-pass een enkel geadjungeerd systeem op $A^\top \lambda = \partial L / \partial x$ . De gradient met betrekking tot de matrix wordt berekend als $-\lambda_i x_j$ alleen op het sparsiteitspatroon.
Niet-lineaire Systemen: Voor $F(u, \theta)=0$ vereist de backward-pass één geadjungeerd lineair oplossen ( $J^\top \lambda = \partial L / \partial u$ ) en een vector-Jacobi-product, onafhankelijk van het aantal niet-lineaire iteraties.
Eigenwaardeproblemen: Gradienten voor eigenwaarden maken gebruik van de Hellmann–Feynman-stelling ( $\partial \lambda / \partial A_{ij} = v_i v_j$ ), waarbij alleen een uitwendig product op het sparsiteitspatroon vereist is.
Complexiteit: Dit framework reduceert de autograd-graaf tot $O(1)$ knopen en het geheugen tot $O(n + \text{nnz})$ , onafhankelijk van het aantal oplositeraties ( $k$ ).

3. Gedistribueerde Laag met Halo-uitwisseling

Voor multi-GPU-uitvoering implementeert torch-sla domeindecompositie waarbij elk proces een rijblok en halo-metadata bezit.

Forward Pass: Processen wisselen halo-waarden uit voordat lokale Sparse Matrix-Vector (SpMV)-operaties worden uitgevoerd.
Backward Pass: De bibliotheek implementeert een getransponeerde halo-uitwisseling. Waar de forward-pass bezette waarden verstrooit naar naburige halos, sommeert de backward-pass inkomende halo-gradienten terug in de grensentries van het bezittende proces.
Implementatie: Zowel forward- als backward-communicatie zijn gewikkeld als torch.autograd.Function-primitieven met behulp van torch.distributed, wat eind-tot-eind differentieerbaarheid over gedistribueerde oplossingen garandeert.

Belangrijkste Bijdragen

Geünificeerde API: Één bibliotheek die autograd-bewuste interfaces biedt voor directe, iteratieve, niet-lineaire en eigenwaarde-oplosprogramma's over vijf backends (SciPy, Eigen, cuDSS, CuPy, PyTorch-native).
Schalbaar Geadjungeerd Framework: Een differentiatiemethode met $O(1)$ -graaf die de geheugenbottleneck van backpropagation via iteratieve oplosprogramma's elimineert, waardoor training mogelijk wordt op problemen met miljoenen onbekenden.
Gedistribueerde Autograd: De eerste implementatie van gedistribueerde oplosprogramma's voor dunne matrices in PyTorch die op native wijze eind-tot-eind differentiatie ondersteunt via getransponeerde halo-uitwisseling.
Uitbreidbaarheid: Een backend-abstractie die het mogelijk maakt externe bibliotheken (bijv. PETSc, Trilinos, hypre) of geleerde preconditioners te integreren door een gemeenschappelijke interface te implementeren.

Experimentele Resultaten

De auteurs hebben torch-sla geëvalueerd op 2D Poisson-benchmarks met NVIDIA H200 GPU's:

Schalbaarheid: De bibliotheek schaalt tot 169 miljoen DOF op één GPU en 400 miljoen DOF over drie GPU's.
Prestaties: Voor een probleem van 1 miljoen DOF presteert de torch-sla CG-backend (190 ms) 38 $\times$ beter dan cuDSS (7,3 s). Boven 2 miljoen DOF raken directe oplosprogramma's het geheugen op, terwijl iteratieve oplosprogramma's een bijna lineaire schaling behouden.
Geheugenefficiëntie: In een vergelijking van naieve backpropagation versus de geadjungeerde methode voor 1.000 CG-iteraties vereiste de naieve aanpak ongeveer 64 GB (OOM bij 2.000 iteraties), terwijl de geadjungeerde methode een constant gebruik van ~328 MB handhaafde, wat resulteert in een 195 $\times$ reductie in geheugengebruik.
Gradientverificatie: Gradienten werden geverifieerd tegen eindige verschillen voor niet-lineaire en eigenwaardeproblemen, met relatieve fouten onder de $10^{-5}$ .
End-to-End Toepassing: Een taak voor het leren van inverse coëfficiënten (het terugvinden van ruimtelijk variërende geleidbaarheid in een Poisson-vergelijking) convergeerde succesvol naar een relatieve L2-fout van 0,23% in minder dan één minuut met behulp van een standaard PyTorch Adam-lus.

Betekenis

Het artikel beweert dat torch-sla een langdurige kloof in het PyTorch-ecosysteem dicht, en capaciteiten biedt die analoog zijn aan JAX's jax.lax.custom_linear_solve en jax.scipy.sparse.linalg, maar dan binnen een geünificeerde PyTorch-bibliotheek. Het stelt onderzoekers in staat om:

Modellen te trainen die betrekking hebben op lineaire oplossingen voor dunne matrices zonder gradients te laten vallen of code te porteren naar JAX.
Workloads voor wetenschappelijk machine learning te schalen naar honderden miljoenen vrijheidsgraden op gedistribueerde GPU-clusters.
Oplosprogramma's voor dunne matrices te behandelen als differentieerbare lagen, wat het onderzoek naar geleerde preconditioners en eind-tot-eind optimalisatie van fysische simulaties faciliteert.

De auteurs merken beperkingen op, waaronder de huidige afhankelijkheid van Jacobi-preconditioning voor de PyTorch-native backend (wat de convergentiesnelheid voor zeer grote systemen beperkt) en een focus op 2D Poisson-benchmarks, waarbij 3D PDE's en indefiniete systemen zijn geïdentificeerd als toekomstig werk.