Profiler votre modèle sur des VM Cloud TPU

Le profilage est l'un des principaux outils permettant d'optimiser les performances de vos modèles sur Cloud TPU. L'outil de profilage principal, XProf, est disponible dans le dépôt GitHub OpenXLA/XProf. XProf permet de profiler tous les frameworks basés sur XLA, y compris JAX, Pytorch XLA et Tensorflow/Keras.

Capturer des profils

Avant d'utiliser XProf, vous devez capturer un profil de votre modèle de l'une des deux manières suivantes :

  1. Capture programmatique
  2. Capture à la demande (ou capture manuelle)

Avec la capture programmatique, vous devez annoter le code de votre modèle pour spécifier les endroits où capturer les profils dans le code. En règle générale, il convient de capturer un profil pour quelques étapes d'entraînement ou de profiler un bloc de code spécifique dans le modèle. Il existe différentes manières de capturer des traces dans les différents frameworks JAX, Pytorch XLA et TensorFlow : soit avec une API qui démarre ou arrête la trace, soit avec un gestionnaire de contexte. Dans les frameworks de niveau supérieur tels que MaxText, vous pouvez activer la capture de profil en activant simplement un flag profiler=xplane lors de l'exécution de MaxText.

Utilisez la capture de profil à la demande si vous souhaitez capturer des profils de manière ad hoc ou si vous n'avez pas activé la capture de profil programmatique. Il peut être utile d'opter pour cette méthode lorsque vous constatez un problème dans les métriques de votre modèle pendant l'exécution, et que vous souhaitez capturer des profils à ce moment précis et pendant un certain temps en vue de diagnostiquer le problème.

Pour activer la capture de profil à la demande, vous devrez toujours démarrer le serveur XProf dans votre code. L'activation de profiler.start_server démarre sur votre charge de travail de ML un serveur XProf qui attend le déclencheur de capture à la demande pour commencer à capturer les profils. Vous pouvez déclencher la capture de profil à la demande via l'UI TensorBoard ou via la CLI avec l'outil XProfiler.

Pour savoir comment activer la capture programmatique et à la demande pour différents frameworks, consultez :

Pour la capture de profil programmatique comme pour la capture de profil à la demande, vous devez spécifier l'emplacement où stocker les profils capturés. Vous pouvez choisir de stocker les profils dans un répertoire de votre VM TPU avec un chemin d'accès semblable à /profiles/run-name ou utiliser Cloud Storage avec un chemin d'accès semblable à gs://bucket-name/run-name/.

Avec les chemins d'accès indiqués ci-dessus, vos profils seront capturés sous :

/profiles/run-name/plugins/profile/session1/

ou

gs://bucket-name/run-name/plugins/profile/session1/.

Une session représente une capture de profils lors d'une exécution de votre modèle. Supposons que, pendant un cycle d'entraînement, vous capturez des profils lors des étapes 1 à 3, puis lors des étapes 8 à 10. Ces profils font partie du même cycle, mais la première capture des étapes 1 à 3 sera la session1 et la deuxième capture des étapes 8 à 10 sera la session2.

Dans chaque cycle, les différentes sessions portent des codes temporels distincts. Les profils capturés lors de ces différentes sessions peuvent l'être de manière programmatique, à la demande ou en combinant les deux, ce qui s'avère utile si vous souhaitez comparer des profils capturés à différents moments de l'entraînement de votre modèle (en comparant par exemple les profils du début de l'entraînement avec ceux de la fin de l'entraînement).

Visualiser les profils avec TensorBoard

Pour afficher les profils XProf dans TensorBoard, vous devez installer le plug-in tensorboard-"plugin-profile".

Pour visualiser les profils capturés sur Google Cloud, nous vous recommandons d'utiliser TensorBoard sur une VM TPU, ou l'outil XProfiler pour héberger TensorBoard sur une VM ou un pod GKE.

Si vous avez capturé vos profils sur la VM TPU, vous pouvez les visualiser sur cette même VM TPU à l'aide de TensorBoard.

Pour installer XProf et TensorBoard sur une VM TPU, vous pouvez exécuter les commandes suivantes :

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Pour visualiser les profils capturés avec TensorBoard, exécutez :

$ tensorboard --logdir=profiles/run-name

ou

$ tensorboard --logdir=profiles

Pour charger plusieurs profils provenant de plusieurs cycles, pointez TensorBoard vers le répertoire racine contenant toutes vos exécutions et tous vos profils, plutôt que vers le sous-répertoire contenant les profils d'une exécution spécifique.

Bibliothèque Cloud-Diagnostics-XProf pour visualiser les profils Tensorboard sur Google Cloud

cloud-diagnostics-xprof (la bibliothèque XProfiler) facilite l'hébergement de TensorBoard et la visualisation des profils sur Google Cloud. Pour en savoir plus, consultez le dépôt GitHub cloud-diagnostics-xprof.

Le dépôt cloud-diagnostics-xprof fournit les améliorations suivantes par rapport à une exécution de XProf et TensorBoard en local :

  • Configuration et packaging des dépendances XProf et TensorBoard.
  • Stockage de vos profils dans Cloud Storage, utile pour la conservation à long terme et l'analyse post-exécution (les profils locaux capturés seront supprimés une fois l'exécution terminée).
  • Chargement rapide de profils volumineux et de plusieurs profils grâce au provisionnement de TensorBoard sur une VM Compute Engine ou un pod GKE, avec la possibilité de modifier le type de machine en fonction des besoins de l'utilisateur en termes de vitesse de chargement et de coût.
  • Création d'un lien facilitant le partager des profils et la collaboration entre les membres d'équipe et les ingénieurs Google.
  • Profilage à la demande des charges de travail sur GKE et Compute Engine facilité : n'importe quel hôte exécutant votre charge de travail peut être sélectionné pour la capture des profils.

Avant d'utiliser la bibliothèque XProfiler, vous devez capturer des profils pour votre code de charge de travail soit de manière programmatique, soit en démarrant le serveur de profil, afin de pouvoir capturer les profils à la demande ultérieurement.

Pour configurer la bibliothèque XProfiler, vous devez avoir déjà configuré Google Cloud CLI et un environnement virtuel Python pour qu'il ne vous reste plus qu'à exécuter la commande suivante :

pip install cloud-diagnostics-xprof

Cette commande installe toutes les dépendances XProf et TensorBoard nécessaires.

Ensuite, pour créer une VM ou un pod GKE qui hébergera TensorBoard, exécutez la commande suivante :

xprofiler create -z $ZONE -l $GCS_PATH

ou

xprofiler create --GKE -z $ZONE -l $GCS_PATH

Remplacez $ZONE par n'importe quelle zone et $GCS_PATH par le chemin d'accès à vos traces de profil. Vous pouvez spécifier le répertoire racine qui contient toutes les traces de profil pour plusieurs exécutions, ou un ensemble spécifique de traces de profil pour une seule exécution.

Par exemple, si vous capturez des profils dans les répertoires suivants :

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

Vous pouvez définir le chemin d'accès GCS au répertoire racine (gs://bucket-name) :

xprofiler create -z $ZONE -l gs://bucket-name/

Dans l'UI TensorBoard, vous verrez tous les profils pour toutes les exécutions et sessions sous la forme run1/session1, run1/session2 et run2/session1.

Par défaut, xprofiler create crée une VM Compute Engine, plus exactement une VM c4-highmem-8. Vous pouvez modifier le type de machine à l'aide du flag -m. Si vous souhaitez créer une instance TensorBoard sur un pod GKE au lieu d'une VM Compute Engine, vous pouvez activer le flag –GKE pour xprofiler create. Le fait d'héberger une instance TensorBoard sur un pod GKE peut simplifier la gestion de TensorBoard avec le reste de votre charge de travail déployée sur GKE.

Le chargement de profils volumineux et de plusieurs profils sur Google Cloud est beaucoup plus rapide si l'instance TensorBoard est hébergée sur une VM ou un pod GKE plutôt qu'en local. Selon nos tests comparatifs, le premier chargement de profils de l'ordre de 1 Go avec la VM c4-highmem-8 par défaut se fait en quelques minutes. Vous pouvez également choisir le type de machine adapté à vos besoins en termes de performances et de coûts.

Après avoir exécuté xprofiler create, un résultat semblable à ce qui suit s'affiche :

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

La première option est un lien cliquable permettant d'afficher vos profils XProf sur un TensorBoard. Ce lien permet de partager les profils avec votre équipe et avec les ingénieurs Google pour optimiser vos performances sur Google Cloud.

C'est vous qui contrôlez qui a accès au lien en fonction des autorisations définies pour le bucket Cloud Storage qui stocke les données de votre profil. Au cas où le lien ne fonctionnerait pas, nous vous offrons la possibilité de vous connecter à l'instance TensorBoard via SSH pour visualiser vos profils à l'aide de la commande xprofiler connect.

Si vous avez activé le serveur Cloud Profiler dans le code de votre charge de travail (ce qui est nécessaire pour les captures programmatiques comme pour celles à la demande) et que vous souhaitez effectuer un profilage à la demande, vous pouvez le faire de deux manières :

a. Cliquez sur le bouton Capturer le profil dans TensorBoard. Vous pouvez désormais choisir de profiler n'importe quel hôte d'appareil sur lequel votre charge de travail s'exécute. Nous prenons en charge la capture à la demande pour les charges de travail exécutées sur GKE ou Compute Engine.

b. Utilisez la commande de capture XProfiler en indiquant bien la zone, le bucket Cloud Storage, le framework, les noms de VM / de pod hôtes et la durée de la capture en millisecondes. Les mêmes informations doivent être saisies dans l'UI TensorBoard.

Pour en savoir plus sur l'utilisation de la bibliothèque cloud-diagnostics-xprof, consultez cette page GitHub.