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.

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

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.

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.

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 = ...

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.

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
(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.

image

image