Genera perfiles de las cargas de trabajo de XLA de PyTorch

La generación de perfiles es una forma de analizar y mejorar el rendimiento de los modelos. Si bien hay mucho más que eso; a veces, es útil pensar en la generación de perfiles como tiempo operaciones 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 sobre cómo generar perfiles para tu código para el entrenamiento o la inferencia. Para más información sobre cómo analizar los perfiles generados, consulta consulta las siguientes guías.

• Depuración del rendimiento de XLA de PyTorch en VMs de TPU (parte 1)
• Depuración del rendimiento de XLA de PyTorch en VMs de TPU (parte 2)
• Depuración del 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

   Es el nombre que quieres usar para la Cloud TPU.

   zone

   Es la zona en la que deseas crear la Cloud TPU.

   project ID

   Es 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 deseas crear. Para 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 el variable de entorno, pero también puedes usar una de la lista parámetros de configuración admitidos.

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 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 TPU. Cada trabajador genera perfiles solo de un dispositivo TPU de los 4 dispositivos disponibles. Crear perfiles de los 4 dispositivos de TPU para cada trabajador pronto se admitirán.

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

Cómo iniciar el servidor de perfiles

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

Nota: Puedes consultar este secuencia de comandos de ejemplo que incluya 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 Lightning módulo, en lugar de iniciar un servidor, pasa profiler="xla" a trainer, y, luego, iniciará automáticamente el servidor en el puerto 9012.

Ahora puedes capturar perfiles como se describe en la siguiente sección. El guion perfila todo lo que sucede en un dispositivo de TPU.

Agregar seguimientos

Si también deseas para generar perfiles de las operaciones en la máquina anfitrión, puedes agregar xp.StepTrace o xp.Trace en tu código. Con estas funciones, se hace un seguimiento del código de Python en la máquina anfitrión. (Esto se puede considerar como una medición del tiempo que lleva ejecutar el comando en el host (CPU) antes de pasar el “gráfico” al dispositivo de TPU. Por lo tanto, es útil sobre todo para analizar la sobrecarga de seguimiento). Puedes agrega esto al bucle de entrenamiento, en el que el código procesa 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. El código falla en ese caso 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 se agrega automáticamente xm.mark_step() cuando sales del permiso. Además, podrías recibir este error cuando Usas torchdynamo porque torchdynamo puede invocar internamente a xm.mark_step. La solución alternativa actual es no usar xp.Trace.

Si usas Iluminación, 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 puedes insertarlos dentro del ciclo de entrenamiento.

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

Captura manual

La secuencia de comandos capture_profile.py del repositorio de XLA de PyTorch te permite capturar un perfil rápidamente. Puedes hacerlo copiando el archivo de perfil de captura directamente a la VM de TPU. El siguiente comando lo 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 el siguiente comando 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 Logging ~/profiles/ a tu ubicación y nombre preferidos. También es posible guardar directamente en el bucket de Cloud Storage. Para hacer eso, configura logdir será gs://your_bucket_name/.

Captura programática

En lugar de capturar el perfil manualmente al activar una secuencia de comandos, puedes configura tu secuencia de comandos de entrenamiento para activar automáticamente un perfil con torch_xla.debug.profiler.trace_detached API en tu secuencia de comandos de entrenamiento.

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

Por ejemplo, para capturar automáticamente un perfil en un ciclo de entrenamiento y paso específicos, puedes configurar tu secuencia de comandos de entrenamiento para que consuma PROFILE_STEP, PROFILE_EPOCH y PROFILE_LOGDIR entorno variables:

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 especificado por el Variable de entorno PROFILE_LOGDIR.

Análisis en TensorBoard

Para analizar aún más los perfiles, 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 este con 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 la máquina remota, instala los paquetes necesarios y, luego, inicia TensorBoard. (suponiendo que tienes perfiles en esa máquina con ~/profiles/). Si almacenaste los perfiles en otro directorio o bucket de Cloud Storage, asegúrate de Especifica rutas de acceso correctamente, por ejemplo, gs://your_bucket_name/profiles.

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

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

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

Nota: Si obtienes "Duplicate plugins" podría deberse a que hay múltiples paquetes con TensorBoard. Desinstala todos esos paquetes y ejecuta la 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

Ejecución de 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 más información sobre la Herramientas de TensorBoard y cómo interpretar el resultado.