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.
- 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
Commencer
Créer un TPU
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.
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
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.