Profila carichi di lavoro XLA PyTorch
La profilazione è un modo per analizzare e migliorare le prestazioni dei modelli. Sebbene c'è molto altro, a volte può essere d'aiuto pensare alla profilazione come ai tempi operazioni e parti di codice in esecuzione su dispositivi (TPU) e host (CPU). Questa guida fornisce una rapida panoramica su come profilare il codice per l'addestramento o l'inferenza. Per ulteriori informazioni su come analizzare i profili generati, consulta alle seguenti guide.
- Debug delle prestazioni di PyTorch XLA sulle VM TPU - Parte 1
- Debug delle prestazioni di PyTorch XLA sulle VM TPU - Parte 2
- Debug delle prestazioni di PyTorch XLA sulle VM TPU - parte 3
Inizia
Crea una TPU
Esporta le variabili di ambiente:
$ 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
Esporta descrizioni variabili
TPU name
- Il nome che vuoi utilizzare per Cloud TPU.
zone
- La zona in cui prevedi di creare la tua Cloud TPU.
project ID
- L'ID progetto che usi per addestrare e profilare il modello.
accelerator-type
- Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi Versioni TPU.
version
- La versione del runtime di Cloud TPU. Nel file esportato viene mostrato un valore predefinito ma puoi anche utilizzarne una dall'elenco configurazioni supportate.
Avvia le risorse 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
Sposta il codice nella home directory sulla VM TPU utilizzando Comando
gcloud scp
. Ad esempio:$ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}: --zone ${ZONE}
Profilazione
È possibile acquisire manualmente un profilo tramite capture_profile.py
oppure
in modo programmatico dallo script di addestramento
torch_xla.debug.profiler
API.
Avvio del server del profilo
Per acquisire un profilo, è necessario che un server di profilo sia in esecuzione all'interno
script di addestramento. Avvia un server con un numero di porta a tua scelta, ad esempio
9012
come mostrato nel seguente comando.
import torch_xla.debug.profiler as xp server = xp.start_server(9012)
Il server può essere avviato direttamente all'inizio della funzione main
.
Ora puoi acquisire profili come descritto nella sezione seguente. Lo script profila tutto ciò che accade su un dispositivo TPU.
Aggiunta di tracce
Se vuoi sapere anche
alle operazioni del profilo sulla macchina host, puoi aggiungere xp.StepTrace
xp.Trace
nel codice. Queste funzioni tracciano il codice Python
della macchina host.
Questo è un po' come misurare il tempo necessario per eseguire
sull'host (CPU) prima di passare il "grafico" al dispositivo TPU. quindi è soprattutto utile per analizzare l'overhead di tracciamento). Puoi
aggiungerlo all'interno del loop di addestramento in cui il codice elabora batch di dati,
ad esempio
for step, batch in enumerate(train_dataloader):
with xp.StepTrace('Training_step', step_num=step):
...
o aggregare singole parti del codice
with xp.Trace('loss'):
loss = ...
Se utilizzi Illuminazione, puoi saltare l'aggiunta di tracce perché l'operazione viene eseguita automaticamente in alcune parti del codice. Tuttavia, se vuoi aggiungere altre tracce, a inserirli nel loop di addestramento.
Potrai acquisire l'attività del dispositivo dopo la compilazione iniziale. attesa finché il modello non inizia i passaggi di addestramento o inferenza.
Acquisizione manuale
Lo script capture_profile.py
del repository Pytorch XLA
consente di acquisire rapidamente un profilo. È possibile farlo copiando
file del profilo di acquisizione
direttamente sulla VM TPU. Il comando seguente lo copia nella home directory.
$ 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"
Durante l'addestramento, esegui il comando seguente per acquisire un profilo:
$ 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"
Questo comando salva .xplane.pb
file in logdir
. Puoi modificare
la directory di logging
~/profiles/
al nome e alla posizione che preferisci. È inoltre possibile
salvare direttamente nel bucket Cloud Storage. Per farlo, imposta
logdir
in modo che sia gs://your_bucket_name/
.
Acquisizione programmatica
Anziché acquisire il profilo manualmente attivando uno script, puoi configurare lo script di addestramento per attivare automaticamente un profilo mediante il comando torch_xla.debug.profiler.trace_detached. all'API Cloud Storage all'interno dello script di addestramento.
Ad esempio, per acquisire automaticamente un profilo in un'epoca e in un passaggio specifici,
puoi configurare lo script di addestramento in modo che utilizzi PROFILE_STEP
,
Ambiente PROFILE_EPOCH
e PROFILE_LOGDIR
variabili:
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)
...
I file .xplane.pb
verranno salvati nella directory specificata
Variabile di ambiente PROFILE_LOGDIR
.
Analisi in TensorBoard
Per analizzare ulteriormente i profili, puoi utilizzare TensorBoard.
con il plug-in TPU TensorBoard
sullo stesso computer o su un'altra macchina (consigliato).
Per eseguire TensorBoard su una macchina remota, connettiti mediante SSH e abilita il port forwarding. Ad esempio,
$ ssh -L 6006:localhost:6006 remote server address
o
$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --ssh-flag="-4 -L 6006:localhost:6006"
Sul computer remoto, installa i pacchetti richiesti e avvia TensorBoard
(supponendo che tu abbia profili sulla macchina in ~/profiles/
). Se hai archiviato
nei profili in un'altra directory o in un bucket Cloud Storage, assicurati di
specificare correttamente i percorsi, ad esempio 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
Esecuzione di TensorBoard
Nel browser locale vai a:
http://localhost:6006/
e scegli PROFILE
dal menu a discesa per caricare i tuoi profili.
Consulta la sezione Strumenti TPU per informazioni sulle gli strumenti TensorBoard e su come interpretare l'output.