Genera perfiles para tu modelo con VMs de Cloud TPU

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

Captura perfiles

Para usar XProf, primero debes capturar un perfil del modelo. A continuación, se indican las dos formas de capturar perfiles:

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

En la captura programática, debes anotar el código del modelo para especificar en qué parte del código quieres capturar perfiles. En general, debes capturar un perfil para unos pocos pasos de entrenamiento o perfilar un bloque de código específico dentro del modelo. Existen diferentes maneras de capturar trazas en los distintos frameworks de JAX, PyTorch XLA y TensorFlow, ya sea con el inicio o la detención de la traza basados en la API o con el administrador de contexto. En frameworks de nivel superior, como MaxText, puedes habilitar la captura de perfiles con solo habilitar una marca profiler=xplane cuando lo ejecutas.

Puedes usar la captura de perfiles bajo demanda para capturar perfiles de forma ad hoc o si no habilitaste la captura de perfiles programática. Esto puede ser útil cuando detectas un problema con las métricas del modelo durante la ejecución y quieres capturar perfiles en ese instante durante un período para diagnosticar el problema.

Para habilitar la captura de perfiles bajo demanda, aún debes iniciar el servidor de XProf en el código. Habilitar profiler.start_server iniciará un servidor de XProf en la carga de trabajo de AA que escuchará el activador de captura bajo demanda para empezar a capturar perfiles. Puedes activar la captura de perfiles bajo demanda a través de la IU de TensorBoard o de la CLI con la herramienta XProfiler.

Conoce cómo se habilita la captura programática y la bajo demanda de diferentes frameworks en los vínculos que se proporcionan más abajo:

Para ambos tipos de capturas, debes especificar la ubicación en la que quieres almacenar los perfiles capturados. Puedes optar por almacenar los perfiles en un directorio de la VM de TPU con una ruta de acceso similar a /profiles/run-name o usar Cloud Storage con una ruta de acceso similar a gs://bucket-name/run-name/.

Con las rutas de acceso mencionadas, los perfiles se capturarán en las ubicaciones que se indican a continuación:

/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 del modelo. Por ejemplo, si capturas perfiles en una ronda de entrenamiento del paso 1 al 3 y, luego, de los pasos 8 a 10, estos perfiles corresponderán a la misma ronda, pero la primera captura del paso 1 al 3 será session1 y la segunda captura del paso 8 al 10 será session2.

Las distintas sesiones se indicarán con marcas de fecha y hora diferentes en cada ronda. Los perfiles se pueden capturar en distintas sesiones de forma programática, bajo de manda o a través de una combinación de ambas. Esto es útil si quieres comparar los perfiles de diferentes partes del entrenamiento del modelo. Por ejemplo, para comparar los perfiles del inicio del entrenamiento con los del final.

Visualiza 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 en un Pod de GKE.

Si capturaste los 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 los comandos que se indican a continuación:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Para ver los perfiles capturados con TensorBoard, ejecuta el comando siguiente:

$ tensorboard --logdir=profiles/run-name

o

$ tensorboard --logdir=profiles

Para cargar muchos perfiles de varias ejecuciones, apunta TensorBoard al directorio raíz que contiene todas las ejecuciones y perfiles, en lugar de dirigirlo al subdirectorio que contiene los perfiles de una ejecución específica.

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

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

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

  • Configuración y empaquetado de las dependencias de XProf y TensorBoard
  • Almacenamiento de los perfiles en Cloud Storage, lo que puede ser útil para la retención a largo plazo y el análisis posterior a la ejecución (los perfiles locales capturados se borrarán después de que el investigador termine la ejecución)
  • Carga rápida de perfiles grandes y de varios perfiles aprovisionando TensorBoard en una VM de Compute Engine o un Pod de GKE, con la opción de cambiar el tipo de máquina según las necesidades del usuario en cuanto a la velocidad de carga y el costo
  • Creación de un vínculo para compartir perfiles y colaborar con miembros del equipo y con ingenieros de Google
  • Generación de perfiles bajo demanda más sencilla de las cargas de trabajo en GKE y Compute Engine con el objetivo de elegir cualquier host que ejecute la carga de trabajo para capturar perfiles

Antes de usar la biblioteca XProfiler, debes capturar perfiles del código de la carga de trabajo de forma programática o iniciar el servidor de perfiles, de modo que puedas capturar perfiles bajo demanda más adelante.

Para configurar la biblioteca XProfiler, debes establecer los ajustes de la Google Cloud CLI y de un entorno virtual de Python. Luego, solo debes ejecutar lo que se indica a continuación:

pip install cloud-diagnostics-xprof

Esto instalará todas las dependencias necesarias de XProf y TensorBoard.

A continuación, con el objetivo de crear una VM o un Pod de GKE para alojar TensorBoard, ejecuta el comando que se proporciona más abajo:

xprofiler create -z $ZONE -l $GCS_PATH

o

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

Reemplaza $ZONE por cualquier zona y $GCS_PATH por la ruta de acceso a las trazas del perfil. Puedes especificar el directorio raíz que contiene todas las trazas del perfil de varias ejecuciones o un conjunto específico de trazas del perfil de una ejecución.

Por ejemplo, si capturas perfiles en los directorios siguientes:

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 establecer la ruta de acceso de GCS al directorio raíz (gs://bucket-name):

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

En la IU 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, específicamente una VM c4-highmem-8. El tipo de máquina se puede cambiar con la marca -m. Si quieres crear una instancia de TensorBoard en un Pod de GKE, en lugar de 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 administración de TensorBoard junto con el resto de la carga de trabajo implementada en GKE.

La VM o el Pod de GKE que aloja TensorBoard permite una carga mucho más rápida de perfiles grandes y de varios perfiles en Google Cloud que en TensorBoard alojado a nivel local. Según nuestras comparativas, verás perfiles con una carga del orden de 1 GB en cuestión de minutos para la primera carga con la VM c4-highmem-8 predeterminada. Además, puedes elegir el tipo de máquina adecuado según tus necesidades de rendimiento y costos.

Después de ejecutar xprofiler create, verás un resultado similar al que se indica a continuación:

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 vínculo en el que puedes hacer clic para ver los perfiles de XProf en un TensorBoard. Este vínculo te permite compartir perfiles con tu equipo y con los ingenieros de Google que te ayudan a optimizar el rendimiento en Google Cloud.

Tú debes controlar quién tiene acceso al vínculo según los permisos establecidos para el bucket de Cloud Storage que almacena tus datos de perfil. En caso de que el vínculo no funcione por algún motivo, también proporcionamos una forma de conectarse a la instancia de TensorBoard con SSH para ver los perfiles con el comando xprofiler connect.

Si habilitaste el servidor de Cloud Profiler en el código de la carga de trabajo (lo que se requiere para las capturas programáticas y bajo demanda) y quieres generar perfiles bajo demanda, puedes hacerlo de dos maneras:

a. Haz clic en el botón Capture profile en TensorBoard. Ahora puedes elegir cualquier host de dispositivo en el que se ejecute la carga de trabajo para crear su perfil. Google admite la captura bajo demanda de cargas de trabajo que se ejecutan en GKE o Compute Engine.

b. Usa el comando de captura de XProfiler y proporciona información como la zona, el bucket de Cloud Storage, el framework, el nombre de la VM host o del Pod, y la duración de la captura en milisegundos. Esta es la misma información que debes ingresar en la IU de TensorBoard.

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