Generación de perfiles de rendimiento de PyTorch

Descripción general

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

Conceptos

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

Perfil de TPU + perfil de PyTorch/XLA

Crea y, luego, inicializa los recursos necesarios relacionados con Cloud TPU.

  1. Crea variables para tu ID del proyecto, tu depósito de Cloud Storage y la zona que usarás para 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. Crea 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 gsutil si aún no la instalaste: instrucciones de instalación.

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

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

    Donde:

    • REGION es la región en la que creaste el 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-1640064440440@cloud-tpu.iam.gserviceaccount.com

  7. Exporta la cuenta de servicio y otorga permisos de la cuenta de servicio en el depósito de almacenamiento. Reemplaza account-number por el PROJECT_NUMBER que se muestra en el resultado de la creación de cuentas 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 9001 de reenvío de VM en tu VM al puerto 9001 de 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 visitas http://localhost:9001/#profile en tu máquina local, deberías ver algo como esto:

image

Si visitas http://localhost:9001, también puedes acceder a la página de perfil anterior si seleccionas la opción PROFILE en el menú desplegable de la esquina superior derecha junto al botón SUBIR. :

image

Genera el perfil del modelo

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

  1. En la ventana de la terminal nueva, exporta el ID del proyecto, el entorno del depósito de almacenamiento y las variables de zona de nuevo, ya que se encuentra en una shell nueva.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh en 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. Verifica que las pruebas de integración funcionen en tu entorno de manera integral:

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

    Cambio:

    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, causar artificialmente los grafos dinámicos y las recuperaciones de tensores, que se exploran más adelante en la sección de análisis de métricas automáticas. Por ahora, dado que solo estás generando el perfil de la TPU, el siguiente ejemplo se ejecuta con un rendimiento nominal.

  5. Inicia una ejecución de entrenamiento que se usa para la generación de 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
    

Perfilado de TPU (servidor)

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

image

La página siguiente se vuelve a cargar de forma automática:

image

Las herramientas compatibles se muestran en el menú desplegable Herramientas (Tools) 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 de nivel del marco de trabajo, estadísticas de PyTorch)
  • Visualizador 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 el tiempo de inactividad muy alto porque entrena un modelo pequeño en el conjunto de datos MNIST.

Lector de memoria

Muestra la memoria del dispositivo (HBM) que se usa por tensor y la operación de PyTorch de HLO. El lector de memoria captura una vista de la memoria por módulo de HLO, por lo que habrá algunos módulos, como los gráficos de asignación de datos del dispositivo (entradas de lote enviados y etiquetas). Para ver el uso de la memoria de un módulo HLO en particular, selecciona el módulo del menú desplegable Hosts a la izquierda:

image

Una vez que visualices el módulo de HLO seleccionado, obtendrás una vista integral del cronograma de huellas digitales de HBM de ejecución del módulo. Se ordenan por tamaño de asignación, tamaño de ejecución del programa y tamaño de relleno.

image

Luego, cada una de las asignaciones de búfer se puede inspeccionar más. Por ejemplo, para saber qué asignación ocupa la mayor parte del dispositivo HBM, haz lo siguiente:

image

En el ejemplo anterior, (1) corresponde a las anotaciones torch_xla.debug.profiler.Trace agregadas por el código de usuario. Inspeccionar el código de test_profile_mp_mnist.py que 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)
   ...
   

También del espacio de nombres de las operaciones test_mnist, puedes indicarle que este módulo de 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 ejecutadas 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 potenciales de optimización.
  • 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.

image

Descripción de cada una:

  1. Descripción general. Muestra el uso de Cloud TPU y proporciona sugerencias de optimización.
  2. Panel de control. Contiene controles que te permiten configurar la cantidad de operaciones que se muestran en la tabla, qué operaciones se muestran y cómo se ordenan.
  3. Tabla de operaciones. Una tabla que enumera las principales categorías de operación 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 las operaciones Muestran detalles acerca de las operaciones, que aparecen cuando colocas el cursor sobre una operación de la tabla. Estos incluyen 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 del lector de pods.

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

Las estadísticas del marco de trabajo proporcionan un desglose de las operaciones de PyTorch y XRT detalladas en ejecución en dispositivos de 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 lector de perfilado de eventos de seguimiento de Chrome, por lo que se 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 el caso de PyTorch / XLA que, por lo general, la máquina host funciona principalmente en la compilación y la asignación/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 tomar medidas para resolverlos. Si desglosas los datos, puedes identificar las operaciones de PyTorch / XLA que más tardan en ejecutarse.

Ten en cuenta que puedes agregar seguimientos directamente para medir cuánto tiempo demoran algunas partes en ejecutar el modelo si agregas anotaciones xp.Trace(NAME). Por ejemplo, se muestra el siguiente seguimiento:

image

  1. Lo genera las anotaciones de usuario explícitas en el código del modelo de test_profile_mp_mnist.py.
  2. Operaciones de PyTorch que se ejecutaron (mejora)
  3. Nombre del módulo HLO generado por PyTorch o XLA
  4. Operaciones de XLA ejecutadas en el dispositivo (fusionadas)

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

Creación de perfiles de cliente de PyTorch / XLA

De manera similar a cuando perfilar el lado de la TPU mientras la ejecución del modelo estaba en curso, ahora vas a perfilar el cliente PyTorch / XLA durante el entrenamiento. La herramienta principal de supervisión que se usa en el lado del cliente es el Lector de seguimiento.

Debes iniciar el servidor de creación de perfiles en tu secuencia de comandos de entrenamiento. Por ejemplo, ver, que puedes consultar desde TensorBoard para capturar un seguimiento.

image

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

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

image

Así como se pueden agregar diferentes seguimientos de espacio 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, como este modelo es pequeño y se usa con imágenes MNIST pequeñas, los tiempos de los pasos serán cortos y no necesariamente uniformes. Como ejercicio, puedes intentar agregar los seguimientos y, luego, iniciar un servidor de generador de perfiles similar al de nuestro ejemplo a test_train_mp_imagenet.py --fake_data para obtener seguimientos 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 se reciben, y su tamaño total:

image

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

image

Además, en lugar de generar perfiles a través de la IU de TensorBoard, también proporcionamos una API para generar perfiles de manera programática de la TPU y el cliente desde PyTorch / XLA: xp.trace().

Análisis de métricas automáticas

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

  • Gráficos dinámicos y compilaciones continuas
  • Compilación de gráficos muy lenta
  • Ejecución de grafo muy lenta
  • XLA frecuentemente: transferencias de CPU
  • HBM repetido del dispositivo al host de RAM
  • Desfragmentación repetida de HBM
  • Operaciones aten:: reducidas

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:

Cambio:

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

Estos cambios generan compilaciones y recuperaciones de tensores de manera artificial. dynamic_graph=True cambia de forma artificial el tamaño del lote para cada paso, lo que causa que los grafos de XLA disminuyan en un paso y cada una de nuevo. fetch_often=True inserta llamadas a loss.item() en cada paso que genera la obtención de valores de 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
   

Cuando se realiza una depuración, se recomienda ejecutar el objeto --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 las métricas automáticas. En este ejemplo, puedes ver TransferFromServerTime se vio con demasiada frecuencia, una vez por paso. Esto se debe al proceso de entrenamiento que recupera el valor de loss.item() en cada paso. Además, es posible que veas una advertencia de CompileTime demasiado frecuente, cuando tu modelo deba volver a compilarse repetidamente debido a formas dinámicas en el grafo. El siguiente fragmento de código genera 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, cuando termine, deberías ver un resumen de las operaciones no reducidas. Ten en cuenta que aten::_local_scalar_dense es una operación especial que corresponde a la recuperación de tensores de XLA al contexto de CPU.

En este informe, puedes ver que hay dos lugares principales en los que 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 secuencia de comandos a continuación, que contiene una operación sin disminuir (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

Si 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 línea 11 de STDIN está causando esta operación no baja (es decir, la línea con torch.nn.CTCLoss()). En este momento, la operación ctc_loss no se redujo. , por eso ve 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 debido a que la operación no se reduce, primero debes transferir los tensores de XLA a La CPU, ejecuta la operación aten:: en la CPU y vuelve a transferirla.

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 la VM y, luego, borra la TPU, la VM y el depósito 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}

VM TPU + perfil de PyTorch/XLA

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

Exportar variables de entorno

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

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

Crear una Cloud TPU

Consulta elGuí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 Instalada.

  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. SSH en la VM, instala la 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 esto:

image

Si ves el resultado de TensorBoard en http://localhost:9001, también puedes acceder a la página de perfil anterior seleccionando la opción PROFILE en el menú desplegable de la esquina superior derecha, junto al botón SUBIR:

image

Genera el perfil del modelo

En unnueva terminal en tu entorno de desarrollo, exporta las mismas variables de entorno que antes yssh en tu VM de TPU:

  1. En la ventana de la terminal nueva, exporta el ID del proyecto y las variables de zona de nuevo, ya que se encuentra en una shell nueva.

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

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
  3. Clona el repositorio PyTorch/XLA y ejecuta nuestra 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:

    Cambio:

        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, causar artificialmente los grafos dinámicos y las recuperaciones de tensores, que se exploran más adelante en la sección de análisis de métricas automáticas. Por ahora, dado que solo estás generando el perfil de la TPU, el siguiente ejemplo se ejecuta con un 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 y perfil de cliente

Una vez que el entrenamiento esté en ejecución, mira el resultado de TensorBoard en http://localhost:9001 y captura un perfil con las siguientes instrucciones:

image

Debería ver que se volvió a cargar la siguiente página:

image

En la configuración de la VM de TPU, solo se selecciona la herramienta de visualizador de seguimiento, por lo que en el menú desplegable Herramientas selecciona trace_viewer y, luego, inspecciona los seguimientos. Puedes ver que en la configuración de VM de TPU se muestran los seguimientos del lado del “cliente” y del dispositivo de la TPU en una vista completa:

image

Limpieza

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

    (vm)$ exit
    

Borra la VM de TPU que creaste:

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

    $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} --zone=${ZONE}
    
  2. Verifica que los recursos se hayan borrado mediante la ejecución del siguiente comando. La eliminación puede tomar 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.