Crear un perfil de tu modelo en VMs de TPU de Cloud

La creación de perfiles es una de las principales herramientas para optimizar el rendimiento de tus modelos en las TPU de Cloud. La herramienta de creación de perfiles principal se llama XProf y está disponible en el repositorio de GitHub OpenXLA/XProf. XProf admite la creación de perfiles de todos los frameworks basados en XLA, incluidos JAX, Pytorch XLA y TensorFlow/Keras.

Capturar perfiles

Para usar XProf, primero debes capturar un perfil de tu modelo. Hay dos formas de capturar perfiles:

  1. Captura programática
  2. Captura bajo demanda (también conocida como captura manual)

Con la captura programática, debes anotar el código de tu modelo para especificar en qué parte del código quieres capturar los perfiles. Normalmente, se captura un perfil durante unos cuantos pasos de entrenamiento o se perfila un bloque de código específico dentro del modelo. Hay diferentes formas de capturar trazas en los distintos frameworks JAX, PyTorch XLA y TensorFlow: se puede iniciar o detener una traza mediante una API o mediante un gestor de contexto. En frameworks de nivel superior, como MaxText, puedes habilitar la captura de perfiles con solo habilitar una marca profiler=xplane al ejecutar MaxText.

Puedes usar la captura de perfiles bajo demanda cuando quieras capturar perfiles de forma puntual o cuando no hayas habilitado la captura de perfiles programática. Esto puede ser útil cuando detectas un problema con las métricas de tu modelo durante la ejecución y quieres capturar perfiles en ese instante durante un periodo determinado para diagnosticar el problema.

Para habilitar la captura de perfiles bajo demanda, debes iniciar el servidor XProf en tu código. Si habilitas profiler.start_server, se iniciará un servidor XProf en tu carga de trabajo de aprendizaje automático que estará atento al activador de captura bajo demanda para empezar a capturar perfiles. Puedes activar la captura de perfiles bajo demanda a través de la interfaz de TensorBoard o de la CLI con la herramienta XProfiler.

Puedes consultar cómo habilitar la captura programática y bajo demanda para diferentes frameworks en los siguientes enlaces:

Tanto para la captura de perfiles programática como para la bajo demanda, debe especificar la ubicación en la que quiere almacenar los perfiles capturados. Puedes almacenar los perfiles en un directorio de tu máquina virtual de TPU con una ruta similar a /profiles/run-name o usar Cloud Storage con una ruta similar a gs://bucket-name/run-name/.

Con las rutas mencionadas, tus perfiles se registrarán en:

/profiles/run-name/plugins/profile/session1/

o

gs://bucket-name/run-name/plugins/profile/session1/.

Una sesión representa una captura de perfiles durante una ejecución de tu modelo. Por ejemplo, si capturas perfiles en una sesión de entrenamiento de los pasos 1 a 3 y, más adelante, capturas perfiles de los pasos 8 a 10, se trata de perfiles de la misma sesión, pero la primera captura de los pasos 1 a 3 será session1 y la segunda captura de los pasos 8 a 10 será session2.

Las diferentes sesiones se indicarán con marcas de fecha diferentes en cada ejecución. Puedes capturar perfiles en diferentes sesiones de forma programática, bajo demanda o una combinación de ambas. Esto resulta útil cuando quiere comparar perfiles de diferentes partes del entrenamiento de su modelo. Por ejemplo, comparar los perfiles del principio del entrenamiento con los del final.

Visualizar perfiles con TensorBoard

Para ver los perfiles de XProf en TensorBoard, debes instalar el complemento tensorboard-plugin-profile.

Para visualizar los perfiles capturados en Google Cloud, te recomendamos que uses TensorBoard en una VM de TPU o la herramienta XProfiler para alojar TensorBoard en una VM o un pod de GKE.

Si has capturado tus perfiles en la VM de TPU, puedes verlos con TensorBoard en esa VM de TPU.

Para instalar XProf y TensorBoard en una VM de TPU, puedes ejecutar estos comandos:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Para ver los perfiles capturados con TensorBoard, ejecuta lo siguiente:

$ tensorboard --logdir=profiles/run-name

o

$ tensorboard --logdir=profiles

Para cargar varios perfiles de varias ejecuciones, dirija TensorBoard al directorio raíz que contiene todas sus ejecuciones y perfiles, en lugar de al subdirectorio que tiene los perfiles de una ejecución específica.

Biblioteca Cloud-Diagnostics-XProf para visualizar perfiles de TensorBoard en Google Cloud

cloud-diagnostics-xprof (también conocida como biblioteca XProfiler) facilita el alojamiento de TensorBoard y la visualización de perfiles en Google Cloud. Para obtener más información, consulta el repositorio de GitHub cloud-diagnostics-xprof.

El repositorio cloud-diagnostics-xprof ofrece las siguientes mejoras con respecto a la ejecución local de XProf y TensorBoard:

  • Configuración y empaquetado de las dependencias de XProf y TensorBoard
  • Almacena tus perfiles en Cloud Storage, lo que puede ser útil para la conservación a largo plazo y el análisis posterior a la ejecución (los perfiles locales capturados se eliminarán cuando el investigador termine la ejecución).
  • Carga rápida de perfiles grandes y múltiples aprovisionando TensorBoard en una VM de Compute Engine o un pod de GKE, con la opción de cambiar el tipo de máquina en función de las necesidades del usuario en cuanto a velocidad de carga y coste.
  • Crea un enlace para compartir perfiles y colaborar con miembros del equipo e ingenieros de Google.
  • Perfilado bajo demanda más sencillo de cargas de trabajo en GKE y Compute Engine para elegir cualquier host que ejecute tu carga de trabajo y capturar perfiles

Antes de usar la biblioteca XProfiler, debes capturar perfiles del código de tu carga de trabajo de forma programática o iniciando el servidor de perfiles para poder capturar perfiles bajo demanda más adelante.

Para configurar la biblioteca XProfiler, debes configurar la CLI de Google Cloud y un entorno virtual de Python. Después, solo tienes que ejecutar el siguiente comando:

pip install cloud-diagnostics-xprof

De esta forma, se instalarán todas las dependencias necesarias de XProf y TensorBoard.

A continuación, para crear una VM o un pod de GKE para alojar TensorBoard, ejecuta el siguiente comando:

xprofiler create -z $ZONE -l $GCS_PATH

o

xprofiler create --GKE -z $ZONE -l $GCS_PATH

Sustituye $ZONE por cualquier zona y $GCS_PATH por la ruta a tus trazas de perfil. Puedes especificar el directorio raíz que contiene todos los rastreos de perfil de varias ejecuciones o un conjunto específico de rastreos de perfil de una ejecución.

Por ejemplo, si captura perfiles en los siguientes directorios:

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

Puedes definir la ruta de GCS en el directorio raíz (gs://bucket-name):

xprofiler create -z $ZONE -l gs://bucket-name/

En la interfaz de usuario de TensorBoard, verás todos los perfiles de todas las ejecuciones y sesiones como run1/session1, run1/session2 y run2/session1.

De forma predeterminada, xprofiler create creará una VM de Compute Engine, concretamente una VM c4-highmem-8. Puedes cambiar el tipo de máquina con la marca -m. Si quieres crear una instancia de TensorBoard en un pod de GKE en lugar de en una VM de Compute Engine, puedes pasar la marca –GKE a xprofiler create. Alojar una instancia de TensorBoard en un pod de GKE puede facilitar la gestión de TensorBoard junto con el resto de la carga de trabajo desplegada en GKE.

La VM o el pod de GKE que aloja TensorBoard hace que la carga de perfiles grandes y de varios perfiles sea mucho más rápida en Google Cloud que en TensorBoard alojado localmente. Según nuestras pruebas de rendimiento, verás perfiles del orden de 1 GB de carga en unos minutos en la primera carga con la VM c4-highmem-8 predeterminada. Además, puedes elegir el tipo de máquina adecuado en función de tus necesidades de rendimiento y coste.

Después de ejecutar xprofiler create, verás un resultado similar al siguiente:

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

La primera opción es un enlace en el que puedes hacer clic para ver tus perfiles de XProf en TensorBoard. Este enlace permite compartir perfiles con tu equipo y con los ingenieros de Google que te ayudan a optimizar el rendimiento en Google Cloud.

Tú controlas quién tiene acceso al enlace en función de los permisos definidos para el segmento de Cloud Storage que almacena los datos de tu perfil. Si el enlace no funciona por algún motivo, también ofrecemos una forma de conectarse a la instancia de TensorBoard mediante SSH para ver tus perfiles con el comando xprofiler connect.

Si has habilitado el servidor de Cloud Profiler en el código de tu carga de trabajo (lo que es obligatorio tanto para las capturas programáticas como para las bajo demanda) y quieres realizar una creación de perfiles bajo demanda, puedes hacerlo de dos formas:

a. Haz clic en el botón Capturar perfil de TensorBoard. Ahora puedes elegir cualquier host de dispositivo en el que se ejecute tu carga de trabajo para crear un perfil. Admitimos la captura bajo demanda para cargas de trabajo que se ejecutan en GKE o Compute Engine.

b. Usa el comando de captura XProfiler y proporciona información como la zona, el segmento de Cloud Storage, el framework, los nombres de la VM o el pod del host y la duración de la captura en milisegundos. Es la misma información que debes introducir en la interfaz de usuario de TensorBoard.

Para obtener más información sobre cómo usar la biblioteca cloud-diagnostics-xprof, consulta esta página de GitHub.