Résumé Technique : ShardTensor – Parallélisme de Domaine pour l'Apprentissage Automatique Scientifique

1. Énoncé du Problème

L'Apprentissage Automatique Scientifique (SciML) fait face à un goulot d'étranglement critique : l'incapacité d'entraîner et d'effectuer des inférences sur des données scientifiques haute résolution en raison des contraintes de mémoire GPU. Contrairement aux Grands Modèles de Langage (LLM), où la mémoire est souvent dominée par le nombre de paramètres, les charges de travail SciML sont fréquemment limitées par les activations intermédiaires.

Le papier identifie une « malédiction computationnelle de la dimensionnalité » dans les données scientifiques. Alors que doubler la longueur d'un texte dans un LLM double approximativement le nombre de jetons d'entrée, doubler la résolution de données scientifiques à $N$ dimensions augmente la taille des données par un facteur de $2^N$. Pour des données 3D ou de dimensions supérieures (par exemple, la dynamique des fluides, l'imagerie médicale, la modélisation climatique), la mémoire requise pour stocker les activations intermédiaires lors de la phase de rétropropagation de l'entraînement croît de manière exponentielle.

Les solutions existantes échouent souvent à passer à l'échelle ou dégradent la précision :

• Précision Réduite : Offre uniquement des économies de mémoire modestes (2x) et peut souffrir de problèmes de plage dynamique dans les substituts scientifiques.
• Sous-échantillonnage Spatial : Réduit la mémoire mais sacrifie la résolution haute fidélité requise pour les percées scientifiques.
• Réduction de Modèle : Diminue la précision.
• Checkpointing des Activations / Délestage vers CPU : Efficace mais introduit une surcharge computationnelle significative (re-calcul ou transferts hôte-périphérique).
• Parallélisme Standard : Le parallélisme de données (division du lot) atteint une limite stricte à une taille de lot de 1. Le parallélisme de modèle (division des poids) est conçu pour les paramètres, et non pour les données d'entrée, et suppose des formes de tenseurs statiques qui ne s'appliquent pas au traitement de données d'entrée dynamiques (par exemple, des convolutions modifiant les formes globales).

Par conséquent, il n'existe pas de cadre généralisé pour paralléliser les données d'entrée en dessous d'une taille de lot de un par appareil, empêchant le SciML d'atteindre les résolutions natives requises par les scientifiques.

2. Méthodologie : ShardTensor

Les auteurs introduisent ShardTensor, un nouveau paradigme de parallélisme de domaine qui découple la dimensionalité spatiale des données d'entrée des contraintes matérielles. C'est une extension de l'abstraction DTensor (Distributed Tensor) de PyTorch, spécifiquement conçue pour gérer les données d'entrée plutôt que les poids du modèle.

Principes de Conception Fondamentaux

• Parallélisme de Domaine : Défini comme un parallélisme à travers les appareils pour les données d'entrée, s'étendant au-delà de la dimension du lot. Lorsque la taille du lot atteint 1, ShardTensor subdivise davantage les dimensions spatiales (par exemple, hauteur, largeur, profondeur) sur plusieurs GPU.
• Gestion des Formes Dynamiques : Contrairement à DTensor, qui suppose des formes de distribution statiques adaptées aux poids du modèle, ShardTensor prend en charge les formes dynamiques. Il introduit un quatrième composant de métadonnées : les formes de découpage (sharding shapes), qui décrivent la forme du fragment local de chaque tenseur le long de l'axe de maillage découpé. Cela permet un découpage arbitraire de données non structurées ou non uniformes.
• Dispatch Impératif : Le cadre est impératif plutôt que statique. Il utilise les mécanismes __torch_dispatch__ et __torch_function__ de PyTorch pour router les opérations dynamiquement. Cela permet au système de gérer un parallélisme de domaine couche par couche sans nécessiter de graphes de calcul précompilés ni de correctifs spécifiques pour chaque couche.
• Extensibilité : Les utilisateurs peuvent étendre le système via des interfaces fonctionnelles de haut niveau (@custom_op) ou une logique de dispatch de bas niveau (opérations aten). Cela assure la compatibilité avec les architectures de modèles en évolution et les noyaux personnalisés.

Détails de l'Implémentation

ShardTensor est construit comme une extension de DTensor de PyTorch.

• Architecture de Dispatch : Il intercepte les opérations au niveau Python. Si une opération spécifique dispose d'un gestionnaire de parallélisme de domaine, elle est exécutée ; sinon, elle revient aux sémantiques de DTensor ou au comportement standard de PyTorch.
• Communication : Le cadre gère les opérations collectives (par exemple, all-reduce, all-gather) et les échanges de « halo » (récupération de pixels adjacents depuis les appareils voisins) requis pour des opérations comme la convolution ou la normalisation afin de maintenir la cohérence numérique.
• Intégration : Il est conçu pour être composable avec les stratégies de parallélisme existantes (Parallélisme de Données et Parallélisme de Modèle), permettant des stratégies de parallélisation 2D ou de dimensions supérieures (par exemple, diviser le lot sur un axe et les dimensions spatiales sur un autre).

3. Contributions Clés

1. Cadre ShardTensor : Un cadre généralisé et open-source (disponible via NVIDIA PhysicsNeMo) permettant le parallélisme de domaine pour les données d'entrée en dessous d'une taille de lot de un.
2. Nouvelle Abstraction : L'introduction des « formes de découpage » pour gérer les distributions dynamiques de tenseurs d'entrée, le distinguant des outils de parallélisme axés sur les poids comme DTensor.
3. Modèle de Programmation Non Invasif : Une conception permettant aux utilisateurs d'envelopper les entrées de modèles avec des modifications de code minimales, évitant la nécessité de réécrire les couches ou les modèles pour le parallélisme.
4. Composabilité : Capacité démontrée à combiner le parallélisme de domaine avec le parallélisme de données et de modèle sur des axes de maillage orthogonaux.

4. Résultats et Benchmarks

Le papier valide ShardTensor via des benchmarks synthétiques et des applications industrielles réelles sur des GPU NVIDIA Blackwell (GB200) et H100.

• Ring Attention : Sur de très petites tailles de séquence, la surcharge de parallélisation domine. Cependant, sur de très grandes tailles de séquence, le système atteint une mise à l'échelle forte quasi-linéaire avec le nombre de GPU, tant pour l'entraînement que pour l'inférence.
• Vision Transformer (ViT) :
  • Mise à l'échelle Forte : À haute résolution (par exemple, $4096^2$), l'inférence est 15 fois plus rapide sur 16 GPU par rapport à un seul GPU. L'entraînement sur $2048^2$ est 5 fois plus rapide avec 8 GPU.
  • Efficacité Mémoire : L'utilisation de la mémoire évolue de manière quadratique pour les données 2D et cubique pour les données 3D, confirmant que les activations intermédiaires dominent. La mise à l'échelle forte avec ShardTensor réduit la mémoire par appareil de manière proportionnelle, permettant l'entraînement de modèles 3D avec plus d'un milliard de points d'entrée sur du matériel standard.
• Transolver (Aérodynamique Automobile) :
  • Entraîné sur le jeu de données DrivaerML avec jusqu'à 1,2 million de points par domaine.
  • Le système a maintenu une stabilité numérique et une précision compétitive (erreur L2 pour la pression et la vitesse) par rapport aux références sur un seul GPU, même lorsque la taille du domaine a été multipliée par des facteurs de 2, 4 et 8.
  • L'ensemble du pipeline de prétraitement a également été parallélisé, assurant une évolutivité de bout en bout.
• StormScope (Prévision Météorologique) :
  • Un modèle à l'échelle continentale (CONUS) à 3 km de résolution (grille de $1024 \times 1792$).
  • L'utilisation maximale de mémoire (114 Go) dépassait la capacité d'un seul GPU H100 (80 Go).
  • ShardTensor a permis un entraînement stable en divisant les activations sur 2 GPU au sein de 16 groupes de parallélisme de données. Le modèle a convergé de manière stable, suivant la trajectoire de perte d'une référence à plus basse résolution (6 km) sur un seul GPU, permettant la résolution de tempêtes convectives individuelles précédemment inaccessibles sans sacrifier l'étendue du domaine.

5. Signification et Revendications

Le papier affirme que ShardTensor répond à un goulot d'étranglement fondamental en SciML : l'incapacité d'entraîner et d'inférer sur des données à la résolution native requise par les scientifiques.

• Débloquer les Charges de Travail Limitées par la Résolution : Il transforme des calculs auparavant impossibles (en raison de plafonds de mémoire) en calculs traitables, permettant un entraînement haute fidélité dans des domaines tels que l'imagerie médicale volumétrique, la dynamique des fluides computationnelle et la modélisation climatique.
• Composabilité : Il élimine le compromis entre le volume de données, la taille du modèle et la résolution. Les scientifiques peuvent désormais évoluer le long de plusieurs axes de parallélisme simultanément.
• Accessibilité : En fournissant un modèle de programmation non invasif, il abaisse la barrière à l'entrée pour les praticiens qui ne sont pas experts en systèmes distribués, contrairement aux efforts précédents de parallélisation sur mesure et spécifiques à l'architecture.

Limites Reconnues

Les auteurs notent modestement que le modèle de dispatch impératif, couche par couche, introduit une surcharge au niveau Python, le rendant moins efficace pour les petites opérations ou les données de faible résolution. Les performances dépendent également de la bande passante d'interconnexion. De plus, toutes les opérations PyTorch ne disposent pas actuellement de chemins de dispatch de parallélisme de domaine, et l'intégration complète avec torch.compile pour les optimisations de graphes statiques est en cours.

En conclusion, ShardTensor représente une étape vers la généralisation du parallélisme de domaine élevé dans l'apprentissage automatique scientifique, aussi courant que le parallélisme de données l'est aujourd'hui, facilitant l'ère des modèles de fondation scientifiques.