🤖 AI

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

Cet article présente torch-sla, une bibliothèque PyTorch open source qui unifie l'algèbre linéaire sparse différentiable en fournissant une API unique consciente de l'autograd pour divers solveurs sur plusieurs backends CPU et GPU, tout en prenant en charge les opérations par lots et l'exécution distribuée évolutive sur plusieurs GPU grâce à un cadre adjoint de graphe en O(1).

Auteurs originaux : Mingyuan Chi, Shizheng Wen

Publié 2026-05-07

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Mingyuan Chi, Shizheng Wen

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de résoudre un puzzle massif et complexe. Dans le monde de l'informatique et de l'apprentissage automatique, ce puzzle est souvent un « système linéaire creux » — une gigantesque grille de nombres où la plupart des cases sont vides (zéro), mais où les quelques cases remplies détiennent la clé de tout.

Pendant longtemps, l'outil de programmation populaire PyTorch (utilisé par de nombreux chercheurs en IA) avait un problème : il était excellent pour gérer des puzzles denses (où chaque case est remplie), mais il était terrible pour résoudre ces puzzles géants, majoritairement vides, tout en apprenant de ses erreurs. Si vous tentiez d'en résoudre un, l'ordinateur plantait soit, fonctionnait incroyablement lentement, soit oubliait comment apprendre du résultat.

Voici torch-sla, une nouvelle bibliothèque open-source qui agit comme un traducteur universel et un mécanicien ultra-efficace pour ces puzzles. Voici comment cela fonctionne, en utilisant des analogies simples :

1. Le « Dispatcher Intelligent » (Backend Unifié)

Imaginez que vous êtes un livreur avec un colis. Vous devez le livrer à une destination, mais les conditions de la route changent : parfois c'est une petite rue de ville, parfois une autoroute immense, et parfois vous êtes dans un pays différent.

Auparavant, vous deviez choisir manuellement un véhicule différent pour chaque trajet (un vélo pour les petits travaux, un camion pour les gros, un bateau pour l'eau). Si vous choisissiez le mauvais, vous perdiez du temps ou restiez bloqué.

torch-sla agit comme un dispatcher intelligent. Vous dites simplement : « Je dois résoudre ce puzzle », et il vérifie automatiquement :

Où êtes-vous ? (Votre ordinateur utilise-t-il un processeur central ou une carte graphique puissante/GPU ?)
Quelle est la taille du puzzle ? (Est-ce un petit puzzle de 100 pièces ou un puzzle de 100 millions de pièces ?)
Quelle est sa forme ? (Est-il parfaitement symétrique ?)

Sur la base de ces indices, il choisit instantanément le meilleur « véhicule » (solveur) dans un garage de cinq options différentes (incluant des outils de SciPy, NVIDIA et PyTorch lui-même). Vous n'avez pas besoin de savoir lequel choisir ; la bibliothèque le fait pour vous.

2. Le « Raccourci Magique » (Différentiation Adjointe)

C'est le tour de magie le plus impressionnant de la bibliothèque.

L'Ancienne Méthode (Rétropropagation Naïve) :
Imaginez que vous marchez dans un labyrinthe sombre, en faisant 1 000 pas pour trouver la sortie. Pour apprendre de votre voyage, vous décidez d'écrire chaque pas que vous avez fait, la position exacte de vos pieds et la pression de l'air à chaque instant. Lorsque vous arrivez au bout, vous avez un carnet de 1 000 pages de notes. Si vous voulez revenir en arrière et voir où vous avez fait une erreur, vous devez feuilleter les 1 000 pages. Si le labyrinthe grossit, votre carnet devient énorme, et votre ordinateur manque de mémoire (RAM) juste pour essayer de contenir les notes.

La Méthode torch-sla (Solveur Adjoint) :
Au lieu d'écrire chaque pas, torch-sla utilise un « raccourci magique » basé sur un principe mathématique appelé le Théorème des Fonctions Implicites.
Pensez-y ainsi : vous parcourez le labyrinthe une fois. Lorsque vous atteignez la sortie, au lieu de retracer vos pas, vous calculez instantanément la « carte inverse » dans votre tête. Vous n'avez besoin de retenir que le départ, l'arrivée et les règles du labyrinthe.

Résultat : Peu importe que le labyrinthe ait pris 10 pas ou 10 000 pas, votre « carnet » reste de la même petite taille.
Avantage : L'article montre que cela économise 195 fois plus de mémoire que l'ancienne méthode. Cela permet de résoudre des puzzles massifs sur un seul ordinateur qui auraient auparavant fait planter le système.

3. Le « Huddle d'Équipe » (Solveurs Distribués)

Que faire si le puzzle est si grand qu'un seul ordinateur ne peut pas le contenir ? Vous avez besoin d'une équipe d'ordinateurs (comme un groupe d'amis travaillant sur une gigantesque fresque murale).

Habituellement, lorsque des amis travaillent sur des parties séparées d'une fresque, ils doivent crier en arrière et en avant pour s'assurer que les bords correspondent. En termes informatiques, cela s'appelle l'« échange de halo ».

Le Problème : Lorsque vous essayez d'apprendre de l'image finale, les anciens outils ne savaient pas comment « crier en arrière » les corrections correctement aux bonnes personnes.
La Solution torch-sla : Elle apprend aux ordinateurs comment transmettre les « notes de correction » en sens inverse. Si l'Ordinateur A envoie un morceau de données à l'Ordinateur B, torch-sla s'assure que lorsque l'apprentissage a lieu, l'Ordinateur B renvoie la correction vers l'Ordinateur A parfaitement. Cela permet à la bibliothèque de passer à l'échelle jusqu'à 400 millions d'inconnues sur trois GPU puissants.

4. Que Peut-Il Réellement Faire ?

L'article démontre que cette bibliothèque peut gérer :

Solveurs Directs : Casser le code instantanément pour des puzzles de taille moyenne.
Solveurs Itératifs : Prendre des hypothèses étape par étape pour des puzzles massifs.
Solveurs Non Linéaires et de Valeurs Propres : Gérer des puzzles où les règles changent au fur et à mesure que vous les résolvez, ou trouver des « vibrations » spécifiques (valeurs propres) dans les données.
Regroupement (Batching) : Résoudre des milliers de ces puzzles en même temps (comme résoudre 1 000 labyrinthes différents simultanément).

Le Test du Monde Réel

Pour prouver que cela fonctionne, les auteurs ont mené un test appelé « Apprentissage de Coefficient Inverse ».
Imaginez que vous voyez la forme finale d'une ombre (la solution) et que vous voulez déterminer quel objet (l'entrée) l'a projetée.

Ils ont utilisé torch-sla pour « apprendre » la forme d'un objet caché en regardant l'ombre qu'il projetait.
La bibliothèque a réussi à identifier l'objet caché avec 99,77 % de précision en moins d'une minute.
Crucialement, le chercheur n'a pas eu à écrire de code mathématique complexe pour que l'apprentissage se produise ; il a simplement appelé une commande simple : A.solve(f).

Résumé

torch-sla comble un énorme vide dans l'écosystème PyTorch. Il prend la tâche difficile et gourmande en mémoire de la résolution de puzzles mathématiques géants et creux et la rend :

Automatique : Il choisit le meilleur outil pour le travail.
Efficace : Il utilise un « raccourci magique » pour économiser d'énormes quantités de mémoire.
Évolutif : Il permet à des équipes d'ordinateurs de travailler ensemble de manière transparente.
Apprenable : Il permet aux modèles d'IA d'apprendre directement de ces problèmes mathématiques complexes sans se briser.

C'est essentiellement le maillon manquant qui permet aux utilisateurs de PyTorch de s'attaquer aux mêmes types de problèmes scientifiques et d'ingénierie qui n'étaient auparavant possibles que dans d'autres environnements de programmation plus spécialisés.

Résumé Technique : torch-sla

Énoncé du Problème

L'algèbre linéaire creuse différentiable est une exigence fondamentale pour l'apprentissage automatique scientifique, notamment pour les réseaux de neurones à graphes (GNN), les opérateurs neuronaux sur des maillages non structurés et les pipelines de simulation différentiables. Cependant, l'écosystème PyTorch manque d'une bibliothèque unifiée pour prendre en charge ces charges de travail. Le support natif de PyTorch est fragmenté : torch.sparse fournit des noyaux de bas niveau et un spsolve non différentiable et uniquement CPU (un wrapper SuperLU), tandis que torch.linalg est restreint aux matrices denses.

Les contournements existants sont insuffisants pour les scénarios d'entraînement à grande échelle ou complexes :

Rupture du Graphe : L'appel à des solveurs externes (par exemple, cuDSS, PETSc) brise le graphe autograd, contraignant les utilisateurs à abandonner les gradients ou à écrire des wrappers torch.autograd.Function personnalisés pour chaque problème.
Différentiation Inefficace : La rétropropagation naïve à travers $k$ itérations d'un solveur itératif crée un graphe de calcul de $O(k)$ nœuds. Pour un problème comportant $10^6$ degrés de liberté (DOF) et 1 000 itérations, le stockage des vecteurs intermédiaires nécessite environ 80 Go de mémoire, rendant l'entraînement à grande échelle inenvisageable.
Limites de Scalabilité : Les problèmes comportant $10^7$ – $10^9$ inconnues nécessitent une décomposition de domaine distribuée. Les primitives distribuées existantes de PyTorch ne prennent pas nativement en charge l'échange de halo transposé requis pour les adjoints d'opérateurs creux dans le passage arrière.
Fragmentation des Backends : Les solveurs directs (par exemple, cuDSS) sont rapides pour les petits problèmes ( $<10^5$ DOF) mais souffrent d'un remplissage $O(n^{1.5})$ et d'un épuisement de la mémoire pour les systèmes plus grands. Les solveurs itératifs s'adaptent mieux à l'échelle mais engendrent une surcharge de lancement de noyaux sur les petits problèmes. Aucune interface unifiée n'existe pour dispatcher entre ces backends hétérogènes en fonction de l'appareil et de la taille du problème.

Méthodologie

Les auteurs présentent torch-sla, une bibliothèque open source qui comble ce vide via trois composants principaux :

1. Abstraction Unifiée des Backends

torch-sla expose une API unique consciente de l'autograd (.solve, .matvec, .eigsh, .det) qui dispatche sur cinq backends interchangeables en fonction de l'appareil et de la taille du problème :

CPU : SciPy (SuperLU, UMFPACK, CG, BiCGStab) et Eigen (CG, BiCGStab).
GPU : NVIDIA cuDSS (LU, Cholesky, LDLT), CuPy (itératif et direct) et un solveur itératif natif de PyTorch.
Auto-Dispatch : Le système sélectionne automatiquement le backend : solveurs directs pour les problèmes de petite à moyenne taille (par exemple, $<2 \times 10^6$ DOF sur GPU) et solveurs itératifs pour les problèmes plus grands afin d'éviter le débordement de mémoire. La symétrie et la définie positivité sont détectées pour optimiser la factorisation (par exemple, passer de LU à Cholesky).

La bibliothèque prend en charge quatre types de tenseurs creux organisés par localité et regroupement : SparseTensor (matrice unique/motif partagé), SparseTensorList (motifs distincts par lot) et leurs variantes distribuées DSparseTensor et DSparseTensorList.

2. Cadre de Différentiation par Adjoint

Pour éviter la complexité mémoire et graphique de $O(k)$ de la rétropropagation naïve, torch-sla emploie une approche basée sur le théorème de la fonction implicite.

Systèmes Linéaires : Pour $Ax=b$, le passage arrière résout un seul système adjoint $A^\top \lambda = \partial L / \partial x$ . Le gradient par rapport à la matrice est calculé comme $-\lambda_i x_j$ uniquement sur le motif de creux.
Systèmes Non Linéaires : Pour $F(u, \theta)=0$ , le passage arrière nécessite une résolution linéaire adjointe ( $J^\top \lambda = \partial L / \partial u$ ) et un produit vecteur-Jacobien, indépendamment du nombre d'itérations non linéaires.
Problèmes de Valeurs Propres : Les gradients pour les valeurs propres utilisent le théorème de Hellmann–Feynman ( $\partial \lambda / \partial A_{ij} = v_i v_j$ ), ne nécessitant qu'un produit extérieur sur le motif de creux.
Complexité : Ce cadre réduit le graphe autograd à $O(1)$ nœuds et la mémoire à $O(n + \text{nnz})$ , indépendamment du nombre d'itérations du solveur ( $k$ ).

3. Couche Distribuée avec Échange de Halo

Pour l'exécution multi-GPU, torch-sla implémente une décomposition de domaine où chaque processus possède un bloc de lignes et des métadonnées de halo.

Passage Avant : Les processus échangent les valeurs de halo avant les opérations locales de Produit Matrice-Vecteur Creux (SpMV).
Passage Arrière : La bibliothèque implémente un échange de halo transposé. Alors que le passage avant disperse les valeurs possédées vers les halos voisins, le passage arrière somme les gradients de halo entrants dans les entrées de frontière du processus propriétaire.
Implémentation : Les communications avant et arrière sont enveloppées en tant que primitives torch.autograd.Function utilisant torch.distributed, assurant une différentiabilité de bout en bout à travers les résolutions distribuées.

Contributions Clés

API Unifiée : Une bibliothèque unique fournissant des interfaces conscientes de l'autograd pour les solveurs directs, itératifs, non linéaires et de valeurs propres sur cinq backends (SciPy, Eigen, cuDSS, CuPy, PyTorch-natif).
Cadre Adjoint Scalable : Une méthode de différentiation à graphe $O(1)$ qui élimine le goulot d'étranglement mémoire de la rétropropagation des solveurs itératifs, permettant l'entraînement sur des problèmes comportant des millions d'inconnues.
Autograd Distribué : La première implémentation de solveurs creux distribués dans PyTorch prenant nativement en charge la différentiation de bout en bout via l'échange de halo transposé.
Extensibilité : Une abstraction de backend permettant d'intégrer des bibliothèques externes (par exemple, PETSc, Trilinos, hypre) ou des préconditionneurs appris en implémentant une interface commune.

Résultats Expérimentaux

Les auteurs ont évalué torch-sla sur des benchmarks de Poisson 2D utilisant des GPU NVIDIA H200 :

Scalabilité : La bibliothèque s'étend jusqu'à 169M DOF sur un seul GPU et 400M DOF sur trois GPU.
Performance : Pour un problème de 1M DOF, le backend CG de torch-sla (190 ms) surpasse cuDSS (7,3 s) par un facteur de 38. Au-delà de 2M DOF, les solveurs directs manquent de mémoire, tandis que les solveurs itératifs maintiennent une mise à l'échelle quasi linéaire.
Efficacité Mémoire : Dans une comparaison de la rétropropagation naïve versus la méthode adjointe pour 1 000 itérations CG, l'approche naïve nécessitait 64 Go (OOM à 2 000 itérations), tandis que la méthode adjointe maintenait une utilisation constante d'**328 Mo**, réalisant une réduction de mémoire de 195×.
Vérification des Gradients : Les gradients ont été vérifiés par rapport aux différences finies pour les problèmes non linéaires et de valeurs propres, montrant des erreurs relatives inférieures à $10^{-5}$ .
Application de Bout en Bout : Une tâche d'apprentissage de coefficient inverse (retrouvant la conductivité variant spatialement dans une équation de Poisson) a convergé avec succès vers une erreur L2 relative de 0,23 % en moins d'une minute en utilisant une boucle Adam standard de PyTorch.

Signification

L'article affirme que torch-sla comble un écart de longue date dans l'écosystème PyTorch, fournissant des capacités analogues à jax.lax.custom_linear_solve et jax.scipy.sparse.linalg de JAX, mais au sein d'une bibliothèque PyTorch unifiée. Il permet aux chercheurs de :

Entraîner des modèles impliquant des résolutions linéaires creuses sans abandonner les gradients ni porter le code vers JAX.
Mettre à l'échelle les charges de travail d'apprentissage automatique scientifique vers des centaines de millions de degrés de liberté sur des clusters GPU distribués.
Traiter les solveurs creux comme des couches différentiables, facilitant la recherche de préconditionneurs appris et l'optimisation de bout en bout des simulations physiques.

Les auteurs notent des limitations, notamment la dépendance actuelle à la préconditionnement de Jacobi pour le backend natif PyTorch (ce qui limite la vitesse de convergence pour les systèmes très grands) et un focus sur les benchmarks Poisson 2D, les EDP 3D et les systèmes indéfinis étant identifiés comme travaux futurs.