Profiler des charges de travail PyTorch XLA

Le profilage permet d'analyser et d'améliorer les performances des modèles. Bien que il y a beaucoup plus à voir, mais il est parfois utile de considérer le profilage comme le timing opérations et parties du code qui s'exécutent à la fois sur les appareils (TPU) et sur les hôtes (processeurs). Ce guide explique rapidement comment profiler votre code pour l'entraînement ou l'inférence. Pour plus d'informations sur l'analyse des profils générés, veuillez consulter 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 de la ressource 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 Versions de TPU.

   version

   Version d'exécution de Cloud TPU. Une valeur par défaut s'affiche dans le fichier , mais vous pouvez aussi en utiliser une de la liste 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 gcloud scp. Exemple :

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

Profilage

Remarque :La version v4-8 comprend quatre appareils TPU. Chaque nœud de calcul ne profile qu'un seul appareil TPU sur les 4 appareils disponibles. Le profilage 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 automatisée depuis le script d'entraînement à l'aide API torch_xla.debug.profiler.

Démarrage du serveur de profils

Pour que vous puissiez capturer un profil, un serveur de profils doit être en cours d'exécution dans la script d'entraînement. Démarrez un serveur avec le numéro de port de votre choix, par exemple. 9012, comme indiqué dans la commande suivante.

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

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

Le serveur peut être démarré directement au début de votre fonction main.

Remarque :Si vous utilisez Lightning au lieu de démarrer un serveur, transmettez profiler="xla" à trainer, et démarrera automatiquement le serveur 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 au profil des opérations sur la machine hôte, vous pouvez ajouter xp.StepTrace ou xp.Trace dans votre code. Ces fonctions tracent le code Python la machine hôte. Vous pouvez considérer que c'est comme mesurer le temps nécessaire pour exécuter sur l'hôte (CPU) avant de transmettre le "graphe" à l'appareil TPU. Il est donc surtout utile pour analyser les frais généraux de traçage. Vous pouvez l'ajouter dans la boucle d'entraînement, où le code traite des lots de données, Exemple :

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

ou encapsulez des parties individuelles du code

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

Important :Évitez d'encapsuler xm.mark_step() avec xp.Trace. Le code plante dans ce cas avec un message d'erreur semblable à ceci: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace ne rencontrera pas ce problème, car il ajoute automatiquement xm.mark_step() lors de la sortie du champ d'application De plus, cette erreur peut se produire Vous utilisez torchdynamo, car torchdynamo peut appeler xm.mark_step en interne. La solution actuelle consiste à ne pas utiliser xp.Trace.

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

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

Capture manuelle

Script capture_profile.py du dépôt Pytorch XLA permet de capturer un profil rapidement. Pour ce faire, copiez capturer le fichier de profil directement à 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 la commande suivante 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 logdir. Vous pouvez modifier Répertoire de journalisation ~/profiles/ à l'emplacement et au nom de votre choix. Il est également possible de l'enregistrer directement dans le bucket Cloud Storage. Pour ce faire, définissez logdir devient gs://your_bucket_name/.

Capture programmatique

Au lieu de capturer le profil manuellement en déclenchant un script, vous pouvez configurer votre script d'entraînement pour déclencher automatiquement un profil à l'aide de torch_xla.debug.profiler.trace_detached, dans votre script d'entraînement.

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

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

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 seront enregistrés dans le répertoire spécifié par le paramètre PROFILE_LOGDIR.

Analyse dans TensorBoard

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

Pour exécuter TensorBoard sur une machine distante, connectez-vous à celle-ci via SSH et activez le réacheminement 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 packages requis et lancez TensorBoard (en supposant que vous ayez des profils sur cette machine sous ~/profiles/). Si vous avez stocké les profils d'un autre répertoire ou bucket Cloud Storage, assurez-vous spécifier correctement les chemins d'accès (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 d'avoir des fichiers .xplane dans la même structure imbriquée que celle est généré lorsque vous capturez un profil. Par exemple: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Remarque :Si le message "Duplicate plugins" (Plug-ins en double) s'affiche, il se peut qu'il y ait plusieurs packages avec TensorBoard. Désinstallez tous ces packages, puis exécutez la à nouveau les commandes pip install. Par exemple, vous devrez peut-être exécuter la commande suivante:

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

Exécuter TensorBoard

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

Reportez-vous à la section Outils TPU pour en savoir plus sur les les outils TensorBoard et comment interpréter les résultats.