Se usó la API de Cloud Translation para traducir esta página.

Genera perfiles de cargas de trabajo de PyTorch XLA

La generación de perfiles es una forma de analizar y mejorar el rendimiento de los modelos. Aunque hay mucho más en el asunto, a veces es útil pensar en la generación de perfiles como operaciones de tiempo 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 la capacitación o la inferencia. Si deseas obtener más información para analizar los perfiles generados, consulta las siguientes guías.

Crear una TPU

Exporta las variables de entorno:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v4-8
export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

Descripciones de las variables de entorno

Variable	Descripción
`PROJECT_ID`	El Google Cloud ID de tu proyecto. Usa un proyecto existente o crea uno nuevo.
`TPU_NAME`	El nombre de la TPU.
`ZONE`	Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU.
`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.
`RUNTIME_VERSION`	La versión de software de Cloud TPU

Inicia los recursos de TPU

$ gcloud compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --version ${RUNTIME_VERSION} \
    --project ${PROJECT_ID}

Usa el siguiente comando para instalar torch_xla en todas las VMs de TPU en una porción de TPU. También deberás instalar cualquier otra dependencia que requiera tu secuencia de comandos de entrenamiento.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --worker=all \
    --command="pip install torch==2.6.0 torch_xla[tpu]==2.6.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html"

Mueve tu código a tu 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}:directory/target-file --zone ${ZONE}

Perfil

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.

Inicia 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 el número de puerto que elijas, por ejemplo, 9012, como se muestra en el siguiente comando.

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

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

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

Cómo agregar seguimientos

Si también deseas generar perfiles de operaciones en la máquina anfitrión, puedes agregar xp.StepTrace o xp.Trace en tu código. Estas funciones rastrean el código de Python en la máquina anfitrión. (Puedes considerar esto como medir cuánto tiempo lleva ejecutar el código de Python en el host [CPU] antes de pasar el "grafo" al dispositivo TPU. Por lo tanto, es más útil para analizar la sobrecarga de seguimiento. Puedes añadir esto dentro del 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 = ...

Si usas Lighting, puedes omitir agregar registros, ya que se hace automáticamente en algunas partes del código. Sin embargo, si deseas agregar registros 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 inicie sus pasos de entrenamiento o inferencia.

Captura manual

La secuencia de comandos capture_profile.py del repositorio de XLA de Pytorch permite capturar un perfil rápidamente. Para ello, copia el archivo de perfil de captura directamente en tu 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 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/ a la ubicación y el nombre que prefieras. También es posible guardar directamente en el bucket de Cloud Storage. Para ello, configura logdir como gs://your_bucket_name/.

Captura programática

En lugar de capturar el perfil de forma manual activando una secuencia de comandos, puedes configurar tu secuencia de comandos de entrenamiento para que active automáticamente un perfil con la API de torch_xla.debug.profiler.trace_detached en tu secuencia de comandos de entrenamiento.

A modo de ejemplo, para capturar automáticamente un perfil en una época y un paso específicos, puedes configurar tu 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)
      ...

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

Análisis en TensorBoard

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

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

$ ssh -L 6006:localhost:6006 remote server address

$ 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 tengas 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 correctamente, por ejemplo, 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

Ejecuta TensorBoard

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

Consulta Perfila tu modelo en VMs de Cloud TPU para obtener información sobre las herramientas de TensorBoard y cómo interpretar los resultados.

Página de captura de perfiles de TensorBoard