Optimisations des performances des TPU7x (Ironwood)

Ce guide décrit plusieurs méthodes pour optimiser les performances avec TPU7x (Ironwood) en gérant efficacement le déplacement des données entre son système de mémoire à plusieurs niveaux. Cela inclut des techniques telles que l'entraînement à faible précision, la segmentation, l'optimisation des communications, la rematérialisation de l'activation, le réglage de la mémoire vectorielle (VMEM) limitée et les noyaux d'accélérateur personnalisés.

Pour optimiser les performances avec TPU7x, vous devez d'abord vous familiariser avec l'architecture Ironwood, en particulier la hiérarchie de la mémoire et la topologie d'interconnexion. Pour en savoir plus, consultez TPU7x (Ironwood).

Entraînement à faible précision avec FP8

FP8 (virgule flottante 8 bits) est un format de données numériques efficace principalement utilisé pour accélérer l'entraînement et l'inférence des modèles. En représentant les nombres sur 8 bits (au lieu des formats standards 16 bits (FP16 ou BF16) et 32 bits (FP32)), les TPU peuvent traiter les données beaucoup plus rapidement et utiliser moins de mémoire.

TPU7x est compatible avec l'accélération matérielle intégrée pour les types de données FP8, offrant des performances théoriques maximales de 4 614 TFLOPS par puce. Cette fonctionnalité peut réduire considérablement les temps d'entraînement de bout en bout. Pour les opérations compatibles, en particulier les multiplications de matrices denses courantes pour les charges de travail d'IA, l'utilisation de FP8 peut améliorer les performances de 1,3 fois par rapport à l'entraînement BF16 standard. Par rapport à BF16, FP8 double le nombre maximal de FLOP et réduit de moitié l'espace mémoire utilisé pour les pondérations et les activations. Le format FP8 doit être un levier de réglage principal pour les charges de travail limitées par le calcul et les scénarios limités par la capacité de mémoire ou la bande passante.

L'utilisation de FP8 offre les avantages suivants en termes de performances :

  •                                                                                                                                                                     Cela évite un déchargement coûteux vers une mémoire hôte plus lente.
  • Augmentation de la taille de lot effective : en réduisant la mémoire requise pour les activations, FP8 permet d'utiliser des tailles de lot plus importantes. Cela améliore le parallélisme des données et peut entraîner un débit plus élevé et une meilleure utilisation des unités de calcul.
  • Réduction des besoins en bande passante mémoire : le déplacement de la moitié de la quantité de données pour chaque opération réduit la demande sur le chemin de données HBM vers MXU. Sur les systèmes où le déplacement des données est un goulot d'étranglement courant, cela permet de maintenir les MXU saturées de travail.

L'utilisation de FP8 avec une dégradation nulle ou limitée des performances nécessite une sélection minutieuse des techniques de quantification. Voici quelques bonnes pratiques à prendre en compte pour l'entraînement FP8 :

  • Précision du scaling : commencez par un scaling par tenseur comme référence. En cas de problèmes de qualité ou de performances, passez à la mise à l'échelle par axe. Le scaling des sous-canaux peut ne pas être nécessaire.
  • Mode de scaling : le scaling dynamique, qui calcule les facteurs de scaling au moment de l'exécution, est une bonne option par défaut pour maintenir la qualité. Bien que le scaling statique puisse offrir une amélioration significative des performances en éliminant les calculs, il nécessite un profilage minutieux pour déterminer les facteurs de scaling corrects et peut ne pas convenir à tous les cas d'utilisation, en particulier lorsque les configurations de modèle changent. À l'inverse, certains modèles et configurations robustes peuvent fixer l'échelle à la limite FP8 pour les pondérations ou les activations, ce qui vous permet de réduire la surcharge de quantification tout en maintenant la justesse et en améliorant les performances.
  • Formats FP8 (E4M3 et E5M2) : une approche courante et efficace consiste à utiliser une combinaison de formats FP8. Par exemple, utilisez E4M3 pour les pondérations et les activations dans la passe avant afin de profiter de la précision supérieure de E4M3, et utilisez E5M2 pour les gradients dans la passe arrière afin de prendre en compte la plage dynamique plus large des gradients.
  • Arrondi : l'utilisation de l'arrondi au pair le plus proche (RNE) au lieu de l'arrondi stochastique pour les gradients peut maintenir la qualité tout en offrant de meilleures performances et une meilleure reproductibilité.
  • Activer FP8 dans MaxText : MaxText est compatible avec l'entraînement FP8 via la bibliothèque de quantification QWIX. Pour activer la quantification, définissez l'option suivante dans votre configuration : use_qwix_quantization=true.

Segmentation et parallélisme

Le partitionnement est un processus qui consiste à diviser un grand modèle ou ses données d'entraînement en plus petits morceaux et à les répartir sur plusieurs puces ou cœurs TPU. Il est important de choisir la bonne stratégie de segmentation pour obtenir des performances élevées sur TPU v7x.

Une approche naïve qui maximise purement le degré de parallélisme entraîne souvent des performances médiocres en devenant liée à la communication. La meilleure approche consiste souvent à sélectionner la stratégie de segmentation la plus simple qui répond aux contraintes de mémoire, car cela minimise les frais de communication et permet d'utiliser efficacement les unités de calcul.

Avant de sélectionner une stratégie de segmentation, la première étape de tout effort d'optimisation des performances doit être une analyse de l'intensité arithmétique. Cette analyse détermine si un calcul donné est limité par le calcul, la bande passante de la mémoire ou la bande passante d'interconnexion. Il est calculé comme le rapport entre les opérations à virgule flottante et les octets de données à déplacer.

Une intensité arithmétique élevée indique une charge de travail liée au calcul. Une faible intensité arithmétique suggère une charge de travail liée à la mémoire ou à la communication, où les performances sont limitées par la vitesse à laquelle les données peuvent être déplacées depuis la mémoire HBM ou sur le réseau ICI. Cette analyse permet de déterminer la taille de lot et la stratégie de segmentation idéales. Par exemple, une charge de travail liée à la communication ne bénéficiera pas d'une stratégie de segmentation qui introduit encore plus de communication, comme le parallélisme de tenseur de haut niveau.

Cadre décisionnel pour la stratégie de segmentation

MaxText propose différentes stratégies de segmentation. Le choix optimal dépend de l'architecture du modèle, de la longueur de la séquence et de la nécessité d'équilibrer la charge de calcul par rapport aux frais de communication.

  • Parallélisme des données entièrement segmentées (FSDP) : il s'agit de la stratégie par défaut privilégiée pour le parallélisme des données. Le FSDP segmente les pondérations, les gradients et les états de l'optimiseur du modèle sur les appareils de parallélisme des données. Pendant le calcul, chaque appareil effectue une opération All-Gather pour récupérer les pondérations complètes nécessaires pour son micro-lot local. Le FSDP est très efficace tant que la taille de lot par appareil est suffisamment grande pour masquer la latence de cette communication All-Gather. Pour les modèles Mixture-of-Experts (MoE), le calcul de l'intensité arithmétique doit tenir compte de la rareté.
  • Parallélisme de tenseur (TP) : le TP segmente les tenseurs individuels sur plusieurs appareils. En règle générale, les tenseurs sont des matrices de pondération dans les blocs de perceptron multicouche (MLP) et d'attention. L'intensité arithmétique élevée du matériel (11, 5 k) impose une exigence très élevée sur les dimensions du modèle pour rendre le TP viable sur l'ICI.Essayer d'utiliser le TP peut entraîner une limitation du système par les communications.
  • Parallélisme d'experts (EP) : il s'agit de la stratégie standard et nécessaire pour entraîner des modèles MoE. EP segmente les couches "expert" sur un ensemble d'appareils, et un collectif de communication All-to-All est utilisé pour acheminer les jetons vers l'appareil expert désigné. Le PE peut être efficace si la dimension MLP du modèle est suffisamment grande pour s'approcher de la limite.
  • Parallélisme de contexte (CP) : le CP est une stratégie spécialisée essentielle pour entraîner des modèles avec des séquences très longues. Sa fonction principale est de gérer la consommation de mémoire des activations, qui augmente de manière quadratique avec la longueur de la séquence et peut dépasser la capacité de la mémoire HBM. Le CP fragmente la dimension de séquence des tenseurs d'activation, ce qui permet d'utiliser une taille de lot fractionnaire par appareil. Étant donné que le CP introduit plus de communication que le FSDP, la règle générale consiste à utiliser le degré minimal de CP nécessaire pour satisfaire les contraintes de mémoire et s'assurer que le segment de l'axe de lot reste un entier.

Le tableau suivant met en correspondance les types de charges de travail courants avec la stratégie de segmentation optimale :

Type de charge de travail Segmentation principale recommandée Segmentation secondaire Principaux goulots d'étranglement Explication
Modèle dense : séquence courte FSDP N/A Rematérialisation, FF Matmuls FSDP offre le meilleur équilibre. Avec des séquences courtes, la mémoire d'activation n'est peut-être pas un problème majeur. La clé est un lot global suffisamment grand pour masquer l'opération All-Gather de poids de FSDP. À mesure que la taille du lot augmente, la taille d'activation augmente, et une règle de rematérialisation appropriée est requise pour s'assurer que cette configuration ne manque pas de mémoire.
Modèle dense : séquence longue FSDP CP FlashAttention, mémoire d'activation La mémoire d'activation devient la principale contrainte. Le CP est requis pour activer les tailles de lot fractionnaires par appareil et éviter les problèmes de mémoire saturée (OOM). FlashAttention est la principale source de calcul et de perte de temps.
Modèle MoE – Séquence courte FSDP + EP N/A All-to-All (routage expert), rematérialisation Les modèles MoE nécessitent un EP pour segmenter les experts. La communication All-to-All pour le routage des jetons est un goulot d'étranglement majeur qui doit être chevauché. La rematérialisation est également une source importante de gaspillage.
Modèle MoE : très grande échelle FSDP + EP + PP Parallélisme des modèles (MP) Tous les goulots d'étranglement mentionnés précédemment, plus les bulles de pipeline Pour les modèles qui dépassent la mémoire d'un seul pod, le PP est nécessaire pour segmenter les couches entre les pods. Cela entraîne des frais généraux de communication DCN et de bulles de pipeline. Il s'agit d'une configuration très complexe qui nécessite un réglage minutieux.

Optimisation de la communication

Le principal mécanisme de chevauchement de la communication et du calcul sur TPU7x est appelé déchargement collectif SparseCore. L'architecture Ironwood inclut des unités SparseCore dédiées, qui agissent comme des threads de contrôle indépendants capables de gérer le déplacement des données sur la structure ICI. Cela permet d'exécuter des opérations de communication collective (comme All-Gather ou Reduce-Scatter) en parallèle des calculs principaux effectués sur les TensorCores. Il s'agit de la méthode recommandée pour les collectifs asynchrones sur TPU7x. Utilisez les options recommandées pour activer le déchargement pour les collectives les plus courantes.

Rematérialisation de l'activation

La rematérialisation de l'activation, également appelée "gradient checkpointing", est une technique fondamentale pour réduire l'empreinte HBM d'un modèle. Au lieu de stocker toutes les activations intermédiaires de la propagation avant dans la mémoire HBM pour les utiliser pendant la rétropropagation, il n'enregistre que quelques activations clés (points de contrôle) et recalcule les autres à la demande pendant la rétropropagation. Cela permet d'économiser une quantité importante de mémoire au détriment d'une augmentation des calculs (environ 25 à 30% de FLOP supplémentaires pour un bloc de transformateur standard).

La décision d'appliquer la rematérialisation de manière agressive est un paramètre de réglage essentiel qui dépend entièrement du goulot d'étranglement principal, qui varie souvent en fonction de la longueur de la séquence.

Pour les charges de travail à longue séquence (par exemple, 128k) : dans ce cas, la taille des tenseurs d'activation est le principal consommateur de mémoire HBM. La charge de travail est généralement limitée par la mémoire. Par conséquent, l'application d'une règle de rematérialisation agressive est très bénéfique. Les économies de mémoire permettent de poursuivre l'entraînement sans erreurs de mémoire insuffisante et d'utiliser des tailles de lot plus importantes. De plus, la surcharge de calcul liée au recalcul est un compromis intéressant.

Pour les charges de travail à séquence courte (8k, par exemple) : dans ce cas, la mémoire d'activation est beaucoup moins préoccupante, et la charge de travail est plus susceptible d'être limitée par les ressources de calcul. La surcharge de calcul de la rematérialisation peut être la principale source d'inefficacité.

Régler les règles de rematérialisation dans MaxText

MaxText offre un contrôle précis de la rematérialisation via un ensemble de règles prédéfinies et personnalisées, configurées à l'aide de l'option remat_policy.

Règles prédéfinies

MaxText propose les règles intégrées suivantes :

  • full : stratégie la plus agressive, qui rematérialise presque tout. Cela minimise l'utilisation de la mémoire HBM, mais maximise la surcharge de recalcul. Idéal pour les scénarios à longue séquence avec une mémoire extrêmement limitée.
  • minimal : règle la moins agressive, qui stocke la plupart des activations. Cela maximise l'utilisation de la mémoire HBM, mais minimise le recalcul. Idéal pour les charges de travail liées au calcul à séquence courte, où la mémoire n'est pas un problème.
  • Règles intermédiaires : des options telles que save_dot_with_context_except_mlp, save_qkv_proj et save_out_proj offrent différents compromis en effectuant des points de contrôle sélectifs sur les résultats des opérations de produit scalaire coûteuses tout en rematérialisant les opérations élément par élément moins coûteuses.

Règles personnalisées

Pour un niveau de contrôle plus élevé, vous pouvez définir remat_policy sur custom. Cela vous permet de spécifier le comportement des couches individuelles dans le module de décodage du modèle. Chaque couche peut se voir attribuer l'un des trois comportements suivants :

  • device : l'activation est stockée dans la mémoire HBM sur l'appareil TPU.
  • remat : l'activation est supprimée et sera rematérialisée lors de la rétropropagation.
  • offload : l'activation est déplacée de la mémoire HBM vers la mémoire de l'hôte du processeur, ce qui libère de la mémoire HBM au détriment de la latence de transfert PCIe.

Réglage de la VMEM limitée

Les performances du kernel, comme l'attention flash, dépendent des tailles de tuiles sélectionnées dans le kernel, dont la taille est limitée par la mémoire vectorielle (VMEM) disponible. Chaque TensorCore d'une puce TPU7x dispose de 64 Mio de mémoire vectorielle (VMEM). Cette capacité VMEM peut être répartie entre le champ d'application actuel (VMEM limitée) et la prélecture de pondération future. L'augmentation de la VMEM limitée permet d'augmenter la taille des tuiles dans le kernel, ce qui peut réduire les blocages de mémoire et améliorer les performances des kernels. Vous pouvez modifier la taille de la VMEM limitée en définissant xla_tpu_scoped_vmem_limit_kib (dans LIBTPU_INIT_ARGS), qui peut être utilisé pour explorer les performances du noyau ainsi que les limites de performances de bout en bout. L'optimisation de la taille de la VMEM limitée peut affecter indirectement les performances du noyau Pallas personnalisé, car l'augmentation de la VMEM limitée débloque un espace de recherche d'hyperparamètres plus important pour les tailles de tuiles dans le noyau.

Noyaux Tokamax

Tokamax, une bibliothèque de kernels JAX hautes performances dotée de nombreux kernels TPU hautement optimisés, résout plusieurs goulots d'étranglement courants spécifiques au matériel :

  • Splash attention : Splash attention est utilisé comme implémentation d'attention principale pour éliminer le goulot d'étranglement de la mémoire HBM de l'attention standard et utilise l'implémentation d'attention la plus efficace sur les TPU.
  • Multiplication matricielle groupée (GMM) Megablox : pour les charges de travail MoE, Megablox gère efficacement les multiplications matricielles groupées en calculant sur la représentation des activations irrégulières. Il mappe efficacement la dimension irrégulière, en calculant les multiplications de matrices entre les groupes de lignes irréguliers dans LHS et la matrice d'expert correspondante, ce qui évite d'avoir à remplir les lots à une taille fixe.
  • Réglage empirique avec tune-jax : la bibliothèque tune-jax dispose d'utilitaires permettant d'effectuer des recherches empiriques pour trouver des tailles de bloc optimales. Les tailles de kernel par défaut sont souvent sous-optimales. Le réglage permet de choisir des tailles de tuiles VMEM adaptées au matériel afin de maximiser l'utilisation du matériel.
  • Estimation des logits max : le noyau d'attention Tokamax Splash peut être optimisé davantage en définissant une valeur pour max_logit_const. S'il est défini, il remplace le calcul de réduction de la fonction logit maximale pendant l'opération softmax d'attention (softmax(Q * KT)), ce qui réduit certains frais de calcul et de synchronisation. Dans MaxText, il est implémenté par la configuration use_max_logits_estimate, qui peut être définie sur None (désactivé) ou sur une valeur à virgule flottante. Vérifiez que la plage de fonctions logit de votre modèle spécifique reste compatible avec l'estimation pour éviter un dépassement numérique. Nous vous recommandons d'effectuer des tests de convergence si cette valeur est définie.