Guide des performances de Cloud TPU
La première étape du dépannage des performances 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 des modèles 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 subordonné aux entrées
Les TPU effectuent des calculs très rapidement. Pour que le TPU ne soit pas inactif, il est important de s'assurer qu'un flux constant de données 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 sur plusieurs cœurs)
L'environnement d'exécution TPU divise un lot sur les huit cœurs d'un appareil TPU (par v2-8 ou v3-8, par exemple). 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 taille de lot la plus grande adaptée à TPU mémoire. Chaque cœur de TPU utilise des registres vectoriels 8 x 128 bidimensionnels pour traiter les multiplications matricielles. En général, votre taille de lot doit être divisible de manière uniforme par 8 ou 128.
Optimisations du compilateur XLA
XLA est un compilateur de code capable de produire des binaires pour les TPU, les CPU, les GPU et d'autres plates-formes. XLA fait partie du code base TensorFlow standard, mais peut aussi être utilisé sur 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 TPU, structurez vos données de sorte qu'elles puissent être regroupées 128 x 8 morceaux. Lorsque les données d'un calcul matriciel ne remplissent pas l'intégralité d'un bloc de 128 x 8, le compilateur XLA remplit les tensors. Cette méthode présente toutefois deux inconvénients :
- Les Tensors rembourrés sous-utilisent le cœur du TPU.
- 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éterminer la quantité de remplissage effectuée à l'aide de l'outil d'affichage de la mémoire. Vous pouvez éviter le remplissage en choisissant des dimensions de Tensors adaptées aux TPU.
Dimensions de Tensor
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'environnement d'exécution TPU dispose des Tensors en mémoire pour optimiser l'efficacité de calcul et à réduire la marge intérieure. Pour minimiser la surcharge de mémoire et optimiser les ressources de calcul l'efficacité, l'une des conditions suivantes doit être remplie:
La taille totale du lot doit être un multiple de 64 (8 par cœur de TPU) et les dimensions des caractéristiques doivent être un multiple de 128.
La taille totale du lot doit être un multiple de 1 024 (128 par cœur de TPU). 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.
Fusion
La Fusion est une technique générale que le compilateur XLA utilise pour 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 = 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 d'allers-retours de la mémoire est réduit, mais pas XLA 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 fusionnées avec leurs clients, peut entraîner de mauvaises performances et une augmentation de 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)`