Présentation de PyTorch Monarch
PyTorch Monarch est une nouvelle approche de programmation distribuée qui permet aux développeurs d'orchestrer des grappes de GPU entières à partir d'un seul programme Python. Avec son modèle d'exécution basé sur des acteurs, une abstraction de mesh de processus et un modèle d'exécution asynchrone, Monarch simplifie la formation distribuée à grande échelle et permet des flux de travail complexes qui combinent la formation, l'évaluation et l'apprentissage par renforcement dans un seul script unifié.
Architecture de PyTorch Monarch
L'architecture de Monarch fonctionne à plusieurs niveaux distincts : l'API Python, le runtime Monarch, le runtime Rust (Tokio) et l'infrastructure. Le runtime Monarch gère les acteurs et les meshes, les arbres de supervision et le partage de tenseurs. Le runtime Rust (Tokio) assure des performances élevées et une sécurité de la mémoire.
Figure 1 : Architecture de PyTorch Monarch
Portage de Monarch sur ROCm
Pour apporter Monarch aux GPU AMD, un effort d'ingénierie important a été nécessaire pour porter le runtime GPU et la pile de communication distribuée sur ROCm. Trois principaux chemins de portage ont été mis en œuvre avec succès : les communications collectives, la gestion de la mémoire GPU et l'intégration de RDMA.
Figure 3 : Portage de Monarch de CUDA à ROCm via hipify_torch et détection automatique
Étude de cas : formation tolérante aux fautes à grande échelle
Pour démontrer la puissance de Monarch sur les GPU AMD, nous l'avons intégré à TorchTitan et TorchFT pour construire une architecture de formation distribuée résiliente et sans point de contrôle. L'architecture se compose de trois couches : Monarch, TorchFT et TorchTitan. Monarch agit comme un orchestrateur, gérant l'orchestration des processus et des grappes. TorchFT gère la tolérance aux fautes au niveau des étapes, tandis que TorchTitan sert de moteur de formation, exécutant les étapes Forward, Backward et Optimizer.
Figure 4 : La pile de formation résiliente sur les GPU AMD intégrant Monarch, TorchFT et TorchTitan