Profiler des charges de travail PyTorch XLA

Le profilage permet d'analyser et d'améliorer les performances des modèles. Bien qu'il y ait beaucoup plus de données, il est parfois utile de considérer le profilage comme des opérations de minutage et des parties du code qui s'exécutent à la fois sur des appareils (TPU) et sur des hôtes (processeurs). Ce guide explique rapidement comment profiler votre code à des fins d'entraînement ou d'inférence. Pour en savoir plus sur l'analyse des profils générés, consultez les guides suivants.

Commencer

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
    

    Descriptions des variables d'exportation

    TPU name
    Nom que vous souhaitez utiliser pour votre Cloud TPU.
    zone
    Zone dans laquelle vous prévoyez de créer la ressource Cloud TPU.
    project ID
    ID du projet que vous utilisez pour entraîner et profiler votre modèle.
    accelerator-type
    Type du TPU à créer.
    version
    Version d'exécution de Cloud TPU. Une valeur par défaut est affichée 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

Un profil peut être capturé manuellement via capture_profile.py ou par programmation depuis le script d'entraînement à l'aide des API torch_xla.debug.profiler.

Démarrage du serveur de profils

Pour capturer un profil, un serveur de profils doit être en cours d'exécution dans le 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 désormais capturer des profils comme indiqué dans la section suivante. Le script décrit tout ce qui se passe sur un appareil TPU.

Ajouter des traces

Si vous souhaitez également profiler des opérations sur la machine hôte, vous pouvez ajouter xp.StepTrace ou xp.Trace dans votre code. Ces fonctions suivent le code Python sur la machine hôte. Vous pouvez considérer cela comme la mesure du temps nécessaire pour exécuter le code Python sur l'hôte (CPU) avant de transmettre le "graphe" à l'appareil TPU. Il est donc surtout utile pour analyser la surcharge de 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 encapsuler des parties individuelles du code avec

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

Si vous utilisez Lighting, vous pouvez ignorer l'ajout de traces, car cette opération est effectuée 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, en attendant que les étapes d'entraînement ou d'inférence du modèle démarre.

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 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 .xplane.pb fichiers dans 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.

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 utilise les 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 alors enregistrés dans le répertoire spécifié par la variable d'environnement PROFILE_LOGDIR.

Analyse dans TensorBoard

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

Pour exécuter TensorBoard sur une machine distante, connectez-vous à celle-ci à 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 packages requis et lancez TensorBoard (en supposant que vous disposez de profils sous ~/profiles/ sur cette machine). Si vous avez stocké les profils dans un autre répertoire ou bucket Cloud Storage, veillez à spécifier correctement les chemins d'accès, par exemple gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow 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/, puis sélectionnez PROFILE dans le menu déroulant pour charger vos profils.

Reportez-vous à la section Outils TPU pour en savoir plus sur les outils TensorBoard et sur l'interprétation du résultat.

Image

Image