Catégorie : Compile Time: SparseCore Allocation Failure
Cette erreur se produit lorsque le compilateur XLA:SparseCore ne parvient pas à allouer un bloc de mémoire contigu dans l'espace mémoire spécifié requis par le programme SparseCore actuel.
Exemples de messages d'erreur :
INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>
Backends XLA : TPU
Présentation
SparseCore (SC) est un processeur spécialisé pour les charges de travail éparses. Il s'appuie sur des hiérarchies de mémoire spécifiques pour gérer efficacement le déplacement des données. Le compilateur XLA tente de dimensionner et d'allouer des tampons de manière statique en fonction des limites matérielles et des formes définies par l'utilisateur. Cette erreur indique une condition Mémoire insuffisante lors de cette phase d'allocation.
Le message d'erreur spécifie généralement un ID d'espace mémoire. Vous trouverez ci-dessous les espaces mémoire courants et leurs encodages entiers :
| ID d'espace mémoire | Nom | Description |
|---|---|---|
| 0 | Smem | Mémoire scalaire locale. Utilisée pour les registres scalaires et le flux de contrôle. |
| 201 | TileSpmem | Mémoire de travail spécifique à la tuile. SRAM locale rapide disponible pour une tuile SC spécifique. |
| 202 | Spmem | Mémoire de travail partagée. Utilisée pour organiser de manière opportuniste les données (entrées, sorties, intermédiaires) afin de masquer la latence HBM. |
| 203 | HBM | Mémoire à haut débit. Grande mémoire partagée utilisée pour intégrer des tables, des tas et des piles. |
| 204 | Sync Flags | Primitives de synchronisation utilisées pour la coordination. |
Pour en savoir plus sur SC et sa hiérarchie de mémoire, consultez la documentation SparseCore.
Débogage
La résolution dépend de l'espace mémoire pour lequel l'allocation a échoué.
Scénario 1 : Échecs d'allocation HBM (ID d'espace mémoire : 203)
Cette erreur se produit si une seule allocation temporaire demandée par le programme SparseCore est trop volumineuse pour tenir dans la mémoire HBM disponible. Dans les charges de travail d'intégration standards et les collectifs déchargés SC, des partitions par cœur extrêmement volumineuses ou des spécifications de sharding incorrectes peuvent forcer le compilateur à demander des tampons massifs.
Actions recommandées :
- Vérifiez le sharding : assurez-vous que vos tables d'intégration et vos tenseurs d'entrée/sortie SC sont correctement partitionnés/fragmentés. Si un seul cœur est responsable d'une trop grande quantité de données, l'allocation peut échouer.
- Ajustez les limites : examinez
max_ids_per_partitionetmax_unique_ids_per_partition. Si ces valeurs sont définies de manière inutilement élevée, le compilateur réserve plus de mémoire que nécessaire. Consultez Comment les limites se traduisent en tables.
Scénario 2 : Échecs de mémoire interne (ID d'espace mémoire : 0, 201, 202, 204)
Les échecs d'allocation dans Smem, TileSpmem, Spmem ou Sync Flags se produisent généralement en raison de bugs du compilateur ou de limites dans la stratégie d'allocation, où le compilateur ne tient pas compte de toutes les exigences de mémoire.
Actions recommandées :
- Isolez l'opération XLA défaillante : pour identifier le noyau SC HLO ou Mosaic spécifique à l'origine de l'échec, générez les représentations intermédiaires du compilateur :
- Videz SparseCore MLIR : définissez l'indicateur
--xla_sc_dump_mlir_to=/path/to/dump. Cela génère le MLIR du programme SparseCore, ce qui vous permet de voir quelle taille d'allocation correspond au message d'erreur. - Videz Mosaic LLO : pour les noyaux personnalisés, utilisez
--xla_mosaic_dump_to=/path/to/dumpafin d'inspecter tous les programmes LLO (Low Level Optimizer) émis par Mosaic.
- Videz SparseCore MLIR : définissez l'indicateur
- Réduisez les tailles de travail (utilisateurs Pallas) : si l'échec se produit dans un noyau Mosaic, examinez votre configuration
scratch_shapes. Assurez-vous que vos requêtespltpu.SMEMrespectent les spécifications matérielles de votre génération de TPU spécifique. - Désactivez le déchargement collectif : si l'erreur provient d'une opération collective déchargée SC, essayez de désactiver les fonctionnalités de déchargement SC :
--xla_tpu_enable_sparse_core_collective_offload_all_gather=false--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
- Signalez un bug : si les étapes ci-dessus ne permettent pas de résoudre le problème, il s'agit probablement d'un bug du compilateur. Veuillez remplir un rapport de bug.