Profila i carichi di lavoro XLA di PyTorch

La profilazione è un modo per analizzare e migliorare le prestazioni dei modelli. Sebbene c'è molto altro da fare, a volte è utile pensare alla profilazione come a operazioni di temporizzazione e parti di codice eseguite su dispositivi (TPU) e host (CPU). Questa guida fornisce una panoramica rapida su come profilare il codice per l'addestramento o l'inferenza. Per ulteriori informazioni su come analizzare i profili generati, consulta le seguenti guide.

Inizia

Crea una TPU

  1. Esporta 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 la tua Cloud TPU.
    zone
    La zona in cui prevedi di creare la Cloud TPU.
    project ID
    L'ID progetto che utilizzi 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 maggiori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.
    version
    La versione del runtime di Cloud TPU. Un valore predefinito viene visualizzato nella variabile esportata, ma puoi anche utilizzarne una dall'elenco delle configurazioni supportate.

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

  3. Sposta il codice nella home directory sulla VM TPU utilizzando il comando gcloud scp. Ad esempio:

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

Profilazione

Un profilo può essere acquisito manualmente tramite capture_profile.py o a livello di programmazione dallo script di addestramento utilizzando le API torch_xla.debug.profiler.

Avvio del server di profili

Per acquisire un profilo, un server del profilo deve essere in esecuzione all'interno dello script di addestramento. Avvia un server con un numero di porta a tua scelta, ad esempio 9012 come mostrato nel comando seguente.

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

Il server può essere avviato all'inizio della funzione main.

Ora puoi acquisire i profili come descritto nella sezione seguente. Lo script profila tutto ciò che accade su un dispositivo TPU.

Aggiunta di tracce

Se vuoi profilare le operazioni anche sulla macchina host, puoi aggiungere xp.StepTrace o xp.Trace nel codice. Queste funzioni tracciano il codice Python sulla macchina host. (Può essere considerato come la misurazione del tempo necessario per eseguire il codice Python sull'host (CPU) prima di passare il "grafico" al dispositivo TPU. pertanto è utile soprattutto per analizzare i costi generali). 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):
       ...

oppure aggregare singole parti del codice

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

Se usi Illuminazione, puoi saltare l'aggiunta di tracce poiché questa operazione viene eseguita automaticamente in alcune parti del codice. Tuttavia, se vuoi aggiungere altre tracce, ti invitiamo a inserirle nel loop di addestramento.

Potrai acquisire l'attività del dispositivo dopo la compilazione iniziale; attendi che il modello inizi i passaggi di addestramento o di inferenza.

Acquisizione manuale

Lo script capture_profile.py dal repository XLA di Pytorch consente di acquisire rapidamente un profilo. Per farlo, copia il file del profilo di acquisizione direttamente nella VM TPU. Il seguente comando 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"

Mentre l'addestramento è in esecuzione, esegui questi passaggi 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 i file .xplane.pb nella cartella logdir. Puoi modificare la directory di logging ~/profiles/ con il nome e la posizione che preferisci. Puoi anche salvare direttamente nel bucket Cloud Storage. Per farlo, imposta logdir su gs://your_bucket_name/.

acquisizione programmatica

Anziché acquisire il profilo manualmente attivando uno script, puoi configurare lo script di addestramento in modo che attivi automaticamente un profilo utilizzando l'API torch_xla.debug.profiler.trace_detached nello script di addestramento.

Ad esempio, per acquisire automaticamente un profilo in un'epoca e un passaggio specifici, puoi configurare lo script di addestramento in modo che utilizzi le variabili di ambiente PROFILE_STEP, PROFILE_EPOCH e 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)
      ...

I file .xplane.pb verranno salvati nella directory specificata dalla variabile di ambiente PROFILE_LOGDIR.

Analisi in TensorBoard

Per analizzare ulteriormente i profili, puoi utilizzare TensorBoard con il plug-in TPU TensorBoard
sulla stessa macchina o su un'altra macchina (consigliato).

Per eseguire TensorBoard su una macchina remota, connettiti alla macchina utilizzando 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"

Sulla macchina remota, installa i pacchetti richiesti e avvia TensorBoard (supponendo che tu disponga di profili su quella macchina in ~/profiles/). Se hai archiviato i profili in un'altra directory o in un altro bucket Cloud Storage, assicurati di specificare correttamente i percorsi, ad esempio 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

TensorBoard in esecuzione

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 sugli strumenti di TensorBoard e su come interpretare l'output.

immagine

immagine