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

Nouvelles ME : le 2 avril (UTC+8), Google a récemment publié un guide d’entraînement pour développeurs destiné à la septième génération d’Ironwood TPU. Ce guide vise à aider les développeurs à tirer pleinement parti des performances au niveau système d’Ironwood TPU, afin d’entraîner et de déployer efficacement des modèles d’IA de pointe. Ironwood TPU est une infrastructure d’IA sur mesure conçue pour répondre aux besoins en puissance de calcul des modèles à plusieurs billions de paramètres. Elle s’appuie notamment sur des technologies telles que les liaisons interpuces (ICI), les commutateurs d’échange de flux optiques (OCS), le réseau de centre de données (DCN) et la mémoire à bande passante élevée (HBM) agrégée, pour construire un système complet prenant en charge jusqu’à 9 216 puces. Le texte présente en détail plusieurs stratégies d’optimisation clés pour ce matériel, notamment : tirer parti de la prise en charge native des unités de multiplication matricielle (MXU) pour l’entraînement FP8 afin d’améliorer le débit ; utiliser la bibliothèque de noyaux JAX Tokamax, spécialement optimisée pour les TPU, qui traite les tenseurs irréguliers dans les modèles à long contexte et les modèles d’experts spécialisés grâce à « attention par éclaboussures » et « multiplication matricielle groupée Megablox » ; utiliser la 4e génération de cœurs clairsemés (SparseCore) pour décharger les opérations de communication collective afin de masquer la latence ; affiner le réglage de l’allocation de la mémoire SRAM rapide à puce des TPU (VMEM) afin de réduire la stagnation mémoire ; et, en fonction de la taille du modèle, de l’architecture et de la longueur de séquence, choisir la meilleure stratégie de partitionnement (comme FSDP, TP, EP). (Source : InFoQ)