Code d'erreur : 0101

Catégorie : Exécution : échec de l'allocation du programme

Cette erreur indique que le runtime XLA sur un appareil TPU n'a pas réussi à charger un exécutable de programme XLA compilé dans la HBM du TPU.

Exemple de message d'erreur :

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Backends XLA : TPU

Présentation

Cette erreur est généralement due à l'une des raisons suivantes :

  • La taille du programme dépasse la mémoire HBM disponible : le programme XLA compilé, y compris ses instructions, ses données statiques et ses constantes intégrées, est plus volumineux que la quantité totale de mémoire HBM libre actuellement disponible sur le ou les cœurs de TPU spécifiques sur lesquels le programme est chargé.
  • Fragmentation de la mémoire HBM : bien que la mémoire HBM libre totale sur l'appareil puisse être suffisante dans l'ensemble, elle n'est pas disponible dans un seul bloc contigu suffisamment grand pour contenir l'intégralité du programme.

Il est important de comprendre comment le runtime TPU hiérarchise la mémoire. Les allocations de mémoire tampon sont privilégiées par rapport aux programmes chargés. Si l'allocation d'un tampon échoue, le runtime évincera les programmes déjà chargés de la HBM pour libérer de l'espace. Cela peut entraîner l'échec d'un programme qui s'était chargé correctement auparavant, avec une erreur OOM, car la HBM est désormais occupée par davantage de tampons de données.

Débogage

  • Réduisez l'empreinte mémoire du tampon : libérer la mémoire utilisée par les tampons de données laissera plus de place au programme lui-même :
    • Diminuer la taille du lot : c'est l'un des moyens les plus efficaces de réduire la quantité de mémoire utilisée pour les activations.
    • Segmentation des paramètres : pour les très grands modèles, utilisez des techniques de parallélisme ou de segmentation de modèle (comme FSDP ou Megascale) pour distribuer les paramètres et les calculs du modèle sur plusieurs cœurs ou hôtes de TPU.
    • Raccourcir la longueur de la séquence/du contexte : pour les modèles traitant des données séquentielles (par exemple, (modèles de traitement du langage naturel), la réduction de la longueur de la séquence peut diminuer considérablement l'utilisation de la mémoire.
    • Don de tampon : utilisez les fonctionnalités du framework (par exemple, jax.jit(..., donate_argnums=...)) pour permettre à XLA de réutiliser la mémoire des tampons d'entrée pour stocker la sortie, ce qui réduit l'utilisation maximale de la mémoire.
  • Réduisez les besoins en mémoire du programme pour les temporaires :
    • Réduisez l'utilisation de la mémoire des programmes pour les temporaires en utilisant l'indicateur tpu_shared_memory_percent. Notez que cela peut avoir un impact négatif sur les performances.
  • Optimiser la stratégie d'exécution/réduire la charge de diffusion :
    • Gérer le chargement du programme : si vous compilez plusieurs fonctions JIT, sachez que chaque fonction peut entraîner le chargement d'un programme. Essayez de structurer votre charge de travail de manière à minimiser le nombre de programmes chargés simultanément.
  • Assurez-vous qu'il n'y a pas de fuite de mémoire :
    • Assurez-vous que les références aux objets jax.Array ne sont pas conservées plus longtemps que prévu. La conservation des objets jax.Array peut empêcher la désallocation automatique, même une fois la compilation du programme terminée.

Outils

  • Activez le flag tpu_log_allocations_on_oom pour lequel l'allocateur générera un rapport détaillé de toutes les allocations actuelles en cas d'erreur de mémoire insuffisante. Cela peut être très utile pour le débogage.
  • Profilez votre programme : utilisez le profileur de mémoire JAX ou le profileur TensorFlow pour obtenir une vue détaillée de l'utilisation de la mémoire de votre programme au fil du temps. Cela peut vous aider à identifier les pics inattendus de consommation de mémoire.