Code d'erreur : 0100

Catégorie : Exécution : échec de l'allocation de mémoire tampon

Cette erreur indique que l'allocateur de mémoire du runtime XLA:TPU n'a pas trouvé de bloc de mémoire approprié sur la HBM de l'accélérateur pour l'allocation demandée.

Exemple de message d'erreur :

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

Backends XLA : TPU

Présentation

Cette erreur est générée sur

  • Échecs d'allocation de mémoire tampon initiée par l'utilisateur via jax.device_put ou
  • échecs des allocations de sortie des programmes planifiés par les utilisateurs.

Ces échecs sont généralement dus à l'une des raisons suivantes :

  • Mémoire insuffisante (OOM, Out Of Memory) : l'utilisateur tente d'allouer un bloc de mémoire plus grand que la quantité totale de mémoire libre disponible sur la HBM du TPU.
  • Fragmentation de la mémoire : l'allocation échoue, car aucun bloc libre contigu dans l'espace mémoire n'est suffisamment grand pour répondre à la taille demandée. La quantité totale de mémoire disponible est suffisante pour l'allocation, mais elle est répartie dans l'espace mémoire en petits blocs non contigus.

Le runtime TPU dispose de plusieurs mécanismes pour réessayer les échecs d'allocation, y compris :

  • S'il existe des désallocations en file d'attente, le runtime réessaie les allocations ayant échoué.
  • En cas d'OOM causée par une fragmentation, le runtime peut automatiquement déclencher une défragmentation et une nouvelle tentative.
  • L'exécution TPU donne la priorité aux allocations de mémoire tampon plutôt qu'au maintien des programmes chargés. Si une allocation de tampon échoue en raison d'une mémoire HBM insuffisante, le système expulsera les programmes TPU chargés jusqu'à ce que suffisamment de mémoire soit disponible pour le tampon.

Par conséquent, une erreur rencontrée après les mesures d'atténuation ci-dessus nécessite généralement une action de l'utilisateur.

Débogage

  • Réduisez l'empreinte mémoire de votre modèle :
    • Diminuer la taille du lot : cela réduit directement l'utilisation de la mémoire.
    • Partitionnement des paramètres : pour les très grands modèles, utilisez des techniques telles que le parallélisme de modèle ou le partitionnement pour distribuer les paramètres sur la mémoire HBM de plusieurs cœurs ou hôtes TPU.
    • Raccourcir la longueur de la séquence/du contexte : pour les modèles qui fonctionnent sur des séquences (comme les modèles de langage), la réduction de la longueur de la séquence d'entrée peut diminuer considérablement l'empreinte mémoire.
    • Don de tampon : utilisez les fonctionnalités du framework (par exemple, jax.jit(..., donate_argnums=...)) pour indiquer à XLA que certains tampons d'entrée peuvent être écrasés et réutilisés pour les sorties.
    • Optimisez la stratégie de point de contrôle : au lieu d'enregistrer l'état complet du modèle en une seule fois, envisagez d'enregistrer uniquement les pondérations du modèle ou d'utiliser une stratégie de point de contrôle fragmentée.
  • Disposition et remplissage de la mémoire d'adresse :
    • La mémoire du TPU est allouée par blocs, et le remplissage peut augmenter la taille réelle des Tensors.
  • 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 (OOM), ce qui peut être très utile pour le débogage.