Code d'erreur : E1000

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: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

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.

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 de TPU : mémoire dynamique requise pour les calculs intermédiaires (par exemple, activations, gradients, etc.).
  • Binaire compilé : code machine pour TensorCore (TC) et SparseCore (SC).
  • Frais généraux du système : espace réservé pour l'environnement d'exécution 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 :

Scénario 1 : Équilibrer l'utilisation de la mémoire HBM par les TC et les SC

Si l'erreur détaille explicitement l'utilisation, par exemple "TC Hbm usage: X, SC Hbm usage Y", cela signifie que l'utilisation agrégée de TensorCore (TC) et SparseCore (SC) dépasse la limite HBM. Comparez les deux valeurs pour identifier le goulot d'étranglement :

  • Utilisation élevée de SparseCore
    • Optimisez 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érifiez la surcharge de remplissage : SparseCore aligne les tables d'intégration sur 32 octets (8 floats). Les tables dont les caractéristiques sont de petite largeur (par exemple, moins de huit valeurs flottantes) entraînent une surcharge de marge intérieure 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érifiez le partitionnement : assurez-vous que les tables d'intégration sont correctement partitionnées par modulo sur toutes les puces. 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
  • É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.

Scénario 2 : Mémoire insuffisante en raison d'allocations de taille inattendue

Si le message d'erreur "Mémoire insuffisante dans l'espace mémoire HBM" s'affiche et qu'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 le libellé XLA (le cas échéant) des grandes allocations pour obtenir des indices 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 au processeur, 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 partitionnement 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.

Scénario 3 : Mémoire insuffisante en raison des allocations agrégées

Si le message d'erreur "Mémoire insuffisante dans l'espace mémoire HBM" s'affiche et qu'aucun tenseur anormalement volumineux n'est présent dans les journaux, cela signifie que le programme manque de capacité, car la somme globale des allocations dépasse la limite HBM. Dans ce cas, 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.

Scénario 3.A : 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 du modèle et les exigences de qualité le permettent. Notez que cette modification peut affecter le comportement du modèle et doit être envisagée avec soin.

Scénario 3.B : Optimiser l'architecture et le partitionnement

Si les modifications de configuration sont insuffisantes, la topologie du modèle peut être trop volumineuse 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 répartir sur un plus grand nombre de 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 : les techniques de déchargement de l'hôte permettent à l'utilisateur de décharger de grands Tensors dans la mémoire du processeur hôte (par exemple, le déchargement de l'activation et le déchargement de l'état de l'optimiseur).

Scénario 3.C : Vérifier le remplissage et l'alignement des Tensors

Les formes de tenseur 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 le visualiseur de mémoire Xprof, la fiche de détails du tampon s'affiche. Elle contient des informations sur le tampon, y compris sur le remplissage.
    • 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 du lot, dimension d'embedding, taille cachée) sont des multiples de 128. Notez que cette modification peut affecter le comportement du modèle et doit être envisagée avec soin.

Scénario 3.D : Ajuster les indicateurs XLA ayant un impact sur la mémoire des clés

Les indicateurs de mémoire clés peuvent être ajustés pour échanger des performances contre une utilisation de mémoire plus faible. Toutefois, cette stratégie ne doit être utilisée qu'en dernier recours, car elle peut nuire aux performances.

Scénario 3.E : Ajuster la passe 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 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.

Vous pouvez également forcer le passage XLA::Rematerialization pour privilégier les économies de mémoire, potentiellement au détriment de compilations plus lentes :

Option Description Impact / Compromis
--xla_tpu_max_hbm_size_mib Définit manuellement la limite de la taille HBM utilisée par la passe de rematérialisation. Force le compilateur à travailler davantage pour adapter le programme à une limite inférieure à la mémoire HBM physique réelle.
--xla_tpu_rematerialization_algo=PEAK_PRIORITY Concentrez vos efforts sur les points d'utilisation maximale de la mémoire. Peut être plus efficace que l'algorithme par défaut pour réduire la mémoire de manière agressive.
--xla_tpu_rematerialization_max_block_size_limit=32 Contrôle le nombre maximal d'instructions d'un bloc pouvant être rematérialisées à la fois. L'augmentation de cette valeur permet d'économiser de la mémoire, mais augmente considérablement le temps de compilation.
--xla_tpu_rematerialization_block_effort_factor=10.0 Définit le niveau d'effort (temps de compilation) consacré à la recherche de blocs à rematérialiser. Des valeurs plus élevées permettent une recherche plus exhaustive des économies de mémoire au détriment de temps de compilation plus longs.
--xla_tpu_pre_fusion_remat=true Permet un passage de rematérialisation supplémentaire avant le passage de fusion. Permet de réaliser davantage d'économies de mémoire, mais augmente les temps de compilation et peut avoir un impact sur la stabilité numérique.

Notez que la modification des indicateurs XLA doit être utilisée en dernier recours, car elle peut nuire aux performances.

Scénario 3.F : Utiliser des outils de profilage avancés

Déboguer les erreurs OOM avec XProf fournit un tutoriel sur l'utilisation du lecteur de mémoire XProf 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.