Présentation de PyTorch Monarch

• PyTorch Monarch est un nouveau framework conçu pour prendre en charge l’entraînement distribué et l’inférence efficaces de modèles de grande taille
• Il étend la structure modulaire existante de PyTorch et fournit des fonctions permettant de partitionner et de gérer automatiquement d’énormes réseaux neuronaux sur plusieurs appareils et nœuds
• Grâce à une API permettant de contrôler de manière unifiée le parallélisme de modèle, le parallélisme de pipeline et le parallélisme de données, il réduit la charge liée aux configurations complexes pour les développeurs
• Monarch se montre particulièrement efficace pour des charges de travail gourmandes en mémoire comme les grands modèles de langage (LLM) et les systèmes de recommandation
• Il s’inscrit dans la volonté d’atteindre à la fois scalabilité et optimisation des performances au sein de l’écosystème PyTorch, en tant que composant clé de la prochaine génération d’infrastructures d’entraînement distribué

---

Aperçu de PyTorch Monarch

• PyTorch Monarch est présenté comme un nouveau composant de PyTorch destiné à simplifier l’entraînement distribué et l’inférence des grands modèles
  • Il est conçu pour conserver la flexibilité de PyTorch tout en permettant de répartir efficacement sur plusieurs GPU et nœuds des modèles comptant des dizaines de milliards de paramètres
  • Il fournit des fonctions automatisées de partitionnement et d’optimisation des communications, sans qu’il soit nécessaire de configurer manuellement des stratégies de parallélisation complexes
• L’objectif central de Monarch est de relever le niveau d’abstraction du parallélisme de modèle, afin de permettre aux développeurs de se concentrer sur la conception de l’architecture des modèles
  • Différentes techniques de parallélisation, comme le parallélisme de données, le parallélisme de pipeline et le parallélisme tensoriel, peuvent être pilotées via une interface unifiée
  • Cela réduit fortement la complexité du code et le surcoût de communication par rapport aux frameworks d’entraînement distribué existants

Fonctions principales et caractéristiques techniques

• Monarch place chaque couche du modèle sur l’appareil optimal grâce à un algorithme de partitionnement automatique
  • La stratégie de partitionnement est déterminée dynamiquement en tenant compte de la capacité mémoire des GPU, de la bande passante de communication et de la charge de calcul
  • Cette automatisation se révèle particulièrement efficace pour les LLM, les modèles basés sur Transformer et les systèmes de recommandation à grande échelle
• Il propose une API de parallélisation unifiée, permettant aux développeurs d’expérimenter différentes stratégies de parallélisation à partir d’une base de code unique
  • Par exemple, un même modèle peut être exécuté avec une combinaison de parallélisme de données et de parallélisme de pipeline, ou basculé vers du parallélisme tensoriel
  • Cette flexibilité facilite la recherche d’optimisations selon la taille du modèle et la configuration matérielle
• Monarch est compatible avec les fonctionnalités existantes de PyTorch, DistributedDataParallel(DDP) et Fully Sharded Data Parallel(FSDP)
  • Il est possible de migrer vers Monarch sans modifier en profondeur la base de code existante
  • Il s’intègre également avec TorchScript et TorchDynamo de PyTorch, avec prise en charge de l’optimisation de la compilation et de l’exécution

Performances et cas d’usage

• D’après les premiers résultats de benchmark, Monarch obtient par rapport à l’entraînement distribué PyTorch existant une amélioration de 20 à 30 % de l’efficacité des communications et une réduction de 15 % de l’utilisation mémoire
  • En particulier, la vitesse d’entraînement et le taux d’utilisation des GPU progressent nettement sur des modèles comptant des dizaines de milliards de paramètres
  • Ces résultats ont été validés expérimentalement sur de grands modèles de langage (par ex. la famille GPT) et des systèmes de recommandation
• Monarch fonctionne à la fois en environnements cloud et on-premise et est compatible avec les principales infrastructures cloud comme AWS, Azure et GCP
  • L’intégration avec des frameworks de niveau supérieur comme PyTorch Lightning et Hugging Face Transformers est également prise en charge

Signification dans l’écosystème PyTorch

• Monarch est considéré comme une extension d’infrastructure clé de PyTorch pour répondre à l’ère des grands modèles d’IA
  • Il fournit les bases permettant de dépasser le paradigme d’entraînement centré sur un seul GPU et de réaliser l’entraînement de modèles géants mobilisant des milliers de GPU
  • Pour les chercheurs comme pour les entreprises, il constitue une solution standardisée d’entraînement distribué permettant d’obtenir simultanément scalabilité et efficacité
• L’équipe PyTorch a publié Monarch en open source et prévoit de le faire évoluer en continu en intégrant les retours de la communauté
  • Des fonctions comme l’optimisation automatique, l’ordonnancement dynamique et le parallélisme hybride devraient être ajoutées à l’avenir
  • En tant que framework distribué de nouvelle génération pour PyTorch, il devrait contribuer à la démocratisation de l’infrastructure IA et à l’amélioration de son accessibilité

[GN⁺]

Avis sur Hacker News

• Ce projet semble viser une couche différente de Tinker
  D’après l’article de présentation de Tinker, Tinker est un service de fine-tuning managé, tandis que Monarch fournit des primitives d’infrastructure
  Je me demande donc s’il serait possible de construire au-dessus de Monarch un service comme Tinker

  • Pour l’instant, quelque chose comme TorchForge fonctionne déjà au-dessus

• On dirait que l’oxydation de PyTorch a commencé
  Monarch est séparé entre un frontend basé sur Python et un backend implémenté en Rust
  Globalement, le projet semble assez intéressant

  • D’après plusieurs sources, Monarch est un framework expérimental de PyTorch, pas un remplaçant
    On pourra donc sans doute toujours profiter des graphes cycliques basés sur std::shared_ptr et des fuites mémoire
    C’est dommage qu’ils ne l’aient pas entièrement réécrit dans un langage fonctionnel
  • Cela ressemble non pas à une version oxydée d’un projet existant, mais à un projet entièrement nouveau

• J’ai moi-même créé une extension PyTorch — mycelya-torch
  Ma version ne prend pas encore en charge la communication entre nœuds, mais la manière dont Monarch obtient ses performances m’a paru intéressante
  Monarch utilise probablement cloudpickle pour partager le code avec tous les nœuds, ce qui ne coûte qu’une phase de configuration initiale et reste efficace
  La structure qui diffuse les messages depuis un contrôleur unique était également impressionnante
  En revanche, je me demande s’il prend en charge les kernels personnalisés et à quel point le contrôle de la communication entre acteurs est fin
  Dans l’ensemble, j’aime davantage cette approche qu’une architecture à multiples contrôleurs

  • Pour utiliser des kernels personnalisés, il faudra peut-être modifier légèrement le code d’initialisation des workers distants
    Mais on peut directement intégrer (bake-in) les kernels ou le code système nécessaires

• Cela ressemble à une architecture proche de Ray

  • L’exemple de code est presque identique à Ray
    Le Actor de Monarch et les classes @ray.remote de Ray suivent le même modèle
  • Je me demande pourquoi utiliser Monarch plutôt que Ray — peut-être à cause d’une intégration plus étroite avec PyTorch ou l’abstraction des tenseurs
  • Dask propose aussi un traitement distribué similaire, mais comme il a été conçu à l’origine pour le HPC, sa prise en charge des GPU est limitée
    Voir le site officiel de Dask
  • J’ai pensé la même chose. C’est d’autant plus frappant quand on voit la récente annonce de collaboration entre PyTorch et Ray
    Billet de blog associé

• Intéressant, cela rappelle essentiellement le concept des coarrays Fortran (2008)

  • Ou peut-être Hadoop (2006)
    Sauf qu’ici c’est bien mieux, puisqu’il n’est pas nécessaire d’utiliser directement MapReduce ou Fortran

• Il y avait cette phrase : « éviter le goulot d’étranglement d’un hôte unique et utiliser l’ensemble du maillage comme cluster distribué pour l’échange de messages »
  Si quelqu’un pouvant modifier cela passe par là, j’aimerais qu’il ajoute des chiffres de référence

• Le projet a l’air bon
  J’ai quelques questions

  • Est-ce comparable à openMPI ?
  • Comment le maillage est-il construit ? Est-ce possible uniquement au sein d’un même hôte ?

• Cela pourrait devenir un projet important dans le monde des coarrays, mais on voit déjà des signes de problème
  Le moteur de tenseurs est lié à CUDA et à RDMA (ibverbs), et comme le code utilise GPUDirect RDMA,
  la dépendance à CUDA risque au final de devenir encore plus forte
  OpenUCX aurait sans doute été un meilleur choix

• Cela semble moins riche fonctionnellement que Jax
  Jax optimise la communication entre nœuds grâce à un compilateur puissant

  • Mais Monarch se concentre sur un paradigme à contrôleur unique, alors que Jax repose sur une architecture SPMD à multiples contrôleurs
    Le contrôleur unique est plus facile à comprendre, tandis que le multi-contrôleur est mieux optimisé pour certains flux de données
    Il existe aussi des tentatives intéressantes pour combiner les deux approches