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 s'appelle XProf. Il est disponible dans le dépôt GitHub OpenXLA/XProf. XProf est compatible avec le profilage de tous les frameworks basés sur XLA, y compris JAX, Pytorch XLA et TensorFlow/Keras.

Capturer des profils

Pour utiliser XProf, vous devez d'abord capturer un profil de votre modèle. Il existe deux façons de capturer des profils :

  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 où vous souhaitez capturer les profils dans votre code. En règle générale, vous capturez un profil pour quelques étapes d'entraînement ou vous profilez un bloc de code spécifique dans votre modèle. Il existe différentes manières de capturer des traces dans les différents frameworks JAX, Pytorch XLA et TensorFlow : soit en démarrant ou en arrêtant la trace à l'aide d'une API, soit à l'aide d'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 indicateur profiler=xplane lors de l'exécution de MaxText.

Vous pouvez utiliser la capture de profil à la demande lorsque vous souhaitez capturer des profils de manière ad hoc ou lorsque vous n'avez pas activé la capture de profil programmatique. Cela peut être utile lorsque vous constatez un problème avec les métriques de votre modèle pendant l'exécution et que vous souhaitez capturer des profils à ce moment précis pendant une certaine période afin de diagnostiquer le problème.

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

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

Pour la capture de profil programmatique et à la demande, vous devez spécifier l'emplacement où vous souhaitez 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 mentionné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 de l'exécution de votre modèle. Par exemple, si vous capturez des profils dans une série d'entraînement à partir des étapes 1 à 3, puis que vous capturez des profils à partir des étapes 8 à 10, il s'agit de profils pour la même série, mais la première capture des étapes 1 à 3 sera session1 et la deuxième capture des étapes 8 à 10 sera session2.

Les différentes sessions seront indiquées par des codes temporels différents sous chaque exécution. Vous pouvez capturer des profils dans différentes sessions de manière programmatique, à la demande ou en combinant les deux. Cela est utile lorsque vous souhaitez comparer des profils provenant de différentes parties de l'entraînement de votre modèle. Par exemple, en comparant 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 afficher à l'aide de TensorBoard sur cette VM TPU.

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 afficher les profils capturés avec TensorBoard, exécutez :

$ tensorboard --logdir=profiles/run-name

ou

$ tensorboard --logdir=profiles

Pour charger plusieurs profils à partir de plusieurs exécutions, 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 (également appelée 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 offre les avantages suivants par rapport à l'exécution de XProf et TensorBoard en local :

  • Configuration et packaging des dépendances XProf et TensorBoard
  • Stockez vos profils dans Cloud Storage, ce qui peut être 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 en provisionnant 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éer un lien pour partager des profils et collaborer avec des membres de l'équipe et des ingénieurs Google
  • Profilage à la demande plus facile des charges de travail sur GKE et Compute Engine pour choisir n'importe quel hôte exécutant votre charge de travail afin de capturer les profils

Avant d'utiliser la bibliothèque XProfiler, vous devez capturer les profils de votre code de charge de travail de manière programmatique ou 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 configurer Google Cloud CLI et un environnement virtuel Python. Ensuite, il vous suffit d'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 pour héberger 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'interface utilisateur 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 précisément une VM c4-highmem-8. Vous pouvez modifier le type de machine à l'aide de l'option -m. Si vous souhaitez créer une instance TensorBoard sur un pod GKE au lieu d'une VM Compute Engine, vous pouvez transmettre l'indicateur –GKE à xprofiler create. L'hébergement d'une instance TensorBoard sur un pod GKE peut faciliter la gestion de TensorBoard avec le reste de votre charge de travail déployée sur GKE.

La VM ou le pod GKE qui héberge TensorBoard permet de charger des profils volumineux et plusieurs profils beaucoup plus rapidement sur Google Cloud que sur un TensorBoard hébergé en local. D'après nos tests comparatifs, vous verrez des profils de l'ordre de 1 Go de charge en quelques minutes pour la première charge à l'aide de la VM c4-highmem-8 par défaut. 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 sur lequel vous pouvez cliquer pour afficher vos profils XProf sur un TensorBoard. Ce lien permet de partager des profils avec votre équipe et avec les ingénieurs Google qui vous aident à optimiser les performances sur Google Cloud.

Vous contrôlez qui a accès au lien en fonction des autorisations définies pour le bucket Cloud Storage qui stocke vos données de profil. Si le lien ne fonctionne pas pour une raison quelconque, nous vous proposons également un moyen de vous connecter à l'instance TensorBoard à l'aide de SSH afin d'afficher 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 et à la demande) et que vous souhaitez effectuer un profilage à la demande, vous pouvez le faire de deux manières :

a. Cliquez sur le bouton Enregistrer le profil dans TensorBoard. Vous pouvez désormais choisir n'importe quel hôte d'appareil sur lequel votre charge de travail s'exécute pour le profilage. Nous acceptons 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 fournissant des informations telles que 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. Il s'agit des mêmes informations que celles que vous devez saisir dans l'UI TensorBoard.

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