Se usó la API de Cloud Translation para traducir esta página.

Cómo perfilar tu modelo en VMs de Cloud TPU

El perfilado es una de las principales herramientas para optimizar el rendimiento de tus modelos en 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 de todos los frameworks basados en XLA, incluidos JAX, PyTorch XLA y TensorFlow/Keras.

Cómo capturar perfiles

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

Captura programática
Captura a pedido (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 deseas capturar perfiles. Por lo general, capturas un perfil para unos pocos pasos de entrenamiento o perfilas un bloque de código específico dentro de tu modelo. Existen diferentes formas de capturar registros en los diferentes frameworks JAX, PyTorch XLA y TensorFlow, ya sea con el inicio o la detención del registro basados en la API, o bien 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 ejecutas MaxText.

Puedes usar la captura de perfiles a pedido cuando quieras capturar perfiles de forma ad hoc o cuando no habilitaste 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 deseas capturar perfiles en ese instante durante un período para diagnosticar el problema.

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

Puedes obtener información para habilitar la captura programática y a pedido para diferentes frameworks en los siguientes vínculos:

Tanto para la captura de perfiles programática como para la captura de perfiles a pedido, debes especificar la ubicación en la que deseas almacenar los perfiles capturados. Puedes almacenar los perfiles en un directorio de tu 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 mencionadas anteriormente, tus perfiles se capturarán en las siguientes ubicaciones:

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

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 ejecución de entrenamiento de los pasos 1 a 3 y, luego, capturas perfiles de los pasos 8 a 10, estos son perfiles para la misma ejecució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 diferentes marcas de fecha y hora en cada ejecución. Puedes capturar perfiles en diferentes sesiones de forma programática, a pedido o una combinación de ambos. Esto es útil cuando deseas comparar perfiles de diferentes partes del entrenamiento de tu modelo. Por ejemplo, comparar los perfiles del comienzo 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 un pod de GKE.

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

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

$ tensorboard --logdir=profiles

Para cargar varios perfiles de varias ejecuciones, apunta TensorBoard al directorio raíz que contiene todas tus 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 de cloud-diagnostics-xprof.

El repositorio de cloud-diagnostics-xprof proporciona las siguientes mejoras en comparación con 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 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 finalice 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 según las necesidades del usuario en cuanto a velocidad de carga y costo
Crea un vínculo para compartir perfiles y colaborar con miembros del equipo y con ingenieros de Google.
Es más fácil realizar la generación de perfiles a pedido de las 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 de XProfiler, debes capturar perfiles para el código de tu carga de trabajo de forma programática o iniciando el servidor de perfiles, de modo que puedas realizar la captura de perfiles a pedido más adelante.

Para configurar la biblioteca de XProfiler, debes configurar Google Cloud CLI y un entorno virtual de Python. Luego, solo debes ejecutar el siguiente comando:

pip install cloud-diagnostics-xprof

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

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

xprofiler create -z $ZONE -l $GCS_PATH

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

Reemplaza $ZONE por cualquier zona y $GCS_PATH por la ruta de acceso a los registros de seguimiento de tu perfil. Puedes especificar el directorio raíz que contiene todos los registros de perfil de varias ejecuciones o un conjunto específico de registros de perfil de una ejecución.

Por ejemplo, si capturas 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 establecer la ruta de acceso de GCS en el 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. Puedes cambiar el tipo de máquina con la marca -m. Si deseas 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 tu carga de trabajo implementada en GKE.

La VM o el pod de GKE que aloja TensorBoard hace que la carga de perfiles grandes y múltiples perfiles sea mucho más rápida en Google Cloud que en TensorBoard alojado de forma local. Según nuestras comparativas, verás perfiles de la carga de 1 GB en unos 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 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 vínculo en el que puedes hacer clic para ver tus perfiles de XProf en un TensorBoard. Este vínculo permite compartir perfiles con tu equipo y con los ingenieros de Google que ayudan a optimizar el rendimiento en Google Cloud.

Tú controlas 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 tus perfiles con el comando xprofiler connect.

Si habilitaste el servidor de Cloud Profiler en el código de tu carga de trabajo (lo que se requiere para las capturas programáticas y a pedido) y deseas realizar la generación de perfiles a pedido, 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 tu carga de trabajo para crear un perfil. Admitimos la captura a pedido para las 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, los nombres de la VM host o el 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.