Présentation détaillée de SparseCore pour les grands modèles d'embedding (LEM)

SparseCore est un processeur en mosaïque spécialisé conçu pour accélérer les charges de travail hautes performances qui impliquent un accès et un calcul irréguliers et clairsemés de la mémoire, en particulier sur les grands ensembles de données stockés dans la mémoire à haute bande passante (HBM). Bien qu'elle excelle dans des tâches telles que les recherches d'intégration, ses capacités s'étendent à l'accélération d'une variété d'autres charges de travail dynamiques et éparses.

1. Présentation de SparseCore

Principales caractéristiques de l'architecture :

• Architecture en mosaïque : elle comprend plusieurs mosaïques de calcul (chaque mosaïque est une unité de flux de données complète avec sa propre mémoire locale et son unité de traitement), ce qui permet un traitement parallèle.
• Exécution dynamique : prend en charge de manière native le flux de contrôle et les accès à la mémoire dépendants des données, ce qui est essentiel pour les données éparses.
• Traitement vectoriel : utilise des tâches de petits vecteurs (8 ou 16 éléments, selon la version matérielle) pour un calcul efficace.
• Contrôle centralisé : un seul séquenceur SparseCore orchestre les tâches sur toutes les tuiles, ce qui garantit des opérations synchronisées.
• Prise en charge de la synthèse des données : inclut des opérations spécialisées entre les voies, utiles pour des tâches telles que le tri, le filtrage et les sommes de préfixes.
• Hiérarchie de mémoire : utilise stratégiquement la mémoire HBM pour stocker de grands ensembles de données et la mémoire scratchpad locale (SPMEM) pour organiser les données fréquemment consultées, ce qui réduit considérablement la latence HBM.

Caractéristiques en un coup d'œil :

2. Prétraitement de l'hôte SparseCore

Une préparation efficace des données est essentielle pour les performances de SparseCore. C'est là que le prétraitement de l'hôte joue un rôle essentiel. Il englobe plusieurs fonctionnalités clés :

• Transformation des données :
  • Appliquez les transformations nécessaires aux données d'entrée brutes.
  • Gérez les transformations d'ID, ce qui est particulièrement important lorsque vous empilez des caractéristiques ou des tables.
  • Convertissez les données d'entrée au format COO (Coordinate) clairsemé, comme indiqué dans la section suivante.
  • Partitionnez les données pour les distribuer efficacement entre les différents SparseCores disponibles sur le chip.
• Validation des limites :
  • Assurez-vous que les caractéristiques des données d'entrée (par exemple, le nombre d'ID) respectent les limites opérationnelles prédéfinies de SparseCore, telles que max_ids_per_partition et max_unique_ids_per_partition.
  • Si les données d'entrée dépassent ces limites, votre couche de prétraitement de l'hôte peut tenter de segmenter les données en plus petits minibatchs qui respectent les contraintes.
• Transfert de données :
  • Copiez efficacement les données traitées et validées dans la mémoire à haut débit (HBM, High Bandwidth Memory) du TPU, ce qui les prépare à l'exécution SparseCore.

Comprendre l'empilement des tableaux :

L'empilement de tables est une technique d'optimisation importante qui consiste à combiner logiquement plusieurs tables d'intégration pour améliorer l'efficacité de la recherche d'intégration. Ce processus est généralement géré automatiquement par le framework de ML sous-jacent.

• Empilement de caractéristiques : cela se produit lorsque plusieurs caractéristiques distinctes partagent la même table d'intégration sous-jacente. Un exemple courant consiste à utiliser un seul dictionnaire d'intégration pour différentes caractéristiques catégorielles, comme les codes postaux de différents contextes.
• Empilement de tables : dans ce scénario, plusieurs tables d'intégration distinctes sont empilées. Les tables qui partagent la même dimension d'embedding et la même configuration d'optimiseur sont souvent regroupées.

L'avantage principal de l'empilement de tables est la création d'une taille de lot effective plus importante pour les opérations sur ces tables empilées. Cela réduit la surcharge de calcul et peut être efficace pour masquer les latences de communication entre les puces (ICI, inter-chip communication). Pour des performances optimales, nous vous recommandons d'utiliser un nombre modéré de tables empilées (généralement entre 5 et 100).

3. Conversion en tenseurs COO

Avant que les données puissent être traitées par SparseCore, elles sont généralement converties au format de Tensor creux COO (Coordinate). Le format COO est un moyen de représenter efficacement les matrices creuses, généralement à l'aide de trois tableaux :

• row_ids : tableau contenant les index de ligne pour chaque élément non nul. Dans le contexte du traitement par lot, cela correspond souvent à la dimension du lot.
• col_ids : tableau contenant les index de colonne pour chaque élément non nul. Pour les embeddings, il s'agit souvent des valeurs de caractéristiques ou d'ID.
• values (facultatif) : tableau contenant les valeurs réelles des éléments non nuls aux coordonnées (row, col) correspondantes. Pour les calculs de limites (abordés plus loin) liés aux nombres d'ID, ces valeurs (gains) ne sont souvent pas prises en compte.

Exemple :

Prenons l'exemple d'une matrice creuse d'entrée représentant des lots d'ID :

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

Après la conversion au format COO (et potentiellement après la déduplication des ID dans le même échantillon) :

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

Cette conversion est essentielle au traitement et à la distribution des tâches par SparseCore. Les col_ids, en particulier, sont essentiels pour déterminer à quelle partition SparseCore spécifique appartient un ID, ce qui permet un partitionnement et une recherche efficaces.

4. SparsecoreConfig : l'API de haut niveau

API d'embedding spécifiques au framework :

• JAX : https://github.com/jax-ml/jax-tpu-embedding
• TensorFlow : https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding
• Keras : https://keras.io/keras_rs/api/embedding_layers/distributed_embedding

SparsecoreConfig, ou des mécanismes équivalents tels que les flags XLA, sert d'interface de haut niveau pour contrôler un large éventail de comportements SparseCore. Il est essentiel de bien comprendre ces paramètres pour ajuster efficacement les performances et assurer le bon fonctionnement de vos modèles.

• disable_table_stacking: bool = False
  • Explication : Ce signalement permet d'empêcher le framework d'empiler les tables, ce qui peut entraîner une baisse des performances en raison de l'augmentation des frais généraux et d'une capacité réduite à masquer la latence Inter-Chip Interconnect (ICI).
  • Par défaut : False (ce qui implique que l'empilement des tableaux est généralement activé par défaut lorsque le framework le prend en charge).
• max_ids_per_chip_per_sample: int = 64
  • Explication : ce paramètre établit une limite supérieure globale sur le nombre total d'ID d'embedding qu'une seule puce peut traiter à partir d'un échantillon du lot d'entrée, agrégé sur toutes les tables. Il s'agit d'un mécanisme de gestion des ressources au niveau du chip, avant que des limites plus précises par table ou par partition ne soient prises en compte. L'ajustement précis de cette valeur dépend généralement des caractéristiques spécifiques du modèle et de la capacité globale du système.
  • Par défaut : 64.
• max_ids_per_table: Optional[Dict[str, int]] = None
  • Explication : ce paramètre spécifie le nombre maximal d'ID d'embedding (qui peuvent inclure des doublons) pouvant être traités pour chaque table logique, en tenant compte de toutes ses partitions dans tous les SparseCores. Il s'agit d'une limite plus large que max_ids_per_partition. Si une table T est divisée en P partitions, cette limite s'applique à la somme des ID dirigés vers toutes les partitions P. Elle est souvent liée à max_ids_per_partition_per_sample et à la taille globale du lot.
  • Paramètre : généralement configuré à l'aide d'un fichier de limites (par exemple, à l'aide de l'indicateur xla_sparse_core_max_ids_file), où max_ids_per_partition est défini. Ce concept au niveau de la table est une méthode permettant de définir ces limites au niveau de la partition (max_ids et max_uniques).
  • Par défaut : None (la valeur peut être déduite des limites par partition ou d'autres configurations si elle n'est pas explicitement fournie).
• max_unique_ids_per_table: Optional[Dict[str, int]] = None
  • Explication : ce paramètre est analogue à max_ids_per_table, mais il spécifie le nombre maximal d'ID uniques pour chaque table logique. Il s'agit d'un paramètre essentiel pour dimensionner correctement les tampons sur l'appareil utilisés dans le traitement des identifiants uniques et les opérations vectorielles ultérieures.
  • Paramètre : également défini dans un fichier de limites ou dérivé de max_unique_ids_per_partition_per_sample.
  • Par défaut : None.
• allow_id_dropping: bool = False
  • Explication : ce indicateur booléen contrôle la suppression des ID lorsque le nombre d'ID rencontrés dans les données d'entrée (limites observées) dépasse les limites définies lors de la compilation (par exemple, max_ids_per_partition).
    • Si True : les ID qui entraîneraient le dépassement des limites sont supprimés sans notification. En règle générale, les ID d'une partition sont traités dans un ordre trié. Tout ID qui dépasserait le nombre en cours par rapport à la limite de son minibatch désigné est supprimé. Cela permet au programme de continuer à s'exécuter, mais peut avoir un impact négatif sur la précision du modèle.
    • Si False : une erreur est déclenchée et le processus s'arrête probablement si les limites observées dépassent les limites compilées. Cette approche garantit le traitement de toutes les données, mais nécessite une configuration plus conservatrice des limites.
  • Par défaut : False (une erreur se produit en cas de dépassement de capacité au lieu d'une perte de données silencieuse).

• initialize_tables_on_host: bool = True

  • Explication : Cet indicateur détermine si les tables d'intégration sont initialisées sur le processeur hôte avant d'être transférées vers la mémoire à haut débit (HBM) du TPU. La pratique standard consiste à initialiser les tables sur l'hôte. Si vous définissez cette valeur sur True, vous respectez cette convention. Si la valeur était définie sur False, cela impliquerait un mécanisme d'initialisation sur l'appareil, qui pourrait avoir des implications différentes en termes de performances ou des conditions préalables d'initialisation spécifiques.

• enable_fast_table_initialization: bool = False

  • Explication : initialise les tables directement sur le TPU. Cela peut aider à réduire les temps de démarrage des modèles.

5. Pipelining pour les performances

Le pipelining est une technique d'optimisation des performances qui permet l'exécution simultanée d'opérations sur TensorCore (TC) et SparseCore (SC). En chevauchant ces calculs, le débit global peut être considérablement amélioré.

• Mécanisme : lors d'une étape d'entraînement standard impliquant des recherches d'embedding creuses (gérées par SC) et des calculs de couches denses (gérés par TC), le pipeline permet au SC de travailler sur sa partie de l'étape i (par exemple, le pass avant ou arrière) tandis que le TC traite simultanément une autre partie de la même étape i, voire des parties d'étapes adjacentes comme i-1 ou i+1.
• Impact sur les gradients : SparseCore peut fonctionner sur des gradients "obsolètes". Par exemple, les gradients calculés lors de la phase de rétropropagation de l'étape i peuvent ne pas être entièrement mis à jour et visibles pour le SC jusqu'à l'étape i+2.
• Compromis entre performances et valeurs numériques : cette exécution chevauchante peut entraîner des accélérations substantielles, potentiellement jusqu'à doubler le temps de pas de l'appareil. Toutefois, les subtiles modifications des valeurs numériques (embedding_weights) résultant de l'utilisation de gradients obsolètes peuvent influencer le comportement de convergence du modèle ou la précision finale obtenue. L'acceptabilité de ce compromis dépend fortement du modèle et nécessite souvent une validation empirique.
• Indicateur de contrôle : le pipeline peut être contrôlé par tf_xla_disable_full_embedding_pipelining. Si vous définissez ce flag sur true, le pipeline complet (chevauchement du calcul TensorCore et SparseCore) est désactivé. Si vous le définissez sur false (ou si la sémantique du flag implique l'activation lorsqu'il est défini sur "false"), il est activé.

Flux de pipeline conceptuel :

• Sans pipeline (flux séquentiel simplifié) :

  Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

• Avec le pipeline (flux simplifié et chevauché) :

  Time ->
  Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
  Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1
  

  Remarque : Les étapes de pipelining réelles implémentées dans le matériel et le compilateur peuvent être plus complexes. Elles impliquent souvent des pré-boucles, des boucles d'exécution principales et des post-boucles pour gérer les dépendances de données et garantir l'exactitude.

6. Rôle de XLA

XLA (Accelerated Linear Algebra) est le compilateur spécifique à un domaine qui traduit les graphiques de calcul de haut niveau, généralement issus de frameworks tels que TensorFlow, en code machine hautement optimisé et adapté aux TPU. Cela inclut la génération des instructions pour les opérations destinées à SparseCore.

Fonctions clés dans le contexte SparseCore :

• Compilation des opérations éparses : XLA est responsable de la compilation des opérations de recherche d'embedding (telles que SparseDenseMatmulOp) et d'autres calculs éparses en programmes SparseCore exécutables de bas niveau.
• Intégration des limites : il utilise les limites opérationnelles configurées (par exemple, max_ids_per_partition, max_unique_ids_per_partition, souvent fournies via un fichier de limites spécifié par des indicateurs tels que xla_sparse_core_max_ids_file) pour déterminer statiquement les tailles et allouer les tampons de mémoire sur l'appareil, en particulier dans le SPMEM.
• Optimisations ciblées : XLA effectue une série d'optimisations spécialement conçues pour l'architecture SparseCore. Il peut s'agir de la planification des instructions, des transformations de la disposition de la mémoire et de la fusion des opérations pour maximiser l'efficacité.
• Contrôle à l'aide de flags : de nombreux aspects du comportement de SparseCore, des paramètres d'ajustement et des stratégies d'optimisation sont exposés et contrôlés à l'aide de flags XLA (par exemple, xla_sparse_core_estimate_max_ids pour l'estimation des limites ou xla_sc_detect_nan pour le débogage).

État Open Source :

Actuellement, l'implémentation de Sparsecore est interne et fournie à l'aide de libtpu.so.

Rapports d'erreur et diagnostics :

Les échecs de compilation liés aux configurations SparseCore ou aux contraintes de ressources se manifestent souvent sous la forme d'erreurs de compilation XLA:TPU. Ces messages d'erreur peuvent fournir des informations précieuses sur des problèmes tels que des limites trop élevées pour la SPMEM disponible ou l'utilisation de configurations non compatibles.

7. Comment les limites se traduisent-elles dans les tableaux sur SparseCore ?

Dans SparseCore, les "limites" sont des paramètres de configuration fondamentaux qui font principalement référence à deux paramètres par partition pour chaque table fragmentée (distribuée) sur les SparseCores disponibles :

• max_ids_per_partition : définit le nombre maximal d'ID totaux (y compris les doublons) qu'un SparseCore donné est censé envoyer à une partition spécifique d'une table donnée ou traiter pour celle-ci au cours d'une seule étape de calcul.
• max_unique_ids_per_partition : définit le nombre maximal d'ID uniques qu'un seul SparseCore est censé envoyer ou traiter pour un

Traduction en disposition et traitement de table physique :

• Stratégie de partitionnement des tables : les tables d'intégration sont généralement partitionnées par modulo sur tous les SparseCores du système. Cela signifie que chaque SparseCore devient responsable d'un sous-ensemble distinct du vocabulaire (lignes) de chaque table. Un ID j est généralement attribué à SparseCore_k en fonction d'une formule telle que k = j % num_total_sparse_cores.
• Définition d'une "partition" : dans ce contexte, une "partition" désigne le segment spécifique d'une table d'embedding pour lequel un seul SparseCore gère les recherches.
• Allocation de tampon SPMEM : ces limites sont utilisées par le compilateur XLA pour dimensionner et allouer statiquement des tampons dans la mémoire scratchpad (SPMEM) de l'appareil. Les tampons sont dimensionnés de sorte que toutes les données nécessaires liées aux ID d'une partition donnée (jusqu'aux limites max_ids et max_unique_ids spécifiées) puissent être chargées dans SPMEM pour le traitement. Cela est particulièrement crucial pour les calculs non élément par élément, tels que la réduction des ID en double dans une partition (par exemple, lors de la création d'une représentation CSR (Compressed Sparse Row)), où l'ensemble des données pertinentes pour les ID de cette partition doit être facilement disponible dans la mémoire rapide.

• Limites compilées et limites observées :

  • Limites observées : il s'agit du nombre réel d'ID rencontrés pour chaque partition lors de l'exécution, en fonction des données d'entrée traitées.
  • Si les limites observées dépassent les limites compilées, cela peut entraîner une perte d'ID (si allow_id_dropping est activé) ou des erreurs.

• Calcul des limites : le processus de détermination des limites appropriées implique une analyse minutieuse de la distribution des données d'entrée. Pour une table donnée (appelons-la T1, qui peut elle-même faire partie d'une table empilée plus grande T) :

  1. Le lot d'entrée (par exemple, un SparseTensor 2D de forme [BatchSize, MaxSequenceLength]) est initialement réparti entre les SparseCores disponibles. Par exemple, si un TensorCore est associé à deux SparseCores, chaque SparseCore peut recevoir un sous-lot de forme [BatchSize/2, MaxSequenceLength].
  2. Ce sous-ensemble est ensuite converti au format COO, ce qui donne row_ids et col_ids.
  3. Les ID en double dans le même échantillon (c'est-à-dire les entrées avec le même row_id et le même col_id) sont supprimés.
  4. Pour chaque col_id unique restant (dans un échantillon), le SparseCore cible responsable de cet ID est déterminé à l'aide de la règle de partitionnement mod : target_sc_id = col_id % num_total_sparse_cores.
  5. Un décompte est effectué du nombre total d'ID (ids_per_sparse_core[target_sc_id]++) et du nombre d'ID uniques (unique_ids_per_sparse_core[target_sc_id]++, après avoir vérifié l'unicité pour ce target_sc_id spécifique) destinés à chaque target_sc_id.
  6. La valeur max_ids_per_partition de la table T1 est ensuite définie sur max(ids_per_sparse_core_array).
  7. De même, le max_unique_ids_per_partition de la table T1 est défini sur max(unique_ids_per_sparse_core_array).
  8. Si la table T1 est un composant d'une table empilée, des transformations supplémentaires telles que des rotations ou des décalages peuvent être appliquées aux distributions d'ID avant de cumuler les statistiques de toutes les tables constituantes. Cela permet d'équilibrer la charge entre les puces.

Définir correctement ces limites est un exercice d'équilibre : des limites plus basses peuvent potentiellement entraîner des performances plus élevées (car moins de données doivent être traitées par étape et la pression SPMEM est réduite), mais si elles sont définies trop bas, elles peuvent entraîner un mini-batching excessif ou une suppression d'ID indésirable.

8. Comment chaque SparseCore communique

La communication SparseCore, en particulier dans le contexte du traitement d'une liste d'ID pour les recherches d'embedding, repose sur plusieurs mécanismes coordonnés :

• Sharding de module et routage implicite :
  • Les tables d'intégration sont mod-shardées sur tous les SparseCores du système.
  • Lorsque l'hôte fournit un lot de données d'entrée (qui est ensuite prétraité au format COO, y compris col_ids), la valeur col_id est utilisée pour déterminer quel SparseCore est responsable de cet ID spécifique : target_sc_id = col_id % num_total_sparse_cores.
  • Chaque SparseCore reçoit et traite uniquement le sous-ensemble d'ID qui correspondent aux partitions de vocabulaire qui lui sont attribuées. L'étape de prétraitement de l'hôte est essentielle pour préparer les données de sorte que chaque SparseCore puisse facilement identifier et traiter ses ID pertinents.
• Distribution des données par hôte :
  • La logique de prétraitement de l'hôte partitionne le lot d'entrée global et distribue les parties pertinentes de row_ids et col_ids (ainsi que les caractéristiques ou pondérations associées, le cas échéant) soit à la mémoire (HBM) directement accessible par chaque SparseCore, soit à une HBM partagée à partir de laquelle les SparseCores récupéreront les données requises.
• Traitement Intra-SparseCore :
  • Une fois qu'un SparseCore a reçu son ensemble d'ID désigné pour une partition de table donnée, il effectue des opérations telles que la déduplication de ces ID et la collecte des vecteurs d'embedding correspondants. Il s'agit principalement de calculs locaux exécutés dans les propres tuiles de SparseCore et utilisant son SPMEM local.
• Communication Inter-SparseCore (All-to-All) :
  • Après la phase de traitement initiale (comme les recherches d'intégration), un modèle de communication "all-to-all" peut être utilisé pour combiner ou redistribuer les résultats entre les SparseCores (par exemple, avant d'alimenter les activations dans une couche TensorCore qui attend une entrée correspondant à toutes les positions d'échantillon d'origine). C'est essentiel pour reconstruire l'ensemble complet des activations si le lot d'entrée d'origine a été distribué pour un traitement parallèle.
• Communication avec les TensorCores :
  • Les SparseCores communiquent avec les TensorCores pour envoyer les activations d'embedding (pendant la propagation avant) et recevoir les gradients (pendant la rétropropagation). Cette interaction est orchestrée par le programme compilé XLA et implique souvent la HBM comme tampon intermédiaire. La stratégie de pipeline (abordée précédemment) influence fortement le timing et la synchronisation de cette communication SC-TC.

En substance, la "distribution" initiale des ID aux SparseCores appropriés est en grande partie gérée par le schéma de sharding et les étapes de prétraitement de l'hôte. La communication ultérieure implique que les SparseCores fonctionnent sur leurs données locales, potentiellement suivies d'opérations de communication collective comme all-to-all si les données doivent être échangées ou réorganisées globalement entre les SparseCores avant d'être traitées par les TensorCores.

9. Gestion de la mémoire SparseCore

Chaque SparseCore gère efficacement plusieurs types de mémoire distincts pour effectuer ses calculs :

• Mémoire scratchpad (SPMEM) :
  • Nature : SRAM local relativement petit, mais très rapide, disponible exclusivement pour chaque SparseCore. Il est important de noter que SPMEM n'est pas un cache. Son utilisation est explicitement gérée et orchestrée par le compilateur XLA.
  • Objectif : SPMEM est utilisé pour "organiser les données de manière opportuniste". Cela inclut les entrées, les sorties et les résultats intermédiaires nécessaires aux calculs de couverture et de fréquence en cours. La préparation des données dans SPMEM réduit considérablement la latence élevée généralement associée à l'accès à la mémoire HBM.
  • Taille : comme indiqué dans la section "Limites", la taille des tampons SPMEM est définie de manière statique au moment de la compilation. Cette taille est basée sur des paramètres tels que max_ids_per_partition et max_unique_ids_per_partition. Cette allocation statique garantit que, pour toute opération donnée sur une partition de table (telle que la réduction du CSR), toutes les données nécessaires pour les ID de cette partition (jusqu'aux limites définies) peuvent tenir dans SPMEM.
  • Optimisations du compilateur : le compilateur XLA intègre des optimisations sophistiquées pour déterminer précisément la quantité de données et les éléments de données spécifiques qui doivent être mis en scène dans SPMEM afin de masquer efficacement la latence HBM et de maximiser les performances.
  • Contrainte d'allocation dynamique : le compilateur SparseCore ne prend actuellement pas en charge l'allocation dynamique de blocs-notes. Cela souligne l'importance cruciale du dimensionnement statique grâce à la configuration minutieuse des limites.
• Mémoire à haut débit (HBM) :
  • Nature : ressource de mémoire partagée de grande taille accessible par tous les SparseCores, TensorCores et le système hôte. Les tables d'embedding principales sont stockées dans HBM.
  • Utilisation de la pile : les opérations SparseCore nécessitent souvent un stockage temporaire dans la mémoire HBM pour les résultats intermédiaires qui ne tiennent pas dans la mémoire SPMEM limitée ou qui doivent être transmis entre les étapes plus importantes du pipeline de traitement. L'utilisation de la pile HBM lors des passes avant et arrière peut être estimée comme suit :
    • Pile HBM de transmission directe (table unique) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 octets
    • Pile HBM de transmission inverse (table unique) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 octets
  • Utilisation du tas : la HBM prend également en charge le tas, qui est géré par l'hôte. Le tas stocke des données telles que les pondérations des couches denses, les constantes utilisées par le modèle et les données d'entrée préchargées. L'utilisation du tas a tendance à augmenter avec le nombre d'étapes pour lesquelles l'hôte précharge les données (contrôlé par l'indicateur maximum_parallel_iterations). Bien qu'une prélecture plus importante puisse améliorer les performances en chevauchant les transferts hôte-appareil avec le calcul de l'appareil, elle consomme également plus de HBM.
  • Sérialisation pour l'optimisation de la HBM : l'indicateur xla_sc_num_serialized_tables_to_optimize_hbm permet de contrôler le nombre de tables dont les données sont conservées "en direct" dans la mémoire de pile HBM à un moment donné. L'augmentation de ce nombre sérialise efficacement le traitement pour un plus grand nombre de tables, ce qui peut réduire l'utilisation maximale de la pile HBM, mais peut se faire au détriment des performances en raison du parallélisme réduit.
• Mémoire vectorielle (VMEM) :
  • La mémoire VMEM est une mémoire de travail locale utilisée exclusivement par le TC (TensorCore). Bien que la VMEM ne soit pas directement gérée par SparseCore, elle fait partie intégrante de l'écosystème de mémoire avec lequel le SC interagit, principalement par le biais de TensorCore.

Stratégie globale de gestion de la mémoire :

La stratégie de gestion de la mémoire principale de SparseCore consiste à utiliser la petite et rapide SPMEM pour les données "chaudes" qui sont activement traitées par un bloc SparseCore, ce qui minimise les accès à la HBM plus lente. Les limites configurées sont le principal mécanisme permettant de s'assurer que SPMEM ne déborde pas. La HBM est utilisée pour stocker de grandes tables d'embedding et des données temporaires qui dépassent la capacité de la SPMEM ou qui doivent être partagées entre différentes unités de traitement ou étapes du pipeline. Le compilateur XLA est responsable de l'orchestration de tous les mouvements de données et de l'allocation de mémoire tampon en fonction de ces principes architecturaux et des limites configurées par l'utilisateur.

10. Goulots d'étranglement liés aux performances et à la mémoire

Pour obtenir des performances optimales avec SparseCore, vous devez bien comprendre les goulots d'étranglement potentiels et savoir comment les résoudre. Elles peuvent survenir sur l'hôte, au sein du SparseCore lui-même ou dans son interaction avec les TensorCores.

Goulots d'étranglement courants :

• Goulot d'étranglement de l'hôte :
  • Problème : le processeur hôte peut ne pas parvenir à prétraiter les données et à les transmettre au TPU assez rapidement, ce qui entraîne une sous-utilisation des SparseCores et des TensorCores. Il s'agit d'un facteur limitant fréquent des performances.
  • Atténuation : surveillez l'utilisation du processeur de l'hôte et les métriques du pipeline d'entrée. Optimisez les routines de chargement et de prétraitement des données côté hôte (consultez les conseils de conversion COO). Ajustez le flag maximum_parallel_iterations pour affiner la prélecture des données.
• Synchronisation TC/SC sous-optimale (manque de pipeline) :
  • Problème : Si le pipeline entre TensorCore et SparseCore est désactivé ou ne fonctionne pas efficacement, une unité peut passer beaucoup de temps à attendre l'autre, ce qui réduit le débit global du système.
  • Atténuation : assurez-vous que le pipeline est activé (par exemple, tf_xla_disable_full_embedding_pipelining = false ou son équivalent).
• Goulots d'étranglement induits par les limites :
  • Problème :
    • Limites trop basses : cela peut déclencher un mini-batching excessif (division des lots d'entrée en de nombreux sous-lots plus petits pour respecter les limites strictes). Bien que cela permette de maintenir l'exactitude, chaque mini-batch introduit une certaine surcharge de traitement, ce qui peut ralentir l'exécution globale. Si allow_id_dropping est défini sur "true", des limites trop basses peuvent également entraîner une perte d'ID, ce qui a un impact sur la précision du modèle.
    • Limites trop élevées (mais toujours adaptées) : bien que des limites très élevées puissent empêcher le mini-batching, elles pourraient augmenter inutilement la pression sur la SPMEM si les caractéristiques réelles des données approchent rarement ces valeurs maximales. Elles peuvent également entraîner une utilisation de la pile HBM plus importante que nécessaire.
    • Échecs de compilation : si les limites configurées nécessitent plus de piles SPMEM ou HBM que la mémoire physique disponible, la compilation échouera.
  • Solution : Assurez-vous que les limites sont correctement définies.
• Asymétrie de la distribution des données :
  • Problème : si certaines partitions SparseCore reçoivent systématiquement un nombre d'ID disproportionnellement plus élevé que d'autres (ce qui indique une mauvaise répartition des ID), ces SparseCores surchargés deviendront des goulots d'étranglement des performances.
  • Atténuation : le mélange des ID lors du processus de mini-batching peut aider à atténuer ce problème pour les tables empilées, en particulier celles avec des tables utilisateur "actives". Analysez attentivement les distributions d'ID pour définir des limites par table appropriées et équilibrées.
• Problèmes d'empilement des tableaux :
  • Problème :
    • Trop peu de tables empilées : cela peut ne pas suffire à masquer efficacement la latence ICI ni à réduire correctement les frais généraux de traitement.
    • Trop de tables empilées : cela peut entraîner la création de tables logiques très volumineuses qui deviennent difficiles à gérer ou qui peuvent dépasser les limites de ressources disponibles.
  • Mesures d'atténuation :
    • Assurez-vous que le nombre de tables à empiler est optimal. Une règle générale suggère un "sweet spot" de 5 à 100 tables pour l'empilement.
• Quantification/valeurs numériques inefficaces :
  • Problème : L'utilisation de la précision FP32 complète alors que des formats de précision inférieure tels que BF16 ou des entiers quantifiés suffiraient (et offriraient un calcul plus rapide) peut constituer un goulot d'étranglement des performances.
  • Atténuation : explorez les options de précision inférieure. Toutefois, sachez que la quantification elle-même a un certain coût et peut nécessiter un réglage minutieux des paramètres de quantification pour maintenir la précision du modèle.
• Saturation de la bande passante HBM :
  • Problème : un mouvement excessif de données vers et depuis la mémoire HBM, potentiellement causé par des largeurs de caractéristiques très petites (entraînant un fort surcoût de remplissage), des modèles d'accès à la mémoire inefficaces ou un nombre extrêmement élevé de recherches, peut saturer la bande passante HBM disponible.
  • Atténuation : la mise à l'échelle du nombre de TPU peut aider à résoudre le problème de saturation de la bande passante HBM.

Goulots d'étranglement courants liés à la mémoire :

• Dépassement de SPMEM (échec de la compilation) :
  • Problème : Si max_ids_per_partition et max_unique_ids_per_partition sont définis sur une valeur trop élevée, il est possible que le compilateur XLA ne puisse pas allouer suffisamment de SPMEM, ce qui entraîne des erreurs de compilation telles que "Fixed size allocations (...) do not fit in TileSpmem (...)". De plus, si le terme (sample_count * feature_width) / kNumTiles (où kNumTiles correspond au nombre de blocs par SC) est trop grand pour les opérandes de collecte de préparation dans la SPMEM du bloc, des erreurs telles que "Gather operand too large..." peuvent se produire.
  • Atténuation : réduisez la taille du lot ou augmentez le nombre de puces utilisées pour le traitement.
• Dépassement de capacité de la pile HBM (exécution ou compilation) :
  • Problème : Si la combinaison de feature_width, max_unique_nz_per_row et logical_replica_count entraîne des exigences de mémoire de pile HBM qui dépassent la HBM disponible, cela peut entraîner des erreurs de mémoire insuffisante (OOM) lors de l'exécution ou de la compilation.
  • Atténuation : ajustez l'indicateur xla_sc_num_serialized_tables_to_optimize_hbm pour réduire l'utilisation de la pile HBM en sérialisant le traitement des tables (cela se fait généralement au détriment des performances).
• Épuisement du tas HBM :
  • Problème : principalement dû à des pondérations de couches denses très importantes, à de nombreuses constantes stockées en mémoire ou à une prélecture d'entrée trop agressive (maximum_parallel_iterations élevé).
  • Atténuation : surveillez l'utilisation du tas à l'aide d'outils tels que XProf Memory Viewer.
• Frais généraux de remplissage :
  • Problème : Les tables d'embedding sont complétées pour être alignées sur 32 octets (équivalent à 8 valeurs flottantes) dans la dimension de la caractéristique. Par conséquent, les petites largeurs de caractéristiques (par exemple, 1 float) entraînent une surcharge de remplissage importante (par exemple, 7/8 de l'espace tampon alloué est du remplissage), ce qui entraîne un gaspillage de la HBM. La dimension du vocabulaire des tables est également complétée pour être un multiple du nombre de SparseCores dans le système. Toutefois, cet impact est généralement négligeable pour les tables dont la taille du vocabulaire est suffisamment élevée.

Facteurs généraux ayant un impact sur les performances et la mémoire :

• Topologie : nombre de puces disponibles et architecture d'interconnexion.
• Taille du lot : affecte directement le sample_count par SparseCore, ce qui influence la consommation de mémoire et la charge de calcul.
• Mise en forme des données : il est essentiel de s'assurer que la mise en page des données sur l'appareil est efficace pour optimiser les performances.

11. Analyser un profil SparseCore

L'analyse d'un profil de performances est une étape clé pour identifier les goulots d'étranglement et les opportunités d'optimisation dans vos charges de travail SparseCore.

1. Obtenir une trace :
   • Utilisez des outils de profilage, tels que XProf, pour capturer une trace d'exécution détaillée pendant l'entraînement ou l'inférence de votre modèle. Cette trace fournit un calendrier des opérations qui se produisent sur l'hôte, les TensorCores et les SparseCores.
2. Examinez le lecteur de traces (par exemple, dans XProf ou TensorBoard) :
   • Activité de l'organisateur : examinez l'activité de l'organisateur. Y a-t-il des lacunes importantes dans l'activité des TPU ? Ces écarts peuvent indiquer que l'hôte est un goulot d'étranglement et qu'il ne parvient pas à fournir les données assez rapidement. Analysez les performances de votre pipeline d'entrée.
   • Activité TensorCore (TC) et SparseCore (SC) :
     • Examinez les chronologies d'exécution pour TC et SC. Fonctionnent-elles en parallèle, ce qui indique un pipeline efficace ? Ou y a-t-il des périodes prolongées où une unité est inactive, en attente de l'autre ?
     • Identifiez les opérations qui consomment le plus de temps (opérations les plus longues) sur le SC et le TC.
     • Les résultats de trace visuels (qui affichent souvent des blocs de couleur représentant différentes opérations au fil du temps, comme TPU:0 SparseCore 1 (pid 1005)) sont très utiles pour identifier visuellement les opérations dominantes et les périodes d'inactivité.
   • Analyse du temps de traitement : observez le temps de traitement global et comprenez comment il est réparti entre le traitement de l'hôte, le calcul SC et le calcul TC.
3. Analyse de la mémoire (XProf Memory Viewer) :
   • Utilisation du segment : utilisez des outils tels que l'onglet "Memory Viewer" (Visualiseur de mémoire) de XProf pour inspecter l'utilisation du segment HBM. Cela peut vous aider à déterminer si des pondérations de modèle importantes, des constantes ou une prélecture d'entrée trop agressive consomment une quantité excessive de HBM. L'activation d'indicateurs tels que --vmodule=best_fit_allocator=1 peut fournir des journaux d'utilisation maximale du tas.
   • Utilisation de la pile (indirecte) : bien que le profilage direct de la pile HBM puisse être complexe, si vous rencontrez des erreurs de saturation de la mémoire et que l'utilisation du tas semble raisonnable, l'épuisement de la pile HBM (souvent dû à des limites ou des largeurs de caractéristiques trop importantes) est un suspect probable. Les formules fournies pour l'utilisation de la pile HBM peuvent vous aider à l'estimer.
4. Recherchez des schémas spécifiques :
   • Mini-batching : si les limites sont fréquemment dépassées, vous pouvez observer des signes de mini-batching dans la trace (par exemple, un nombre plus élevé d'opérations SC plus petites que prévu pour la taille de lot globale). Cela peut souvent être déduit des journaux ou en observant le nombre d'appels de certaines opérations.
   • Suppression d'ID : si la suppression d'ID est activée et se produit, les journaux système peuvent en fournir des indications. Cela indiquerait également clairement que les limites configurées sont trop restrictives pour les données d'entrée.
   • Temps de compilation : des temps de recompilation prolongés, en particulier si l'optimisation guidée par les commentaires (FDO) est activée et que les limites sont fréquemment ajustées, peuvent ajouter une surcharge importante au temps d'entraînement global.
5. Corrélation avec les indicateurs et la configuration :
   • Reliez le comportement observé dans le profil à vos configurations SparseCore (paramètres dans les fichiers de limites, indicateurs XLA). Par exemple, si xla_sc_num_serialized_tables_to_optimize_hbm est défini sur une valeur élevée, vous pouvez vous attendre à des performances SC plus lentes, mais à une consommation de pile HBM plus faible.
6. Processus itératif :
   • Le profilage est souvent un processus d'affinage itératif. Apportez une modification spécifique (ajustez une limite, activez ou désactivez une fonctionnalité), capturez un nouveau profil, puis comparez-le au profil précédent pour voir l'impact de votre modification.

12. Options de débogage générales

Plusieurs indicateurs peuvent être activés pour faciliter le débogage des problèmes liés à l'exécution de SparseCore. Il est important de noter que l'activation de ces vérifications entraîne souvent une perte de performances. Elles doivent donc généralement être désactivées pour les exécutions de production.

• Contrôles d'identité (hors plage) :
  • Signalement : xla_sparse_core_enable_id_bound_check = true
  • Objectif : permet de vérifier sur le système hôte si des ID d'intégration dans les données d'entrée se trouvent en dehors de la plage de vocabulaire valide définie pour une table d'intégration donnée. Cela permet de détecter les problèmes liés à des données d'entrée incorrectes ou corrompues.
• Vérificateur NaN :
  • Signalement : xla_sc_detect_nan = true
  • Objectif : permet de détecter les valeurs NaN (Not a Number) dans les données à virgule flottante traitées sur SparseCore. Si un NaN est détecté dans les entrées ou les sorties de différentes passes de compilation, ce flag entraînera une erreur. Ces erreurs fournissent généralement des informations sur l'emplacement où la valeur NaN a été rencontrée.
• Vérificateur de limites (accès à la mémoire) :
  • Signalement : xla_sc_assert_level=bounds
  • Objectif : cet indicateur permet d'utiliser un outil de type ASAN (AddressSanitizer) qui réécrit les instructions d'accès à la mémoire (telles que les chargements/stockages VMEM et les opérations DMA) pour inclure des vérifications dynamiques. Ces vérifications permettent de vérifier si l'accès à la mémoire se situe dans les limites allouées de la région de mémoire cible.
  • Comportement : si un accès à la mémoire hors limites est détecté, l'exécution échoue.
  • Attention : Il est possible que ce vérificateur produise des faux positifs, par exemple en raison de schémas d'accès complexes avec un pas qui ne sont pas entièrement compris par le vérificateur. Cette transformation est appliquée à une étape tardive du processus de compilation du backend.
• Vérificateur de tampon (corruption de mémoire) :
  • Indicateurs :
    • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
    • xla_tpu_verify_launch_id_across_cores=true
  • Objectif : ces indicateurs permettent de s'assurer que les tampons de mémoire ne sont pas corrompus ou écrasés par inadvertance par des opérations non liées. Le Buffer Sanitizer vérifie le contenu des tampons pour s'assurer qu'ils ne changent pas de manière inattendue.

13. Compatibilité avec la quantification

SparseDenseMatmulOp de SparseCore est conçu pour prendre en charge les opérations sur les tables d'intégration à l'aide de types de données à virgule flottante (FP32) et entiers sur 32 bits. Bien que l'entraînement du modèle soit généralement effectué avec une précision FP32 pour les tables d'embedding, la quantification post-entraînement (PTQ) peut être appliquée. La quantification post-entraînement permet d'utiliser des types de données de précision inférieure (comme des entiers de 8 bits) pour l'inférence, ce qui peut potentiellement améliorer les performances et réduire l'empreinte mémoire.

Quantification simulée :

Le SparseDenseMatmulOp peut être configuré pour effectuer une "quantification simulée". Dans ce mode opérationnel, les vecteurs d'embedding sont d'abord quantifiés à une précision inférieure, puis déquantifiés à une précision supérieure (par exemple, FP32) avant d'être utilisés dans les calculs ultérieurs. Cette technique permet d'entraîner des modèles en tenant compte des effets du bruit de quantification. L'entraînement avec une quantification simulée peut améliorer la précision du modèle final lorsqu'il est entièrement quantifié pour l'inférence.

Attributs de configuration pour SparseDenseMatmulOp (pour la quantification) :

• quantization_config_num_buckets = 256
  • Cet attribut spécifie le nombre de buckets ou de niveaux discrets dans lesquels un nombre à virgule flottante de 32 bits sera quantifié. Par exemple, lors de la quantification en entiers de 8 bits, on spécifie généralement 2^8 =256 buckets.
• quantization_config_low = -X.X
  • Cet attribut définit la valeur à virgule flottante minimale dans la plage de quantification. Toutes les valeurs d'entrée inférieures à ce minimum spécifié seront écrêtées à cette valeur minimale lors de la quantification.
• quantization_config_high = Y.Y
  • Cet attribut définit la valeur à virgule flottante maximale dans la plage de quantification. Toutes les valeurs d'entrée supérieures à ce maximum spécifié seront écrêtées à cette valeur maximale lors de la quantification.

Interaction entre les valeurs numériques et le pipelining :

Le comportement numérique du modèle peut changer selon que le pipeline entre TensorCore et SparseCore est activé ou non. Si le pipeline est actif, les gradients traités par SparseCore peuvent être "obsolètes" (issus d'une itération précédente). Cela peut interagir avec le processus de quantification et potentiellement affecter la dynamique d'entraînement du modèle ou la précision finale.

14. Fonctionnalités à venir et améliorations récentes

L'écosystème SparseCore est en constante évolution et amélioration.

Feuille de route :

• Mini-batching de la dimension d'échantillon :
  • Cette fonctionnalité est prévue pour compléter les capacités de minibatch existantes pour la dimension du vocabulaire.
  • Cela permettrait de partitionner davantage les entrées d'embedding selon la dimension de l'échantillon. Pour ce faire, nous allons introduire des boucles sur l'appareil qui peuvent filtrer et traiter les recherches à partir d'un sous-ensemble d'échantillons à la fois. Cette fonctionnalité pourrait être utile pour gérer un très grand nombre d'ID par échantillon ou pour améliorer l'équilibrage de charge entre les unités de traitement.
• Amélioration de la prise en charge de l'intégration avec moins de huit entiers par ligne (petites largeurs de caractéristiques) :
  • La conception actuelle utilise souvent une marge intérieure importante pour les largeurs de caractéristiques d'intégration inférieures à huit floats (ce qui correspond à 32 octets). Ce remplissage peut entraîner un gaspillage de la mémoire HBM et une sous-utilisation potentielle des ressources de calcul. Les améliorations futures visent à atténuer cette inefficacité pour les tables dont les dimensions des caractéristiques sont petites.

Améliorations récentes :

• Mise en scène des opérandes de collecte dans HBM :
  • Cette optimisation permet de réduire la pression sur la mémoire scratchpad partagée (SPMEM) en autorisant la mise en scène de certaines entrées ou sorties d'opérations de collecte dans la plus grande mémoire HBM.
• Réduction de l'utilisation de la mémoire de pile :
  • Des améliorations ont été apportées pour réduire la consommation de mémoire de la pile HBM lors des opérations SparseCore, idéalement sans impacter négativement les performances ou le débit globaux.

Ces améliorations visent à améliorer les performances, l'efficacité de la mémoire et la flexibilité opérationnelle de SparseCore pour un éventail encore plus large de charges de travail éparses.