Guide des performances Cloud TPU

La première étape du dépannage des performances des TPU consiste à profiler votre modèle. Pour en savoir plus sur la capture d'un profil de performances, consultez Profiler votre modèle sur Cloud TPU.

Performances du modèle TPU

Cette section décrit les problèmes généraux susceptibles de réduire les performances des modèles et explique comment les résoudre.

1. Le modèle est limité par les entrées

   Les TPU effectuent des calculs très rapidement. Pour s'assurer que la TPU n'est pas inactive, il est important de s'assurer qu'un flux constant de données est chargé sur la TPU. La procédure à suivre dépend de la façon dont vous chargez et prétraitez votre ensemble de données. Par exemple, vous pouvez lire des fichiers de données en parallèle à l'aide de tf.data.TFRecordset() et du paramètre num_parallel_reads.

2. Taille de lot trop faible pour cause de segmentation (répartition des lots sur les cœurs)

   Le runtime TPU divise un lot sur les huit cœurs d'un appareil TPU (par exemple, v2-8 ou v3-8). Si vous spécifiez une taille de lot globale de 128, chaque cœur reçoit une taille de lot de 16 (128 / 8).

   Pour une utilisation optimale de la mémoire, utilisez la plus grande taille de lot adaptée à la mémoire TPU. Chaque cœur de TPU utilise des registres vectoriels bidimensionnels de 8 x 128 pour traiter les multiplications matricielles. En général, votre taille de lot doit être divisible de manière uniforme par 8 ou 128.

3. Optimisation de la gestion de la mémoire

   Vous pouvez utiliser les variables d'environnement TPU_PREMAPPED_BUFFER_SIZE pour affiner les comportements d'exécution de bas niveau.

• Description : TPU_PREMAPPED_BUFFER_SIZE définit la taille du tampon de mémoire hôte (en octets) qui est prémappé et épinglé pour être utilisé par le runtime TPU pour les transferts de données (par exemple, DMA). La valeur par défaut est de 4 294 967 296 octets. La valeur doit être un multiple de 2^12 (4 Ko = 4 * 1 024 octets = 4 096 = 2^12).

  Les exemples suivants sont des valeurs TPU_PRE_MAPPED_BUFFER_SIZE valides.

      17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
      40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
  

• Impact : augmenter cette taille peut potentiellement améliorer les performances de transfert de données entre l'hôte et l'appareil TPU, en particulier pour les charges de travail avec de grands tenseurs ou une communication fréquente entre l'hôte et l'appareil. Toutefois, cela augmente également la quantité de mémoire hôte épinglée, ce qui réduit la mémoire disponible pour d'autres processus.

  Taille de la mémoire tampon

  Si la région de tampon prémappée n'est pas assez grande pour allouer de la mémoire pendant l'exécution du programme, la charge de travail échouera et renverra une erreur RESOURCE_EXHAUSTED semblable à celle-ci :

  "L'allocation de mémoire tampon à partir de la région prémappée a échoué avec l'erreur suivante : RESOURCE_EXHAUSTED. Tentative d'allocation de allocation_size. Cela n'a pas été possible. available_size sont disponibles gratuitement."

  Si le tampon est excessivement grand, l'initialisation du TPU peut prendre beaucoup plus de temps (potentiellement plus de 15 secondes), ce qui peut donner l'impression que le TPU est bloqué.

  Pour diagnostiquer ce problème, inspectez les journaux d'exécution du TPU. Ces journaux détaillent les opérations effectuées, y compris le prémappage des tampons. Vous trouverez les journaux dans /tmp/tpu_logs/tpu_driver.INFO ou vous pouvez les imprimer directement dans la console en définissant la variable d'environnement TPU_STDERR_LOG_LEVEL=0. Ce paramètre génère un résultat semblable à celui-ci :

   I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
   I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
   I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
   I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
   ```
  
  This output will tell you how long it took to initialize the TPU and
  the size of the premapped buffer.
  

• Utilisation : si la mémoire tampon prémappée est trop petite ou trop grande, vous pouvez définir manuellement sa taille à l'aide des variables d'environnement suivantes.

  TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
  pre-mapped buffer region.
  TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
  a single buffer that can be allocated from the pre-mapped region.
  

  Par exemple, vous pouvez :

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  pour définir la taille de la mémoire tampon :

   export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
   ```
   to enable it.
  
   This export sets the size to the default.
  

• Conseil : Ajustez la valeur de TPU_PREMAPPED_BUFFER_SIZE si vous pensez que le transfert de données entre l'hôte et l'appareil est un goulot d'étranglement. Surveillez l'utilisation de la mémoire hôte et les performances du modèle pour trouver un équilibre optimal. La valeur par défaut est généralement suffisante pour la plupart des cas d'utilisation.

Optimisations du compilateur XLA

XLA est un compilateur destiné au machine learning, qui peut produire des fichiers binaires pour les TPU, les processeurs, les GPU et d'autres plates-formes. Bien que XLA fasse partie du codebase TensorFlow standard, il peut également être utilisé sur les modèles PyTorch et JAX. Les modèles pour Cloud TPU sont convertis en graphe XLA, que XLA compile ensuite en exécutable TPU. Pour en savoir plus sur XLA, consultez XLA : optimiser le compilateur pour le machine learning.

Remplissage

Pour utiliser efficacement la mémoire du TPU, structurez vos données de manière à ce qu'elles puissent être divisées en blocs de 128 x 8. Lorsque les données d'un calcul matriciel ne remplissent pas un bloc entier de 128 x 8, le compilateur XLA remplit les Tensors. Cette méthode présente toutefois deux inconvénients :

1. Les Tensors remplis sous-utilisent le cœur du TPU.
2. Le remplissage augmente la quantité de mémoire sur puce requise pour un Tensor. Cette technique peut également entraîner une erreur de mémoire insuffisante.

Bien que le remplissage soit automatiquement effectué par le compilateur XLA en cas de besoin, vous pouvez définir la proportion de remplissage à l'aide de l'outil d'affichage de la mémoire. Vous pouvez éviter le remplissage en choisissant des dimensions de Tensor adaptées aux TPU.

Dimensions du Tensor

Pour atteindre le pic de FLOPS, les dimensions de la multiplication matricielle doivent être supérieures à la taille de l'unité MXU pour la version de TPU que vous utilisez. La taille de l'unité matricielle est de 256 x 256 pour la version 6e et de 128 x 128 pour les versions antérieures. Pour en savoir plus, consultez Architecture du système Cloud TPU.

Taille de lot

Le compilateur XLA arrondit les tailles des Tensors stockés dans la mémoire HBM du TPU pour effectuer les calculs plus efficacement. Ce remplissage se produit de manière transparente au niveau du matériel et n'affecte pas les résultats. Toutefois, dans certains cas, le remplissage peut entraîner une augmentation significative de l'utilisation de la mémoire et du temps d'exécution.

L'exécution TPU dispose les Tensors en mémoire afin de maximiser l'efficacité des calculs et de minimiser le remplissage. Pour minimiser la surcharge de mémoire et optimiser l'efficacité de calcul, l'une des conditions suivantes doit être remplie :

1. La taille totale du lot doit être un multiple de 64 (huit par cœur de TPU) et les dimensions des caractéristiques doivent être un multiple de 128.

2. La taille totale du lot doit être un multiple de 1 024 (128 par cœur de TPU), et les dimensions des caractéristiques doivent être un multiple de 8.

L'efficacité maximale est atteinte avec une taille de lot de 1 024 et des dimensions de caractéristiques multiples de 128, bien que cela ne soit pas toujours possible pour tous les modèles.

Remarque : La dimension de caractéristique désigne la taille cachée d'une couche entièrement connectée ou le nombre de canaux de sortie dans une convolution. Certaines couches ne peuvent pas se conformer à cette règle, en particulier les première et dernière couches du réseau. Ce n'est pas un problème, la plupart des modèles nécessitent un certain remplissage.

Fusion

La fusion est une technique générale qui permet au compilateur XLA d'optimiser les programmes. Une opération fusionnée est la combinaison de plusieurs opérations constitutives devant être exécutées ensemble.

Par exemple, considérons la série d'opérations suivante :

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Ce code est à peu près équivalent au pseudo-code suivant :

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result[i] = tmp[i] * z[i];
    }

Avec la fusion, les accès au tableau surviennent en même temps :

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

Dans cet exemple, le nombre d'allers-retours en mémoire est réduit et XLA n'a pas besoin d'allouer d'espace pour "tmp".

La fusion est une optimisation critique et profite à Cloud TPU de plusieurs manières :

• Elle réduit les transferts mémoire en éliminant la nécessité de stocker des résultats intermédiaires dans la mémoire principale, qui est lente.
• Elle permet un usage accru des unités matérielles, qui autrement seraient inutilisées.
• Elle peut réduire l'utilisation de la mémoire d'un modèle, car le nombre de tampons devant être actifs en même temps est moins élevé.

Broadcasting

Le broadcasting survient implicitement lorsque deux Tensors de formes différentes, mais compatibles, sont combinés.

Par exemple, tf.add(vector, matrix) nécessite que le vecteur soit broadcasté vers la forme de la matrice. Le résultat de l'opération a la même forme que la matrice. Pour en savoir plus, consultez le guide sur le broadcasting des tableaux.

Bien que les broadcasts puissent souvent être fusionnés avec leurs destinataires, le fait de les forcer peut entraîner des performances médiocres et une utilisation accrue de la mémoire.

Dans l'exemple suivant, le broadcast implicite dans l'ajout d'un vecteur et d'une matrice ne peut pas être fusionné avec l'argmax, ce qui entraîne sa matérialisation :

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`