Code d'erreur : 1000

Catégorie : Temps de compilation : HBM OOM

Cette erreur indique que le programme nécessite plus de mémoire à haut débit (HBM) que ce qui est physiquement disponible sur l'appareil TPU.

Exemples de messages d'erreur :

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

Backends XLA : TPU

Présentation

XLA effectue des vérifications pour s'assurer que la taille globale de toutes les allocations statiques nécessaires tient dans la HBM de l'appareil.

Le compilateur gère la capacité HBM fixe du TPU pour plusieurs types d'allocations :

• Entrées et sorties du programme : lots d'entraînement, états de l'optimiseur, etc.
• Temporaires TensorCore + SparseCore : mémoire dynamique requise pour les calculs intermédiaires (par exemple, les activations, les gradients, etc.).
• Binaire compilé : code machine pour TensorCore (TC) et SparseCore (SC).
• Frais généraux du système : espace réservé au runtime XLA (par exemple, les tampons Infeed sur les anciennes générations de TPU).
• Constantes : les valeurs constantes intégrées dans l'IR HLO sont allouées sur HBM.
• Internes du compilateur : allocations au niveau du programme et par HLO (par exemple, informations de routage pour les nœuds du maillage)

Cette erreur se produit lorsque le compilateur XLA ne peut pas insérer toutes les allocations ci-dessus dans la HBM de l'appareil.

Débogage

Analysez attentivement le message d'erreur et les journaux pour déterminer la catégorie d'erreur HBM OOM ci-dessous qui décrit le mieux votre erreur :

• L'utilisation de la mémoire HBM TensorCore (TC) + SparseCore (SC) dépasse la limite : Si l'erreur détaille explicitement l'utilisation, par exemple : "Utilisation de TC Hbm : X, utilisation de SC Hbm : Y". → Accéder à la section 1. Équilibrez l'utilisation de la mémoire HBM pour les TC et les SC.
• Allocations de taille inattendue : si l'erreur Mémoire insuffisante dans l'espace mémoire HBM s'affiche, consultez les journaux pour obtenir la liste des allocations les plus importantes sur HBM. Si un ou plusieurs Tensors de taille inattendue (par exemple, > 50 % de la limite HBM) sont présents, passez à la section 2. Attributions anormalement importantes.
• Les allocations agrégées dépassent la limite HBM : Si l'erreur indique "Mémoire insuffisante dans l'espace mémoire HBM", mais qu'aucun tenseur de taille anormalement grande n'est présent dans les journaux, passez à la section 3. Le total des allocations dépasse la limite HBM.

Section 1. Équilibrer l'utilisation de la mémoire HBM TC et SC

Si l'erreur détaille explicitement l'utilisation, par exemple : "Utilisation de Hbm TC : X, utilisation de Hbm SC : Y" : comparez les deux valeurs pour identifier le goulot d'étranglement.

• Utilisation élevée de SparseCore
  • Optimiser l'utilisation de la pile HBM : la consommation de mémoire de la pile HBM est proportionnelle à feature_width, max_unique_nz_per_row et logical_replica_count. Vous pouvez réduire l'utilisation maximale de la pile en ajustant l'indicateur --xla_sc_num_serialized_tables_to_optimize_hbm, qui sérialise le traitement des tables. Cela se fait au détriment d'un parallélisme réduit.
  • Vérifier la surcharge de remplissage : SparseCore aligne les tables d'intégration sur 32 octets (8 floats). Tables avec des largeurs de caractéristiques réduites (par exemple, < 8 flottants) entraînent une surcharge de remplissage importante, ce qui gaspille de la HBM.
  • Réduisez l'utilisation du tas : des valeurs élevées pour maximum_parallel_iterations augmentent la quantité de données d'entrée préchargées dans le tas HBM. Diminuer cette valeur peut libérer une quantité importante de mémoire.
  • Vérifier le partitionnement : assurez-vous que les tables d'intégration sont correctement partitionnées par modulo sur tous les chips. Consultez Comment les limites se traduisent-elles dans les tables ?
  • Pour obtenir d'autres idées, consultez SC : Goulots d'étranglement des performances et de la mémoire.
• Utilisation élevée des TensorCores
  • Passez à la section 2.
• Équilibré
  • Si aucune des valeurs n'est excessive individuellement, mais que leur somme est trop élevée, vous avez atteint la capacité du chip. Vous devez essayer de réduire l'utilisation des deux composants. Suivez les recommandations dans les trois sections.

Section 2. Allocations anormalement importantes

Si une ou plusieurs allocations de taille inattendue sont présentes dans les journaux (> 50 % de la limite HBM), il s'agit presque toujours d'un problème de capacité matérielle. Il s'agit généralement d'une erreur de configuration. Vérifiez l'étiquette XLA (le cas échéant) des grandes allocations pour obtenir des indications sur leur code source JAX.

• Supprimer les artefacts de débogage
  • L'utilisation de jax.debug.print() dans des exécutions à grande échelle peut forcer le compilateur à matérialiser le Tensor complet dans HBM pour le transférer vers le CPU, ce qui interrompt la fusion et augmente l'utilisation maximale de la mémoire. Supprimez les jax.debug.print() restants.
• Corriger les formes de maillage ou le sharding inefficaces :
  • Des formes de maillage incorrectes ou des annotations de partitionnement manquantes peuvent entraîner le compilateur à utiliser par défaut la réplication, ce qui le force à essayer d'adapter des tenseurs très volumineux sur un seul chip.
  • Vérifiez les formes des grandes allocations et assurez-vous que le sharding est correctement spécifié et propagé par XLA.

Section 3. Les allocations agrégées dépassent la limite HBM

Si le programme manque de capacité en raison de la somme globale des allocations qui dépasse la limite HBM, il est souvent utile de visualiser le profil de mémoire pour identifier les tampons spécifiques qui contribuent à l'utilisation maximale. Consultez Déboguer les erreurs OOM avec XProf pour obtenir un guide pas à pas sur l'identification des principaux contributeurs à la mémoire.

Une fois que vous avez identifié certains des principaux contributeurs, suivez les étapes ci-dessous pour optimiser l'espace mémoire.

A. Vérifier le remplissage et l'alignement du Tensor

Les formes de tenseurs inefficaces sont une cause courante et silencieuse d'erreurs OOM sur les TPU. Pour obtenir des performances optimales sur les TPU, XLA complète les dimensions des Tensor, généralement à des multiples de 128 pour la dimension la plus petite et de 8 pour la deuxième plus petite. Ce remplissage affecte à la fois les tableaux d'entrée et les Tensors intermédiaires (temporaires HLO), ce qui peut augmenter considérablement l'utilisation de la mémoire, en particulier avec de petites tailles de dimension. Consultez Mises en page de tableaux.

• Vérifiez les formes des grands tampons (sur TPU v5 avec des mises en page par défaut) :
  • Si vous pointez sur un tampon dans Xprof Memory Viewer, la fiche de détails du tampon s'affiche. Elle contient des informations sur le tampon, y compris des informations sur la marge intérieure.
  • Exemple : Une forme de (129, 1024) peut être complétée à (256, 1024), ce qui entraîne un gaspillage de mémoire de près de 50 %.
  • Correction : Une forme de (128, 1024) ne nécessite aucun remplissage et entraîne une perte de mémoire de 0 %.
• Aligner les dimensions : assurez-vous que toutes les grandes dimensions de Tensor (taille de lot, dimension d'embedding, taille cachée) sont des multiples de 128.

B. Ajuster la configuration

Vous pouvez souvent résoudre les erreurs OOM en ajustant la configuration comme suit :

• Réduisez la taille du lot : la mémoire nécessaire pour les activations et les gradients intermédiaires est directement proportionnelle à la taille du lot. Réduire la taille de lot peut souvent aider à réduire l'utilisation de la mémoire.
• Donner des tampons d'entrée : lorsque vous utilisez jax.jit, spécifiez donate_argnums pour les paramètres de votre modèle. Cela permet à XLA d'écraser la mémoire d'entrée avec la sortie.
• Activez la précision mixte (bfloat16) : utilisez bfloat16 ou la quantification (int8, etc.) pour les Tensors les plus volumineux du programme si l'architecture et les exigences de qualité du modèle le permettent.

C. Optimiser l'architecture et le partitionnement

Si les modifications de configuration sont insuffisantes, il est possible que la topologie du modèle soit trop grande pour la configuration matérielle actuelle.

• Utilisez des générations de TPU plus récentes : les TPU plus récents offrent généralement plus de HBM par puce. Passez à des générations de TPU plus récentes si elles sont disponibles.
• Exécutez le modèle sur une topologie de puces plus grande : si les pondérations du modèle sont trop importantes pour la topologie existante, vous pouvez essayer de les partitionner sur plusieurs puces.
• Implémentez des techniques de sharding avancées :
  • Explorez des approches plus avancées de parallélisme des données, des tenseurs ou des pipelines.
  • Spécifiez des indices de partitionnement pour les valeurs et les sorties intermédiaires.
• Utiliser le déchargement de l'hôte JAX : décharger les grands Tensors dans la mémoire du processeur hôte (par exemple, déchargement de l'activation et déchargement de l'état de l'optimiseur).

D. Régler les indicateurs XLA ayant un impact sur la mémoire clé :

Les indicateurs de mémoire clés peuvent être ajustés pour échanger des performances contre une utilisation de mémoire plus faible. Toutefois, elles ne doivent être utilisées qu'en dernier recours, car elles peuvent nuire aux performances.

E. Régler le pass de rematérialisation XLA / le point de contrôle manuel

Si le modèle est sur le point de tenir dans la mémoire, vous pouvez forcer le pass XLA::Rematerialization à privilégier les économies de mémoire, potentiellement au détriment de compilations plus lentes :

Vous pouvez également utiliser le décorateur jax.checkpoint avec jax.grad pour contrôler manuellement les intermédiaires qui sont enregistrés lors de la passe avant par rapport à ceux qui sont recalculés lors de la passe arrière, en échangeant des cycles de calcul contre de la mémoire HBM.

F. Utiliser des outils de profilage avancés

Déboguer les erreurs OOM avec XProf fournit un tutoriel sur l'utilisation du XProf Memory Viewer pour visualiser la vue du compilateur sur l'utilisation de la mémoire HBM.

Cet outil vous permet de visualiser l'allocation de mémoire maximale et la durée de vie des tampons, ce qui est essentiel pour comprendre exactement ce qui consomme de la mémoire HBM au moment de l'utilisation maximale. Pour la configuration générale du profilage, consultez Premiers pas avec Xprof et Profilage TensorBoard.