Développé depuis plusieurs années par Google, à l’origine par des ingénieurs de Deepmind, Cloud Pathways est désormais accessible à tous les utilisateurs. Disponible sur Google Kubernetes Engine (GKE), cette solution est taillée pour la conception de systèmes IA à grande échelle basés sur plusieurs milliers – ou dizaines de milliers – d’accélérateurs. Google Cloud l’utilise d’ailleurs pour entrainer ses propres modèles Gemini les plus volumineux.
Cloud Pathways alloue dynamiquement différentes quantités de ressources de calcul pour une variété de tâches IA en fonction de leurs exigences de traitement. Il simplifie les calculs ML à grande échelle en permettant à un seul client JAX d’orchestrer des charges de travail s’étendant sur des milliers de puces TPU. Selon le fournisseur, l’implémentation de Cloud Pathways apporte des gains en termes de développement d’architectures d’IA à grande échelle : unification du cycle de développement grâce à un unique interpréteur client qui simplifie le débogage complexe sur des milliers de nœuds, maximisation du taux d’utilisation réel des puces IA, et éliminer les limites de synchronisation Single Program Multiple Data (SPMD) sur lesquelles reposent d’autres frameworks d’entraineent distribués comme PyTorch DDP ou FSDP.
Dépasser les limites de SPMD
En mode SPMD, chaque hôte et chaque accélérateur (TPU ou GPU) exécute exactement la même copie du script d’entraînement en mode synchrone (en lockstep). « La coordination entre les puces s’effectue ainsi exclusivement par des opérations dites collectives, telles que AllReduce ou AllGather, pour moyenner les gradients ou synchroniser les paramètres », explique GCP. Si ce mécanisme s’avère extrêmement efficace au sein d’un groupe de puces physiquement proches et reliées par une interconnexion propriétaire très rapide (comme l’interconnexion TPU), dès que le modèle grandit et doit s’étendre sur plusieurs groupes à travers le réseau global du datacenter, les performances s’effondrent, indique le fournisseur. Pour y remédier, Cloud Pathways apporte ainsi plusieurs améliorations comme un système de répartition asynchrone (pré-allocation de buffers des mémoires de bloc), un ordonnanceur centralisé (séquençage unique et cohérent d’opérations pour l’ensemble d’un cluster), tout en assurant aussi une meilleure tolérance aux pannes.
Cloud Pathways est disponible dans toutes les régions GCP prenant en charge les Cloud TPU v5e, TPU v5p, Trillium (TPU v6e) et Ironwood (TPU 7x) sur GKE, avec des images de conteneurs versionnées (us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways) prenant en charge JAX 0.10.1+. Cloud Pathways est également disponible en préversion sur les clusters Slurm TPU dans Cluster Director.
