Guide sur les performances de Cloud TPU
Pour résoudre les problèmes de performances du TPU, la première étape consiste à profiler votre modèle. Pour en savoir plus sur la capture d'un profil de performances, consultez la page Profiler votre modèle sur Cloud TPU.
Performances du modèle de 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é aux entrées
Les TPU effectuent des calculs très rapidement. Pour s'assurer que le TPU n'est pas inactif, il est important de s'assurer qu'un flux de données constant y est chargé. La méthode dépend de la manière 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 faible en raison de la segmentation (répartition des lots sur plusieurs cœurs)
L'environnement d'exécution TPU répartit 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, la taille de votre lot doit être divisible par 8 ou 128.
Optimisations du compilateur XLA
XLA est un compilateur pour le machine learning capable de produire des binaires pour les TPU, les processeurs, les GPU et d'autres plates-formes. Bien que XLA fasse partie du code base TensorFlow standard, il peut également être utilisé sur les modèles PyTorch et JAX. Les modèles de Cloud TPU sont traduits en graphe XLA, que XLA compile ensuite en exécutable TPU. Pour en savoir plus sur XLA, consultez XLA: Optimizing Compiler for Machine Learning.
Remplissage
Pour utiliser efficacement la mémoire TPU, structurez vos données de sorte qu'elles puissent être divisées en fragments de 128 x 8. Lorsque les données d'un calcul matriciel ne remplissent pas l'intégralité d'un fragment de 128 x 8, le compilateur XLA remplit les Tensors. Cette méthode présente deux inconvénients:
- Les Tensors remplis sous-utilisent le cœur du TPU.
- Le remplissage augmente la quantité de mémoire sur puce requise pour un Tensor, ce qui peut 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 marge intérieure effectuée à l'aide de l'outil de visualisation de mémoire. Vous pouvez éviter le remplissage en choisissant des dimensions de Tensor adaptées au TPU.
Dimensions du Tensor
Le compilateur XLA arrondit les tailles des Tensors stockés dans la mémoire du TPU HBM 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é de calcul et minimiser le remplissage. Pour réduire la surcharge de la mémoire et optimiser l'efficacité des calculs, l'une des conditions suivantes doit être remplie:
La taille totale de 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 de 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.
Si vous utilisez une taille de lot de 1 024 et des dimensions de caractéristiques multiples de 128, l'efficacité est optimale, 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 qui doivent ê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 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 annonces puissent souvent être fusionnées avec leurs consommateurs, forcer une annonce peut entraîner de mauvaises performances et une utilisation de la mémoire plus importante.
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)`