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.

Conceptos

Si recién comienzas a usar PyTorch / XLA, consulta los documentos de API_guid y solución de problemas de PyTorch. Para Cloud TPU, consulta el documento de conceptos.

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

Crea e inicializa los recursos necesarios relacionados con Cloud TPU.

  1. Crea variables para el ID del proyecto, el bucket de Cloud Storage y la zona que se usará para los 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 sus 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 instalaste: instrucciones de instalación.

  5. Usa el comando gsutil mb para crear un bucket de Cloud Storage en el que se almacenan todos los artefactos de generación de perfiles. Reemplaza las variables de región y nombre del bucket por 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 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 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 en el puerto de reenvío de VM 9001 en la VM al puerto 9001 en la 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 la 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 algo similar a lo siguiente:

imagen

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

imagen

Perfila el modelo

Para mantener activo el servidor de TensorBoard, desde tu máquina local, inicia una nueva ventana de la terminal 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 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 de extremo a extremo en tu entorno:

    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 generan de forma artificial grafos dinámicos y recuperaciones de tensores, que se exploran más adelante en la sección Análisis de métricas automáticas. Por ahora, como recién generas el perfil de la TPU, el siguiente ejemplo se ejecuta con el 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
    

Generación de perfiles de TPU (servidor)

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

imagen

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

imagen

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

  • Página de descripción general (esto no incluye mostrar el analizador de canalización de entrada, que es muy diferente entre TensorFlow/TPU y PyTorch / TPU de XLA).
  • Lector de memoria
  • Perfil de operaciones
  • Visor de pod
  • Estadísticas de TensorFlow (estadísticas de nivel de framework, es decir, estadísticas de PyTorch)
  • Visor de Trace (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 está entrenando 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 HLO PyTorch. El visualizador de memoria captura una vista de la memoria por módulo de HLO, por lo que habrá algunos módulos, como los grafos de asignación de datos del dispositivo (entrada y etiquetas enviados por lotes). Para ver el uso de memoria de un módulo HLO en particular, selecciona el módulo desde el menú desplegable Hosts a la izquierda:

imagen

Una vez que estés viendo el módulo HLO seleccionado, obtendrás una vista integral del cronograma de ejecución del módulo HBM del módulo. Esto se ordena 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

Cada una de las asignaciones de búfer se puede inspeccionar aún más mediante el desplazamiento sobre ellas. Por ejemplo, para determinar qué asignación ocupa la mayor parte del HBM del dispositivo, haz lo siguiente:

imagen

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

   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 desde el espacio de nombres de operación test_mnist, puedes 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 estadísticas de rendimiento de las operaciones de XLA que se ejecutan durante un período de generación de perfiles. El perfil de operación muestra lo siguiente:

  • La eficacia con la que tu aplicación usa Cloud TPU como un porcentaje de tiempo dedicado a las operaciones por categoría y al uso de FLOPS de TPU.
  • Las operaciones que más tiempo consumieron. Estas operaciones son objetivos potenciales de optimización.
  • Detalles de las operaciones individuales, incluidas la forma, el padding 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 qué operaciones de XLA tardan más en ejecutarse y cuántos FLOPS de TPU consumen.

imagen

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 establecer 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 en la que se enumeran las categorías de operación de TensorFlow principales asociadas con las operaciones de XLA. Estas operaciones se ordenan según el porcentaje de uso de Cloud TPU.
  4. Tarjetas de detalles de operaciones Son los detalles sobre la operación que aparece cuando se desplaza sobre una operación en la tabla. Estos incluyen el uso de FLOPS, la expresión en la que se usa la operación y el diseño de la operación (ajuste).

Visor de pod

Consulta las herramientas de TPU para obtener una descripción completa de la herramienta de visor de pods.

Estadísticas del framework (estadísticas de TensorFlow/PyTorch)

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

Lector de seguimiento

El visualizador de seguimientos es una herramienta de análisis de rendimiento de Cloud TPU. El visor de seguimiento utiliza el visor 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ó el modelo de TensorFlow.
  • La parte del sistema (máquina host o TPU) que ejecutó una operación. Por lo general, para PyTorch/XLA, la máquina anfitrión 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 visualizador de seguimientos te permite identificar problemas de rendimiento en el modelo y, luego, tomar medidas para resolverlos. La perforación, puede identificar qué operaciones de PyTorch / XLA están tardando más en ejecutarse.

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

imagen

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

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

Generación de perfiles de clientes de PyTorch / XLA

De forma similar a cuando generaste un perfil de la TPU mientras la ejecución del modelo estaba en curso, ahora harás el perfil del cliente de PyTorch / XLA durante el entrenamiento. La herramienta principal de supervisión en el lado del cliente es el visualizador de Trace.

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

imagen

Para capturar seguimientos de varios procesos, cada proceso puede iniciar servidores de generación de perfiles 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 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

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 pequeñas de MNIST, los tiempos de los pasos serán cortos y no necesariamente serán uniformes. Como ejercicio, puedes intentar agregar los seguimientos e iniciar un servidor del 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 TransferToServer y TransferFromServer muestran la cantidad exacta de tensores que se envían y reciben, y su tamaño total:

imagen

En las compilaciones de grafos de XLA, puedes ver el hash del 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 generar perfiles de manera programática de la TPU y del 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:

  • Grafos dinámicos / compilaciones continuas
  • Compilación de grafos muy lenta
  • Ejecución muy lenta del grafo
  • Transferencias frecuentes de XLA→CPU
  • HBM de dispositivo repetido para alojar intercambio de RAM
  • Desfragmentación de HBM repetida
  • Se disminuyeron las operaciones de aten::

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 provocan compilaciones y recuperaciones de tensores de forma artificial. dynamic_graph=True cambia artificialmente el tamaño del lote para cada paso, lo que hace que los grafos reducidos de XLA sean diferentes en cada paso y se vuelva a compilar. fetch_often=True inserta llamadas a loss.item() en cada paso, lo que genera valores de tensor desde el 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 realices la depuración, la práctica recomendada es ejecutar con --num_cores=1, ya que simplifica el proceso de depuración. El resultado de muestra se verá 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 vio 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 demasiado frecuente, cuando se deba volver a compilar el modelo repetidas veces debido a formas dinámicas en el gráfico. 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, al final, 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 en el contexto de CPU.

En este informe, verás que hay dos lugares principales en los que se llama a la operación aten::_local_scalar_dense, y 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 el análisis de métricas automáticas en la siguiente secuencia de comandos que contiene una operación no reducida (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 11.a línea de STDIN está causando esta operación no disminuida (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 están inicialmente en XLA:TPU antes de la ejecución. Sin embargo, como la operación no se reduce, primero debes transferir los tensores de XLA a 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 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 + TPU/XLA de TPU

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

Exporta variables de entorno

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

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

Crear una Cloud TPU

Consulta el Guía del usuario de VM de TPU y, después de configurar, crear una VM TPU v3-8, que viene contorch ,torch_xla, torchvision ytensorboard preinstalado.

  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 si seleccionas la opción PERFIL del menú desplegable en la esquina superior derecha junto al botón SUBIR:

imagen

Perfila el modelo

En una ventana de la terminal nueva en el entorno de desarrollo, exporta las mismas variables de entorno que antes y ssh a la VM de TPU:

  1. En la ventana de la terminal nueva, vuelve a exportar las variables de ID y zona del proyecto, 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 de 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 generan de forma artificial grafos dinámicos y recuperaciones de tensores, que se exploran más adelante en la sección Análisis de métricas automáticas. Por ahora, como recién generas el perfil 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 se esté ejecutando el entrenamiento, visualiza el resultado de TensorBoard en http://localhost:9001 y captura un perfil mediante las siguientes instrucciones:

imagen

Deberías ver la siguiente página recargada:

imagen

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

imagen

Limpieza

  1. Sal de la 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 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. Ejecuta el siguiente comando para verificar que los recursos se hayan borrado. La eliminación puede tomar varios minutos. Una respuesta como la que se muestra a continuación indica que tus instancias se borraron correctamente.

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