Generación de perfiles de rendimiento de PyTorch XLA

Descripción general

En esta guía, se explica cómo usar las herramientas de rendimiento de Cloud TPU y la función de análisis automático de métricas con PyTorch. Estas herramientas te ayudan a depurar y optimizar el rendimiento de la carga de trabajo de entrenamiento.

Para obtener más información sobre el rendimiento de Cloud TPU con PyTorch, consulta las siguientes entradas de blog:

Si eres nuevo en PyTorch / XLA, consulta los documentos de API_GUIDE y de solución de problemas de PyTorch. Para Cloud TPU, consulta el documento sobre conceptos.

Generación de perfiles de nodo TPU + PyTorch/XLA

Crea e inicializa los recursos relacionados con Cloud TPU necesarios.

  1. Crea variables para el ID de tu proyecto, tu bucket de Cloud Storage y la zona a fin de usarla en tus recursos TPU.

    export PROJECT_ID=PROJECT_ID
    export BUCKET_NAME=BUCKET_NAME
    export ZONE=ZONE
    
    gcloud --project=$PROJECT_ID compute project-info add-metadata \
    --metadata BUCKET_NAME=$BUCKET_NAME
    
  2. Crear una instancia de VM de Compute Engine Aquí es donde se almacenan todas tus secuencias de comandos y modelos de Python.

    gcloud compute instances create profiler-tutorial-vm \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --machine-type=n1-standard-16 \
      --image-project=ml-images \
      --image-family=torch-xla \
      --boot-disk-size=300GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    
  3. Crea un recurso de TPU.

    gcloud compute tpus create profiler-tutorial-tpu \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --network=default \
     --version=pytorch-1.8 \
     --accelerator-type=v3-8
    
  4. Crea un bucket de Cloud Storage

    Primero, instala la CLI de gsutil si aún no la tienes instalada: instrucciones de instalación.

  5. Usa el comando gsutil mb para crear un bucket de Cloud Storage en el que se almacenen todos los artefactos de generación de perfiles. Reemplaza las variables de región y nombre de bucket con los valores que usarás para el entrenamiento.

    gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
    

    En el ejemplo anterior, se ilustra lo siguiente:

    • REGION es la región en la que creaste la Cloud TPU, por ejemplo, europe-west4.
  6. Crea una cuenta de servicio para el proyecto de Cloud TPU.

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    El comando muestra una cuenta de servicio de Cloud TPU con el siguiente formato:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

    Por ejemplo: service-164006649440@cloud-tpu.iam.gserviceaccount.com

  7. Exporta la cuenta de servicio y otorga permisos de la cuenta de servicio en el bucket de almacenamiento. Reemplaza account-number por el PROJECT_NUMBER que se muestra en el resultado de la creación de la cuenta de servicio.

    export SERVICE_ACCOUNT=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
    gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
    

Configura TensorBoard

  1. ssh al puerto de reenvío de VM 9001 en tu VM al puerto 9001 en tu máquina local. Este puerto se usa para abrir la IU de TensorBoard en tu navegador local.

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
  2. Crea un entorno conda dedicado a la instalación de TensorBoard:

    conda create -y -n tensorboard python=3.6
    conda activate tensorboard
    pip install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
  3. Para probar tu instalación, ejecuta el servidor de TensorBoard en tu VM de Compute Engine y, luego, intenta conectarte al servidor; para ello, visita http://localhost:9001/#profile en tu máquina local:

     # Get bucket name
     BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")
    
     tensorboard --logdir gs://${BUCKET_NAME} --port 9001
    

Cuando visites http://localhost:9001/#profile en tu máquina local, deberías ver un resultado similar al siguiente:

imagen

Si visitas http://localhost:9001, también puedes acceder a la página de perfil anterior. Para ello, selecciona la opción PROFILE del menú desplegable en la esquina superior derecha, junto al botón UPLOAD:

imagen

Genera perfiles del modelo

Para mantener el servidor de TensorBoard activo, desde tu máquina local, inicia una ventana de la terminal nueva y ssh en la VM de GCE nuevamente (esta vez sin usar -L) para la redirección de puertos).

  1. En la ventana de la terminal nueva, vuelve a exportar el ID del proyecto, el entorno del bucket de almacenamiento y las variables de zona, ya que se encuentra en una shell nueva.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh a la VM:

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
    conda activate torch-xla-1.8
    PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
    export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
    --format="value(ipAddress)")
    echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
    export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    
  3. Verifique que las pruebas de integración funcionen de extremo a extremo en su entorno:

    python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
    >
  4. Antes de comenzar el entrenamiento, edita las siguientes líneas en /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

    Cambiar:

    accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    Para:
    accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    Los dos argumentos anteriores en train_mnist provocan de forma artificial grafos dinámicos y recuperaciones de tensores, que se exploran más adelante en la sección de Análisis automático de métricas. Por ahora, como solo estás generando perfiles de la TPU, el siguiente ejemplo se ejecuta con el rendimiento nominal.

  5. Inicia una ejecución de entrenamiento que se use para generar perfiles del servidor:

    PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

Generación de perfiles de TPU (servidor)

Una vez que el entrenamiento se esté ejecutando, visita http://localhost:9001/#profile y captura un perfil mediante las siguientes instrucciones:

imagen

La siguiente página se vuelve a cargar automáticamente:

imagen

Las herramientas compatibles se muestran en el menú desplegable Tools (Herramientas) en el panel izquierdo:

  • Página de descripción general (esto no incluye mostrar el analizador de la canalización de entrada, que es completamente diferente entre TensorFlow/TPU y PyTorch / XLA TPU).
  • Lector de memoria
  • Perfil de operaciones
  • Visor de pod
  • Estadísticas de Tensorflow (estadísticas del nivel de framework; es decir, estadísticas de PyTorch)
  • Lector de seguimiento (requiere el navegador Chrome)

Página de descripción general

En esta página, se muestra una descripción general de un perfil capturado. En este ejemplo, se muestra un tiempo de inactividad muy alto porque se entrena un modelo pequeño en el conjunto de datos MNIST.

Lector de memoria

Muestra la memoria del dispositivo (HBM) usada por tensor y op. Para ver el uso de memoria de un módulo HLO en particular, selecciona el módulo en el menú desplegable Hosts a la izquierda:

imagen

Una vez que visualizas el módulo HLO seleccionado, obtienes una vista integral del cronograma de huellas de HBM de ejecución de módulos. Se ordenan según el tamaño de la asignación, el tamaño de la ejecución del programa y el tamaño del padding.

imagen

Luego, puedes inspeccionar cada una de las asignaciones de búfer si te desplazas sobre ellas. Por ejemplo, para saber qué asignación ocupa la mayoría de los dispositivos HBM:

imagen

En el ejemplo anterior, (1) corresponde a las anotaciones torch_xla.debug.profiler.Trace que agrega el código de usuario. Cuando inspeccionas el código test_profile_mp_mnist.py, corresponde a esta línea:

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...
   

Además, desde el espacio de nombres test_mnist, se puede saber que este módulo HLO corresponde al bucle de evaluación, ya que tiene el administrador de contexto xp.trace('test_mnist').

Perfil de operaciones de XLA

El perfil de operaciones es una herramienta de Cloud TPU que muestra las estadísticas de rendimiento de las operaciones de XLA que se ejecutan durante un período de perfilado. En el perfil de operaciones, se muestra lo siguiente:

  • Qué tan bien tu aplicación usa Cloud TPU, lo que se expresa como el porcentaje de tiempo empleado en las operaciones por categoría y la utilización de los FLOPS de TPU.
  • Las operaciones que más tiempo consumieron. Estas operaciones son objetivos de optimización potenciales.
  • Los detalles de las operaciones individuales, incluidas la forma, el relleno y las expresiones que usan la operación.

Puedes usar el perfil de operaciones para encontrar buenos objetivos de optimización. Por ejemplo, si tu modelo solo logra el 5% de los FLOPS máximos de TPU, puedes usar la herramienta para identificar las operaciones de XLA que tardan más tiempo en ejecutarse y cuántos FLOPS de TPU consumen.

imagen

Descripción de cada una:

  1. Sección Resumen. Muestra el uso de Cloud TPU y proporciona sugerencias de optimización.
  2. Panel de control. Contiene controles que te permiten establecer la cantidad de operaciones que se muestran en la tabla, las operaciones que se muestran y cómo se ordenan.
  3. Tabla de operaciones. Tabla en la que se enumeran las principales categorías de operaciones de TensorFlow asociadas con las operaciones de XLA. Estas operaciones se ordenan según el porcentaje del uso de Cloud TPU.
  4. Tarjetas de detalles de la operación. Muestran detalles acerca de las operaciones, que aparecen cuando te desplazas sobre una operación de la tabla. Se incluye el uso de FLOPS, la expresión en la que se usó la operación y el diseño (ajuste) de la operación.

Visor de pod

Consulta las herramientas de TPU para obtener una descripción completa de la herramienta de visualización de Pods.

Estadísticas del marco de trabajo (estadísticas de Tensorflow/PyTorch)

Las estadísticas del framework proporcionan un desglose detallado de las estadísticas de operación de PyTorch y XRT que se ejecutan en dispositivos TPU y hosts de TPU.

Lector de seguimiento

El lector de seguimiento es una herramienta de análisis del rendimiento de Cloud TPU. El lector de seguimiento usa el visor de creación de perfiles de eventos de seguimiento de Chrome, por lo que requiere el uso del navegador Chrome.

El lector de seguimiento muestra un cronograma que presenta la siguiente información:

  • Duración de las operaciones que ejecutó tu modelo de TensorFlow.
  • La parte del sistema (máquina host o TPU) que ejecutó una operación. En general, para PyTorch/XLA, la máquina host funciona en la compilación y la asignación o desasignación de búfer, mientras que la TPU ejecuta el entrenamiento del modelo real.
  • El lector de seguimiento te permite identificar problemas de rendimiento en tu modelo y, luego, tomar medidas para resolverlos. Con un desglose, puedes identificar las operaciones de PyTorch / XLA que tardan más en ejecutarse.

Ten en cuenta que puedes agregar registros directamente para medir cuánto tiempo tardan en ejecutarse ciertas partes del modelo si agregas anotaciones xp.Trace(NAME). Por ejemplo, en el siguiente seguimiento, se muestra lo siguiente:

imagen

  1. Generado por anotaciones de usuario explícitas en el código del modelo test_profile_mp_mnist.py.
  2. Operaciones de PyTorch ejecutadas (disminución previa).
  3. Nombre del módulo de HLO generado automáticamente por PyTorch / XLA.
  4. Operaciones de XLA que se ejecutan en el dispositivo (fusionada).

Si deseas obtener información más detallada, consulta la documentación genérica de TPU para el visualizador de seguimiento, pero ignora las secciones sobre la canalización de entrada y otras partes específicas de TensorFlow, ya que no son relevantes en este contexto.

Creación de perfiles de cliente de PyTorch / XLA

De forma similar al perfil de la TPU cuando la ejecución del modelo estaba en curso, ahora perfilarás el cliente de PyTorch o XLA durante el entrenamiento. La herramienta de supervisión principal del lado del cliente es el lector de seguimiento.

Debes iniciar el servidor de generación de perfiles en tu secuencia de comandos de entrenamiento. Para ver un ejemplo, consulta, que puedes consultar desde TensorBoard a fin de capturar un seguimiento.

imagen

Para capturar seguimientos de varios procesos, cada proceso puede iniciar servidores de perfiles en diferentes puertos (por ejemplo, si agregas xm.get_ordinal() a un número de puerto base) y, luego, proporcionas una lista de localhost:port concatenado por ','. Tensorboard no admite la visualización de seguimientos de varios procesos a la vez, por lo que verás diferentes menús desplegables de Host para cada proceso.

En el siguiente diagrama, se muestra un seguimiento de muestra:

imagen

De manera similar a cómo se pueden agregar diferentes seguimientos de espacios de nombres para los seguimientos de TPU, puedes usar la misma API a fin de agregarlos a los seguimientos del cliente (xp.Trace(NAME)). Ten en cuenta que este modelo es pequeño y se usa con pequeñas imágenes MNIST, por lo que los tiempos de los pasos serán cortos y no necesariamente uniformes. Como ejercicio, puedes intentar agregar los seguimientos e iniciar un servidor de generador de perfiles similar al de nuestro ejemplo en test_train_mp_imagenet.py --fake_data a fin de obtener registros de ResNet50.

Los seguimientos tienen metadatos adicionales que se pueden inspeccionar. Por ejemplo, los seguimientos de TransferToServer y TransferFromServer, muestran la cantidad exacta de tensores que se envían y reciben, y su tamaño total:

imagen

Para las compilaciones de grafos de XLA, puedes ver el hash de grafo que puede ser útil para diagnosticar problemas:

imagen

Además, en lugar de generar perfiles mediante la IU de TensorBoard, también proporcionamos una API para la generación de perfiles programáticos de la TPU y del cliente desde PyTorch/XLA: xp.trace().

Análisis automático de métricas

En esta sección, verás cómo usar el modo de depuración para detectar problemas de rendimiento, como los siguientes:

  • Compilaciones dinámicas o gráficos dinámicos
  • Compilación de gráficos muy lenta
  • Ejecución de grafo muy lenta
  • XLA frecuente: Transferencias de CPU
  • HBM repetido del dispositivo para intercambiar RAM
  • Desfragmentación repetida de HBM
  • Operaciones de aten:: sin disminución

Antes de comenzar el entrenamiento, revierte las siguientes líneas en /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

Cambiar:

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   
Para:
   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

Estos cambios provocan compilaciones y recuperaciones de tensores de forma artificial. dynamic_graph=True cambia, de manera artificial, el tamaño del lote para cada paso, lo que provoca que los grafos reducidos de XLA sean diferentes en cada paso y recompilación. fetch_often=True inserta loss.item() llamadas en cada paso, lo que genera la recuperación de valores del tensor del dispositivo en cada paso, lo que ralentiza el rendimiento.

Ejecuta una secuencia de comandos de entrenamiento de ejemplo:

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1
   

Durante la depuración, se recomienda ejecutar con --num_cores=1, ya que simplifica el proceso de depuración. Algunos de los resultados de muestra se ven de la siguiente manera:

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

Las líneas con el prefijo pt-xla-profiler corresponden al resultado del análisis de métricas automáticas. En este ejemplo, puedes ver que TransferFromServerTime se detectó con demasiada frecuencia, una vez por paso. Esto se debe al bucle de entrenamiento que recupera el valor de loss.item() en cada paso. Además, es posible que veas una advertencia CompileTime demasiadas veces cuando tu modelo debe volver a compilarse varias veces debido a las formas dinámicas del gráfico. El siguiente fragmento de código provoca este tipo de problema: test_profile_mp_mnist.py.

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

Luego, puedes presionar Ctrl^C fuera de la secuencia de comandos de entrenamiento y, al final, deberías ver un resumen de las operaciones no disminuyedas. Ten en cuenta que aten::_local_scalar_dense es una operación especial que corresponde a la recuperación de los tensores XLA en el contexto de la CPU.

En este informe, verás que hay dos lugares principales donde se llama a la operación aten::_local_scalar_dense; ambos corresponden al código fuente de loss.item():

  • test/test_profile_mp_mnist.py:158
  • test/test_profile_mp_mnist.py:61
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

Ahora, ejecuta un análisis de métricas automáticas en la siguiente secuencia de comandos, que contiene una operación no menor (a partir de nuestra versión 1.8), _ctc_loss:

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

Cuando ejecutas la secuencia de comandos anterior con PT_XLA_DEBUG=1, el resultado debería ser similar al siguiente:

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

El analizador de métricas automáticas muestra que la 11.a línea de STDIN está provocando esta operación no disminuyeda (es decir, la línea con torch.nn.CTCLoss()). Actualmente, la operación ctc_loss no se redujo, por lo que ves el informe anterior. También puedes ver algunas advertencias para TransferFromServerTime, ya que los tensores se encuentran inicialmente en XLA:TPU antes de la ejecución, pero como la operación no se reduce, primero debes transferir los tensores XLA a la CPU, ejecutar la operación aten:: en la CPU y transferirla de nuevo.

Si, en cambio, deseas escribir el resultado de pt-xla-profiler en un archivo, configura PT_XLA_DEBUG=1 y PT_XLA_DEBUG_FILE=$PATH_TO_FILE.

Limpieza

Sal de tu VM y, luego, borra la TPU, la VM y el bucket de Cloud Storage mediante la ejecución de los siguientes comandos:

(vm)$ exit
gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}
gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async
gsutil rm -fr gs://${BUCKET_NAME}

Generación de perfiles de VM de TPU y PyTorch/XLA

Usa esta sección para generar un perfil de PyTorch/XLA con la arquitectura de VM de TPU.

Exportar variables de entorno

  1. Crea variables para el ID de tu proyecto y la zona que se usará en tus recursos TPU.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

Crear una Cloud TPU

Consulta la Guía del usuario de VM de TPU y, después de la configuración, crea una VM de TPU v3-8, que viene contorch ,torch_xla, torchvision ytensorboard preinstaladas.

  1. Crea un recurso de TPU.

    gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --version=v2-alpha \
     --accelerator-type=v3-8
    

Inicio del servidor de TensorBoard

  1. Establece una conexión SSH a la VM, instala tensorboard-plugin-profile y, luego, inicia un servidor de TensorBoard.

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
      pip3 install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
      tensorboard --logdir ./tensorboard --port 9001
    

Cuando veas el resultado de TensorBoard en http://localhost:9001 en tu máquina local, deberías ver algo como lo siguiente:

imagen

Si ves el resultado de TensorBoard en http://localhost:9001, también puedes acceder a la página de perfil anterior. Para ello, selecciona la opción PROFILE de la lista desplegable que se encuentra en la esquina superior derecha junto al botón UPLOAD:

imagen

Genera perfiles del modelo

En una ventana de terminal nueva del entorno de desarrollo, exporta las mismas variables de entorno que se mencionaron anteriormente y ssh en la VM de TPU:

  1. En la ventana de la terminal nueva, vuelve a exportar el ID del proyecto y las variables de zona, ya que se encuentran en una shell nueva.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh a la VM:

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
  3. Clona el repositorio de PyTorch/XLA y ejecuta la prueba e2e:

      git clone -b r1.8 https://github.com/pytorch/xla
      export XRT_TPU_CONFIG="localservice;0;localhost:51011"
      python3 xla/test/test_profiler.py  # takes <1 min
    >
  4. Antes de comenzar el entrenamiento, edita las siguientes líneas en xla/test/test_profile_mp_mnist.py:

    Cambiar:

        accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    Para:
        accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    Los dos argumentos anteriores en train_mnist provocan de forma artificial grafos dinámicos y recuperaciones de tensores, que se exploran más adelante en la sección de Análisis automático de métricas. Por ahora, como solo estás generando perfiles de la TPU, el siguiente ejemplo se ejecuta con el rendimiento nominal.

  5. Inicia una ejecución de entrenamiento:

     XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU + generación de perfiles de cliente

Una vez que el entrenamiento se esté ejecutando, consulta la salida de TensorBoard en http://localhost:9001 y captura un perfil con las siguientes instrucciones:

imagen

Deberías volver a cargar la siguiente página:

imagen

Actualmente, en la configuración de la VM de TPU, solo se selecciona la herramienta de visor de seguimiento, por lo que en el menú desplegable Herramientas, debes seleccionar trace_viewer e inspeccionar los seguimientos. Puedes ver que, en la configuración de la VM de TPU, ves los registros del lado del cliente y del dispositivo de TPU en una vista completa:

imagen

Limpieza

  1. Sal de tu VM y, luego, borra la TPU, la VM y el bucket de Cloud Storage mediante la ejecución de los siguientes comandos:

    (vm)$ exit
    

Borra la VM de TPU que creaste:

  1. Borra tus recursos de Cloud TPU y Compute Engine.

    $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} --zone=${ZONE}
    
  2. Ejecuta el siguiente comando para verificar que se hayan borrado los recursos. La eliminación puede tardar varios minutos. Una respuesta como la que se muestra a continuación indica que tus instancias se borraron de forma correcta.

    $ gcloud alpha compute tpus tpu-vm list --project ${PROJECT_ID} --zone=${ZONE}
    
    Listed 0 items.