Google publie le guide de formation pour les développeurs de la septième génération d'Ironwood TPU, détaillant l'optimisation des performances au niveau du système

ME News Actualités, le 2 avril (UTC+8), Google a récemment publié un guide de formation pour les développeurs concernant le TPU Ironwood de septième génération. Ce guide vise à aider les développeurs à exploiter pleinement les performances système du TPU Ironwood pour entraîner et déployer efficacement des modèles d’IA de pointe. Le TPU Ironwood est une infrastructure d’IA personnalisée conçue pour répondre aux besoins en puissance de calcul de modèles de milliards de paramètres, construite à l’aide de technologies telles que l’interconnexion entre puces (ICI), le commutateur optique (OCS), le réseau de centres de données (DCN) et la mémoire haute bande passante agrégée (HBM), supportant un système complet pouvant comporter jusqu’à 9 216 puces. L’article détaille plusieurs stratégies d’optimisation clés pour ce matériel, notamment : l’utilisation de l’unité de multiplication matricielle (MXU) native pour supporter la formation FP8 afin d’améliorer le débit ; l’adoption de la bibliothèque de noyaux JAX optimisée pour TPU, Tokamax, qui traite les tenseurs irréguliers dans les modèles à contexte long et aux experts hybrides via “attention éclaboussante” et “Megablox multiplication de matrices groupée” ; l’utilisation du quatrième génération de cœurs clairsemés (SparseCore) pour décharger les opérations de communication collective afin de masquer la latence ; l’affinement de l’allocation de la SRAM rapide du TPU (VMEM) pour réduire les blocages mémoire ; ainsi que la sélection de la meilleure stratégie de partitionnement (comme FSDP, TP, EP) en fonction de la taille du modèle, de l’architecture et de la longueur de la séquence. (Source : InFoQ)