Genera perfiles de las cargas de trabajo de PyTorch XLA

La creación de perfiles es una forma de analizar y mejorar el rendimiento de los modelos. Aunque hay mucho más al respecto, a veces es útil pensar en la generación de perfiles como operaciones de sincronización y partes del código que se ejecutan en dispositivos (TPU) y hosts (CPU). En esta guía, se proporciona una descripción general rápida de cómo generar perfiles de tu código para entrenamiento o inferencia. Si deseas obtener más información para analizar los perfiles generados, consulta las siguientes guías.

• Depuración de rendimiento de PyTorch XLA en VMs de TPU: parte 1
• Depuración de rendimiento de PyTorch XLA en VMs de TPU: parte 2
• Depuración de rendimiento de PyTorch XLA en VMs de TPU: parte 3

Comenzar

Crea una TPU

1. Exporta las variables de entorno:

   $ 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
   

   Exportar descripciones de variables

   TPU name

   El nombre que quieres usar para tu Cloud TPU.

   zone

   La zona en la que planeas crear tu Cloud TPU.

   project ID

   El ID del proyecto que usas para entrenar tu modelo y generar perfiles de él.

   accelerator-type

   El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que quieres crear. Si quieres obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.

   version

   La versión del entorno de ejecución de Cloud TPU. Se muestra un valor predeterminado en la variable exportada, pero también puedes usar uno de la lista de configuraciones compatibles.

2. Inicia los recursos de 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. Mueve el código al directorio principal en la VM de TPU con el comando gcloud scp. Por ejemplo:

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

Generación de perfiles

Nota: Una v4-8 incluye 4 dispositivos de TPU. Cada trabajador genera perfiles de solo un dispositivo de TPU de los 4 dispositivos disponibles. Pronto se admitirá la creación de perfiles de los 4 dispositivos de TPU para cada trabajador.

Se puede capturar un perfil de forma manual a través de capture_profile.py o de manera programática desde la secuencia de comandos de entrenamiento con las APIs de torch_xla.debug.profiler.

Cómo iniciar el servidor de perfiles

Para capturar un perfil, se debe ejecutar un servidor de perfiles dentro de la secuencia de comandos de entrenamiento. Inicia un servidor con un número de puerto que elijas, por ejemplo, 9012, como se muestra en el siguiente comando.

Nota: Puedes consultar esta secuencia de comandos de ejemplo, que incluye los siguientes pasos.

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

El servidor se puede iniciar al comienzo de la función main.

Nota: Si usas el módulo Lightning, en lugar de iniciar un servidor, pasa profiler="xla" a trainer, y este iniciará automáticamente el servidor en el puerto 9012.

Ahora puedes capturar perfiles como se describe en la siguiente sección. La secuencia de comandos genera perfiles de todo lo que sucede en un dispositivo de TPU.

Agrega seguimientos

Si también deseas generar perfiles de las operaciones en la máquina anfitrión, puedes agregar xp.StepTrace o xp.Trace a tu código. Estas funciones realizan un seguimiento del código de Python en la máquina anfitrión. (Se puede considerar que se trata de medir el tiempo que lleva ejecutar el código de Python en el host (CPU) antes de pasar el "grafo" al dispositivo de TPU. por lo que es útil principalmente para analizar la sobrecarga de seguimiento). Puedes agregarlo dentro del bucle de entrenamiento en el que el código procesa los lotes de datos, por ejemplo,

for step, batch in enumerate(train_dataloader):
    with xp.StepTrace('Training_step', step_num=step):
       ...

o unir partes individuales del código con

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

Importante: Evita unir xm.mark_step() con xp.Trace. En ese caso, el código fallará con un mensaje de error similar a "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace no tendrá este problema, ya que agrega xm.mark_step() automáticamente cuando se sale del alcance. Además, es posible que recibas este error cuando uses torchdynamo, ya que torchdynamo puede invocar internamente a xm.mark_step. La solución alternativa actual es no usar xp.Trace.

Si usas Lighting, puedes omitir la adición de seguimientos, ya que se hace automáticamente en algunas partes del código. Sin embargo, si deseas agregar seguimientos adicionales, puedes insertarlos dentro del bucle de entrenamiento.

Podrás capturar la actividad del dispositivo después de la compilación inicial. Espera hasta que el modelo comience sus pasos de entrenamiento o inferencia.

Captura manual

La secuencia de comandos capture_profile.py del repositorio de Pytorch XLA permite capturar un perfil con rapidez. Para ello, copia el archivo de perfil de captura directamente en la VM de TPU. Con el siguiente comando, se copia en el directorio principal.

$ 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"

Mientras se ejecuta el entrenamiento, ejecuta lo siguiente para capturar un perfil:

$ 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"

Este comando guarda archivos .xplane.pb en logdir. Puedes cambiar el directorio de registro ~/profiles/ según tu ubicación y nombre preferidos. También se puede guardar directamente en el bucket de Cloud Storage. Para hacerlo, configura logdir como gs://your_bucket_name/.

Captura programática

En lugar de capturar el perfil de forma manual con la activación de una secuencia de comandos, puedes configurar la secuencia de comandos de entrenamiento para activar un perfil de forma automática mediante la API torch_xla.debug.profiler.trace_detached en la secuencia de comandos de entrenamiento.

Nota: La API de trace_detached solo está disponible en compilaciones nocturnas o después de la versión 2.2. Para realizar la captura de manera programática en versiones anteriores de torch_xla, puedes copiar el cuerpo del método a fin de crear directamente el subproceso de captura en segundo plano.

Por ejemplo, para capturar automáticamente un perfil en un ciclo de entrenamiento y un paso específicos, puedes configurar la secuencia de comandos de entrenamiento para que consuma las variables de entorno PROFILE_STEP, PROFILE_EPOCH y 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)
      ...

Se guardarán los archivos .xplane.pb en el directorio que especifique la variable de entorno PROFILE_LOGDIR.

Análisis en TensorBoard

Para analizar perfiles en profundidad, puedes usar TensorBoard con el complemento de TPU TensorBoard
en la misma máquina o en otra (recomendado).

Para ejecutar TensorBoard en una máquina remota, conéctate a ella mediante SSH y habilita la redirección de puertos. Por ejemplo,

$ 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"

En tu máquina remota, instala los paquetes necesarios y, luego, inicia TensorBoard (suponiendo que tienes perfiles en esa máquina en ~/profiles/). Si almacenaste los perfiles en otro directorio o bucket de Cloud Storage, asegúrate de especificar las rutas de acceso de forma correcta, por ejemplo, gs://your_bucket_name/profiles.

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

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

Nota: Es importante tener archivos .xplane en la misma estructura anidada que se genera cuando capturas un perfil. Por ejemplo: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Nota: Si se muestra el error “Duplicate plugins”, puede deberse a que hay varios paquetes con TensorBoard. Desinstala todos esos paquetes y vuelve a ejecutar los comandos pip install. Por ejemplo, es posible que debas ejecutar lo siguiente:

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

Ejecutando TensorBoard

En tu navegador local, ve a http://localhost:6006/ y elige PROFILE en el menú desplegable para cargar tus perfiles.

Consulta las herramientas de TPU para obtener información sobre las herramientas de TensorBoard y cómo interpretar el resultado.