Profila i carichi di lavoro PyTorch XLA

Il profiling è un modo per analizzare e migliorare le prestazioni dei modelli. Anche se c'è molto di più, a volte può essere utile pensare al profiling come al temporizzazione delle operazioni e delle parti di codice che vengono eseguite sia sui dispositivi (TPU) sia sugli host (CPU). Questa guida fornisce una rapida panoramica su come eseguire il profiling del codice per l'addestramento o l'inferenza. Per maggiori informazioni su come analizzare i profili generati, consulta le seguenti guide.

• Debug delle prestazioni di PyTorch XLA sulle VM TPU - parte 1
• Debug delle prestazioni di PyTorch XLA su VM TPU - parte 2
• Debug delle prestazioni di PyTorch XLA su VM TPU - parte 3

Inizia

Crea una TPU

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

   Esportare le descrizioni delle 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 ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.

   version

   La versione del runtime Cloud TPU. Nella variabile esportata viene mostrato un valore predefinito, ma puoi anche utilizzarne uno dall'elenco delle configurazioni supportate.

2. Lancia 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 della VM TPU utilizzando il comando gcloud scp. Ad esempio:

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

Profilazione

Nota: una versione v4-8 include 4 dispositivi TPU. Ogni utente crea il profilo di un solo dispositivo TPU tra i 4 disponibili. A breve sarà supportato il profiling di tutti e quattro i dispositivi TPU per ogni utente.

Un profilo può essere acquisito manualmente tramite capture_profile.py o in modo programmatico dallo script di addestramento utilizzando le API torch_xla.debug.profiler.

Avvio del server di profili

Per acquisire un profilo, è necessario che un server di profili sia in esecuzione nello script di addestramento. Avvia un server con un numero di porta a tua scelta, ad esempio9012 come mostrato nel seguente comando.

Nota: puoi fare riferimento a questo script di esempio che include i seguenti passaggi.

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

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

Nota: se utilizzi il modulo Lightning, anziché avviare un server, passa profiler="xla" a trainer, che avvierà automaticamente il server sulla porta 9012.

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

Aggiungere tracce

Se vuoi anche eseguire il profiling delle operazioni sulla macchina host, puoi aggiungere xp.StepTrace o xp.Trace nel codice. Queste funzioni tracciano il codice Python sulla macchina ospitante. Puoi considerare questa misurazione come il tempo necessario per eseguire il codice Python sull'host (CPU) prima di passare il "grafo" al dispositivo TPU. Pertanto, è utile soprattutto per analizzare l'overhead del monitoraggio. 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 racchiudere singole parti del codice con

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

Importante: evita di applicare il wrapping a xm.mark_step() con xp.Trace. In questo caso il codice uscirà con un messaggio di errore simile al seguente: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace non avrà questo problema perché aggiunge automaticamente xm.mark_step() quando esce dall'ambito. Inoltre, potresti visualizzare questo errore quando utilizzi torchdynamo perché torchdynamo potrebbe invocare internamente xm.mark_step. La soluzione alternativa attuale è non utilizzare xp.Trace.

Se utilizzi Lighting, puoi saltare l'aggiunta delle tracce perché viene eseguita automaticamente in alcune parti del codice. Tuttavia, se vuoi aggiungere altre tracce, puoi inserirle all'interno del loop di addestramento.

Potrai acquisire l'attività del dispositivo dopo la compilazione iniziale. Attendi fino a quando il modello non avvia i passaggi di addestramento o di inferenza.

Acquisizione manuale

Lo script capture_profile.py del repository Pytorch XLA consente di acquisire rapidamente un profilo. A tal fine, puoi copiare 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"

Durante l'addestramento, esegui i seguenti 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 in logdir. Puoi modificare la directory di logging~/profiles/ impostando la posizione e il nome che preferisci. È anche possibile 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 all'interno dello script di addestramento.

Nota: l'API trace_detached è disponibile solo nelle build notturne o dopo la release 2.2. Per eseguire l'acquisizione tramite programmazione nelle versioni precedenti di torch_xla, puoi copiare il corpo del metodo per creare direttamente il thread di acquisizione in background.

Ad esempio, per acquisire automaticamente un profilo in un'epoca e un passaggio specifici, puoi configurare lo script di addestramento in modo da utilizzare 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 (opzione consigliata).

Per eseguire TensorBoard su una macchina remota, connettiti utilizzando SSH e attiva 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 abbia profili sulla macchina in ~/profiles/). Se hai archiviato i profili in un'altra directory o in un bucket Cloud Storage, assicurati di specificare i percorsi correttamente, ad esempio gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile 

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

Nota: è importante avere file .xplane nella stessa struttura nidificata che viene generata quando acquisisci un profilo. Ad esempio: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Nota: se ricevi l'errore "Plugin duplicati", è possibile che ci siano più pacchetti con TensorBoard. Disinstalla tutti questi pacchetti ed esegui di nuovo i comandi pip install. Ad esempio, potresti dover eseguire:

(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

Eseguire TensorBoard

Nel browser locale, vai alla pagina: http://localhost:6006/ e scegli PROFILE dal menu a discesa per caricare i tuoi profili.

Consulta gli strumenti TPU per informazioni sugli strumenti di TensorBoard e su come interpretare l'output.