Code d'erreur : E1001

Catégorie : Compile Time: Scoped Vmem OOM

Cette erreur indique que le programme nécessite plus de mémoire vectorielle 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 vectorielle pour différents types d'allocations :

  • Allocations limitées aux instructions : stockage temporaire dans la mémoire vectorielle lors de l'exécution d'une seule instruction HLO. Cela inclut le tampon de portée des opérandes (par exemple, pour la double mise en mémoire tampon) et les débordements de registre.
  • Allocations limitées au 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 mémoire vectorielle limitée (OOM) au moment de la compilation se produit lorsque les allocations limitées aux instructions dépassent la limite d'allocation pour cette instruction. Cette limite est contrôlée

et

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" ...
  • Mémoire vectorielle limitée (OOM) au niveau du noyau personnalisé : si l'erreur pointe vers un noyau personnalisé, passez à la section Réglage du noyau.
  • Problèmes de mémoire vectorielle non liés au noyau : si l'erreur de mémoire vectorielle limitée (OOM) 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 envoyer un rapport de bug sur XLA avec un vidage HLO.

Réglage du 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 vectorielle limitée.
  • Définir des limites de mémoire vectorielle limitée par 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/limitez explicitement les entrées/sorties du noyau à la mémoire vectorielle à l'aide de pallas.tpu.with_memory_space_constraint. Veillez à ne pas colorer trop d'entrées/sorties dans la mémoire vectorielle, car cela pourrait entraîner une erreur de mémoire vectorielle limitée (OOM) globale.
  • Ajuster la limite de mémoire vectorielle : si le réglage spécifique au noyau est difficile ou si le problème affecte de nombreux noyaux, vous pouvez ajuster la limite globale de mémoire vectorielle à l'aide de l'indicateur --xla_tpu_scoped_vmem_limit_kib.