Profilage des charges de travail PyTorch XLA

Le profilage permet d'analyser et d'améliorer les performances des modèles. Bien que le profilage soit bien plus complexe, il est parfois utile de le considérer comme des opérations de chronométrage et des parties du code qui s'exécutent à la fois sur les appareils (TPU) et les hôtes (CPU). Ce guide présente brièvement comment profiler votre code pour l'entraînement ou l'inférence. Pour en savoir plus sur l'analyse des profils générés, consultez les guides suivants.

• Débogage des performances PyTorch XLA sur les VM TPU – Partie 1
• Débogage des performances PyTorch XLA sur les VM TPU – Partie 2
• Débogage des performances PyTorch XLA sur les VM TPU – Partie 3

Premiers pas

Créer un TPU

1. Exportez les variables d'environnement:

   $ export TPU_NAME=your_tpu_name
   $ export ZONE=us-central2-b
   $ export PROJECT_ID=project-id
   $ export ACCELERATOR_TYPE=v4-8
   $ export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

   Exporter les descriptions des variables

   TPU name

   Nom que vous souhaitez utiliser pour votre Cloud TPU.

   zone

   Zone dans laquelle vous prévoyez de créer votre Cloud TPU.

   project ID

   ID du projet que vous utilisez pour entraîner et profiler votre modèle.

   accelerator-type

   Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez la section Versions de TPU.

   version

   Version d'exécution de Cloud TPU. Une valeur par défaut s'affiche dans la variable exportée, mais vous pouvez également en utiliser une dans la liste des configurations compatibles.

2. Lancer les ressources TPU

   $ gcloud compute tpus tpu-vm create ${TPU_NAME} \
   --zone us-central2-b \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --version ${RUNTIME_VERSION} \
   --project $PROJECT_ID \
   --subnetwork=tpusubnet

3. Déplacez votre code vers votre répertoire d'accueil sur la VM TPU à l'aide de la commande gcloud scp. Exemple :

   $ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}: --zone ${ZONE}

Profilage

Remarque :Un v4-8 comprend quatre appareils TPU. Chaque travailleur ne profile qu'un seul appareil TPU sur les quatre disponibles. Le profilage de l'ensemble des quatre appareils TPU pour chaque nœud de calcul sera bientôt disponible.

Un profil peut être capturé manuellement via capture_profile.py ou de manière programmatique à partir du script d'entraînement à l'aide des API torch_xla.debug.profiler.

Démarrer le serveur de profils

Pour capturer un profil, un serveur de profil doit s'exécuter dans le script d'entraînement. Démarrez un serveur avec un numéro de port de votre choix, par exemple 9012, comme indiqué dans la commande suivante.

Remarque :Vous pouvez vous reporter à cet exemple de script qui inclut les étapes suivantes.

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

Le serveur peut être démarré dès le début de votre fonction main.

Remarque :Si vous utilisez le module Lightning, au lieu de démarrer un serveur, transmettez profiler="xla" à trainer. Le serveur démarrera automatiquement sur le port 9012.

Vous pouvez maintenant capturer des profils comme décrit dans la section suivante. Le script profile tout ce qui se passe sur un appareil TPU.

Ajouter des traces

Si vous souhaitez également profiler les opérations sur la machine hôte, vous pouvez ajouter xp.StepTrace ou xp.Trace dans votre code. Ces fonctions tracent le code Python sur la machine hôte. (Vous pouvez considérer cela comme une mesure du temps nécessaire pour exécuter le code Python sur l'hôte (processeur) avant de transmettre le "graphe" à l'appareil TPU. Il est donc principalement utile pour analyser les coûts liés au traçage. Vous pouvez l'ajouter dans la boucle d'entraînement où le code traite des lots de données, par exemple :

for step, batch in enumerate(train_dataloader):
    with xp.StepTrace('Training_step', step_num=step): 
       ...

ou encapsulant des parties individuelles du code avec

 with xp.Trace('loss'): 
    loss = ...

Important :Évitez d'encapsuler xm.mark_step() avec xp.Trace. Dans ce cas, le code plante avec un message d'erreur semblable à: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace ne présente pas ce problème, car il ajoute automatiquement xm.mark_step() à la sortie du champ d'application. De plus, vous pouvez rencontrer cette erreur lorsque vous utilisez torchdynamo, car torchdynamo peut appeler xm.mark_step en interne. La solution actuelle consiste à ne pas utiliser xp.Trace.

Si vous utilisez Lighting, vous pouvez ignorer l'ajout de traces, car cela se fait automatiquement dans certaines parties du code. Toutefois, si vous souhaitez ajouter des traces supplémentaires, vous pouvez les insérer dans la boucle d'entraînement.

Vous pourrez capturer l'activité de l'appareil après la compilation initiale. Attendez que le modèle commence ses étapes d'entraînement ou d'inférence.

Capture manuelle

Le script capture_profile.py du dépôt Pytorch XLA permet de capturer rapidement un profil. Pour ce faire, copiez le fichier de profil de capture directement sur votre VM TPU. La commande suivante le copie dans le répertoire d'accueil.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

Pendant l'entraînement, exécutez les commandes suivantes pour capturer un profil:

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

Cette commande enregistre les fichiers .xplane.pb dans le logdir. Vous pouvez remplacer le répertoire de journalisation ~/profiles/ par l'emplacement et le nom de votre choix. Il est également possible d'enregistrer directement dans le bucket Cloud Storage. Pour ce faire, définissez logdir sur gs://your_bucket_name/.

Capture programmatique

Plutôt que de capturer le profil manuellement en déclenchant un script, vous pouvez configurer votre script d'entraînement pour qu'il déclenche automatiquement un profil à l'aide de l'API torch_xla.debug.profiler.trace_detached dans votre script d'entraînement.

Remarque :L'API trace_detached n'est disponible que dans les builds quotidiens ou après la version 2.2. Pour effectuer une capture par programmation dans les versions antérieures de torch_xla, vous pouvez copier le corps de la méthode pour créer directement le thread de capture en arrière-plan.

Par exemple, pour capturer automatiquement un profil à une époque et à une étape spécifiques, vous pouvez configurer votre script d'entraînement pour qu'il consomme des variables d'environnement PROFILE_STEP, PROFILE_EPOCH et PROFILE_LOGDIR:

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

Les fichiers .xplane.pb sont ainsi enregistrés dans le répertoire spécifié par la variable d'environnement PROFILE_LOGDIR.

Analyse dans TensorBoard

Pour analyser plus en détail les profils, vous pouvez utiliser TensorBoard avec le plug-in TensorBoard TPU
sur la même machine ou sur une autre (recommandé).

Pour exécuter TensorBoard sur une machine distante, connectez-vous-y à l'aide de SSH et activez le transfert de port. Par exemple,

$ ssh -L 6006:localhost:6006 remote server address

ou

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --ssh-flag="-4 -L 6006:localhost:6006"

Sur votre machine distante, installez les paquets requis et lancez TensorBoard (en supposant que vous disposez de profils sur cette machine sous ~/profiles/). Si vous avez stocké les profils dans un autre répertoire ou un autre bucket Cloud Storage, veillez à spécifier correctement les chemins, par exemple gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile 

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

Remarque :Il est important que les fichiers .xplane se trouvent dans la même structure imbriquée générée lorsque vous capturez un profil. Par exemple : ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Remarque :Si l'erreur "plugins en double" s'affiche, cela peut être dû à la présence de plusieurs packages avec TensorBoard. Désinstallez tous ces paquets et exécutez à nouveau les commandes pip install. Par exemple, vous devrez peut-être exécuter les commandes suivantes:

(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

Exécuter TensorBoard

Dans votre navigateur local, accédez à http://localhost:6006/, puis sélectionnez PROFILE dans le menu déroulant pour charger vos profils.

Pour en savoir plus sur les outils TensorBoard et sur l'interprétation de la sortie, consultez la section Outils TPU.