Guide des performances de Cloud TPU

Pour résoudre les problèmes de performances des TPU, la première étape consiste à profiler votre modèle. Pour savoir comment capturer 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.

Le modèle est lié à la saisie

Les TPU réalisent des calculs très rapidement. Pour garantir que le TPU n'est pas inactif, il est important de s'assurer qu'un flux de données stable est chargé sur le 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.
La taille de lot est trop petite en raison de la segmentation (répartition des lots entre cœurs)

L'environnement d'exécution 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 optimiser l'utilisation 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 8X128 multidimensionnels pour traiter les multiplications de matrices. En général, la taille de votre lot doit être équitablement divisée par 8 ou 128.

Optimisations du compilateur XLA

XLA est un compilateur de machine learning qui peut produire des binaires pour les TPU, les processeurs, les GPU et d'autres plates-formes. Bien que XLA fasse partie de la base de code TensorFlow standard, il peut également être utilisé sur les modèles PyTorch et JAX. Les modèles de Cloud TPU sont convertis en un graphique XLA, que XLA compile ensuite dans un exécutable TPU. Pour plus d'informations sur XLA, consultez la page XLA: Optimization Compiler for Machine Learning (XLA : Optimiser le compilateur pour le machine learning).

Remplissage

Pour utiliser la mémoire TPU de manière efficace, structurez vos données de sorte qu'elles puissent être carrelées en morceaux de 128 x 8. Lorsque les données d'un calcul de matrice ne remplissent pas un fragment entier de 128 x 8, le compilateur XLA sollicite les Tensors. L'ajout d'une marge intérieure présente deux inconvénients:

Des Tensors rembourrés sous-utilisent le cœur de TPU.
L'ajout d'une marge extérieure augmente la quantité de stockage sur mémoire nécessaire pour un Tensor et peut entraîner une erreur de mémoire insuffisante.

Bien que le remplissage soit automatiquement effectué par le compilateur XLA lorsque nécessaire, vous pouvez le déterminer à l'aide de l'outil de visualisation de la mémoire. Vous pouvez éviter la marge intérieure en choisissant des dimensions de Tensor adaptées aux TPU.

Dimensions Tensor

Le compilateur XLA arrondit la taille des Tensors stockés dans la mémoire HBM du TPU pour effectuer des calculs plus efficacement. Cette marge intérieure s'effectue 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'environnement d'exécution TPU dispose de Tensors en mémoire pour maximiser l'efficacité du calcul et minimiser la marge intérieure. Pour minimiser la surcharge mémoire et optimiser l'efficacité de calcul, l'une des conditions suivantes doit être respectée:

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

L'efficacité de la taille de lot est de 1 024 avec des dimensions de caractéristiques multiples de 128, mais cela peut ne pas être possible pour tous les modèles.

Fusion

Fusion est une technique générale utilisée par le compilateur XLA pour optimiser les programmes. Une opération fusionnée est la combinaison de plusieurs opérations constitutives devant être exécutées en même temps.

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 = 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 = (x[i] + y[i]) * z[i];
    }

Dans cet exemple, le nombre de trajets aller-retour en mémoire est réduit, et XLA n'a pas besoin d'allouer de l'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 annonces puissent souvent être confondues avec leurs clients, le fait de les forcer peut entraîner des performances médiocres et augmenter l'utilisation 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)`