Code d'erreur : E3000

Catégorie : Échec de l'allocation SparseCore au moment de la compilation

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 de mémoire insuffisante (OOM) pendant cette phase d'allocation.

Le message d'erreur indique généralement un ID d'espace mémoire. Vous trouverez ci-dessous les espaces mémoire courants et leurs encodages entiers :

ID de l'espace Mémoire Nom Description
0 Smem Mémoire scalaire locale. Utilisé pour les registres scalaires et le flux de contrôle.
201 TileSpmem Mémoire du bloc-notes spécifique à la vignette. SRAM local et rapide disponible pour une tuile SC spécifique.
202 Spmem Mémoire du bloc-notes partagé. Utilisé pour organiser les données (entrées, sorties, intermédiaires) de manière opportuniste afin de masquer la latence HBM.
203 HBM Mémoire à haut débit. Grande mémoire partagée utilisée pour les tables d'intégration, les tas et les piles.
204 Indicateurs de synchronisation 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 de l'espace de mémoire : 203

Cette erreur se produit si une allocation temporaire unique demandée par le programme SparseCore est trop volumineuse pour tenir dans la HBM disponible. Dans les charges de travail d'embedding standard 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 partitionnement : assurez-vous que vos tables d'embedding et vos Tensors d'entrée/sortie SC sont correctement partitionnés/fragmentés. Si un seul cœur est responsable d'un trop grand nombre de données, l'allocation peut échouer.
  • Ajuster les limites : consultez max_ids_per_partition et max_unique_ids_per_partition. Si ces valeurs sont définies trop haut, le compilateur réserve plus de mémoire que nécessaire. Consultez Comment les limites se traduisent-elles dans les tableaux ?.

Scénario 2 : Défaillances de la 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 :

  1. Isoler 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 :
    • Vider le MLIR SparseCore : définissez le flag --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.
    • Vider le LLO Mosaic : pour les noyaux personnalisés, utilisez --xla_mosaic_dump_to=/path/to/dump pour inspecter tous les programmes LLO (Low Level Optimizer) émis par Mosaic.
  2. Réduire la taille des zones de travail (utilisateurs Pallas) : si l'échec se produit dans un noyau Mosaic, vérifiez la configuration de votre scratch_shapes. Assurez-vous que vos requêtes pltpu.SMEM respectent les spécifications matérielles de votre génération de TPU.
  3. Désactiver le déchargement collectif : si l'erreur provient d'opérations collectives déchargées sur le SC, essayez de désactiver les fonctionnalités de déchargement du SC :
    • --xla_tpu_enable_sparse_core_collective_offload_all_gather=false
    • --xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
  4. Signaler 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.