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

Actualités ME, 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 inclure jusqu’à 9 216 puces.
L’article décrit en détail 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’augmenter 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 matricielle 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’optimisation fine de l’allocation de la SRAM rapide (VMEM) sur le TPU pour réduire les blocages mémoire ; ainsi que le choix 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)