Code d'erreur : 1001

Catégorie : Temps de compilation : OOM Vmem à portée limitée

Cette erreur indique que le programme nécessite plus de mémoire vectorielle à portée limitée (Vmem) que ce qui a été alloué.

Exemples de messages d'erreur :

RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...

Backends XLA : TPU

Présentation

Les TPU disposent d'une mémoire vectorielle (VMEM), qui est une mémoire de travail locale utilisée exclusivement par le TensorCore (TC). Le compilateur gère la mémoire virtuelle pour différents types d'allocations :

  • Allocations à portée d'instruction : stockage temporaire dans Vmem lors de l'exécution d'une seule instruction HLO. Cela inclut le tampon de portée des opérandes (par exemple, pour la mise en mémoire tampon double) et les déversements de registre.
  • Allocations à portée du programme : allocations qui dépassent la portée d'une seule instruction HLO. Il s'agit généralement de résultats intermédiaires et temporaires HLO qui sont des entrées et/ou des sorties d'instructions HLO.

Une erreur de type "Compile Time Scoped Vmem OOM" se produit lorsque les allocations à portée d'instruction dépassent la limite d'allocation pour cette instruction. Cette limite est contrôlée

Ces erreurs sont généralement dues à un bug interne du compilateur ou à un noyau personnalisé qui dépasse sa limite d'allocation.

Débogage

Analysez attentivement le message d'erreur pour déterminer s'il provient d'un noyau personnalisé ou d'un HLO standard. Une erreur due à un noyau personnalisé doit avoir la signature suivante :

Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
  • Erreur de mémoire insuffisante (OOM) Vmem à portée du noyau personnalisé : si l'erreur pointe vers un noyau personnalisé, passez à Réajuster le noyau.
  • Problèmes de Vmem non liés au noyau : si l'erreur OOM de Vmem se produit en raison d'une opération non liée au noyau personnalisé, il s'agit probablement d'un bug interne du compilateur. Veuillez signaler un bug sur XLA avec un vidage HLO.

Régler à nouveau le noyau

Si l'erreur provient d'un noyau personnalisé, utilisez les techniques suivantes pour réduire les besoins en mémoire du noyau :

  • Ajuster les tailles de bloc : réduisez les tailles de bloc (tailles de tuile) dans la configuration de votre noyau pour réduire l'utilisation de la mémoire virtuelle à portée limitée.
  • Définissez des limites de mémoire virtuelle à portée de noyau : demandez explicitement la quantité de mémoire requise pour ce noyau spécifique à l'aide du paramètre vmem_limit_bytes.
  • Modifier la coloration de la mémoire : colorez/contraint explicitement les entrées/sorties du noyau à VMEM à l'aide de pallas.tpu.with_memory_space_constraint. Toutefois, veillez à ne pas colorer trop d'entrées/sorties sur Vmem, car cela pourrait entraîner une erreur OOM (Out Of Memory) globale de VMEM.
  • Si le réajustement spécifique au noyau est difficile ou si le problème affecte de nombreux noyaux, vous pouvez ajuster la limite Vmem globale à l'aide de l'indicateur --xla_tpu_scoped_vmem_limit_kib.