Cómo portar cargas de trabajo de JAX a Pathways

Debido a la naturaleza distribuida de JAX con Pathways, es posible que algunas operaciones no escalen bien debido a las sobrecargas de comunicación. Si bien Pathways minimiza estos gastos con funciones como el envío asíncrono, hay algunas cosas que debes tener en cuenta cuando portas cargas de trabajo de JAX a Pathways o cuando escalas una carga de trabajo de JAX con Pathways a una gran cantidad de aceleradores.

Antes de comenzar

Asegúrate de tener lo siguiente:

Índice de procesos

JAX con Pathways trata todos los dispositivos de tu clúster de Pathways como locales. Esto simplifica la administración de dispositivos y permite que JAX use todos los recursos disponibles. En la práctica, esto significa lo siguiente:

  • jax.process_index() siempre es 0 para todos los dispositivos.
  • jax.devices() y jax.local_devices() devuelven todos los dispositivos TPU de todo el trabajo.

Tipo de hardware y coubicación

Para obtener el mejor rendimiento, coloca todos los componentes de Pathways y el trabajo del usuario en la misma zona de Google Cloud Cloud. Usa una CPU grande, como el proxy de IFRT y el administrador de recursos. Recomendamos al menos un n2-standard-64 dedicado que incluya 64 CPU virtuales y 256 GB de memoria.

PathwaysUtils

Pathways-utils es un repositorio de GitHub basado en Python que proporciona utilidades y herramientas esenciales que te permiten optimizar la implementación y la ejecución de cargas de trabajo de JAX en la arquitectura de Pathways on Cloud. Este paquete controla las adaptaciones necesarias para el entorno de nube, lo que permite que los desarrolladores de JAX se enfoquen en sus flujos de trabajo principales de aprendizaje automático con una configuración mínima específica de la plataforma. Específicamente, ofrece lo siguiente:

  • Un backend de JAX "proxy": este backend personalizado permite que tu aplicación de JAX use la infraestructura de Pathways configurando la variable de entorno JAX_PLATFORMS=proxy.
  • Utilidades de generación de perfiles integradas: Funciones de generación de perfiles que te permiten comprender el rendimiento de tu aplicación. Si usas las APIs de generación de perfiles de JAX estándar, como jax.profiler.start_trace y jax.profiler.start_server, puedes generar perfiles no solo de tu código de JAX, sino también de los componentes de Pathways subyacentes, lo que proporciona una vista integral de la ejecución dentro del entorno de la nube.
  • Puntos de control distribuidos con Orbax: un controlador de puntos de control de Orbax personalizado que te permite usar puntos de control distribuidos y restablecerlos cuando usas la biblioteca de Orbax en el entorno de Pathways. El objetivo de esta integración es funcionar sin necesidad de realizar cambios en tu código de puntos de control de Orbax existente, siempre que importe pathwaysutils.
  • Primitivas de entrenamiento elástico: Proporciona primitivas de entrenamiento elástico fundamentales que puedes usar para compilar flujos de trabajo de entrenamiento sólidos y escalables con Pathways. Estas primitivas permiten que tus trabajos de entrenamiento se adapten de forma dinámica a los cambios en los recursos disponibles, lo que mejora la eficiencia y la resiliencia en los entornos de nube.

Creación de puntos de control

Orbax se prueba de forma exhaustiva con Pathways para crear puntos de control distribuidos y restablecerlos con Cloud Storage. Cuando configuras la variable de entorno ENABLE_PATHWAYS_PERSISTENCE=1 y llamas a import pathwaysutils; pathwaysutils.initialize() en tu train.py, se registra un ArrayHandler personalizado que controla de manera eficiente las operaciones de punto de control a través del proxy de IFRT, lo que permite que los trabajadores de Pathways en aceleradores guarden y restablezcan datos directamente.

Python ubicado en el mismo lugar

Python ubicado en el mismo lugar es una API de JAX de código abierto que te permite ejecutar código de Python especificado por el usuario directamente en los hosts de TPU o GPU, lo que es más sencillo en JAX con varios controladores. Con este enfoque, puedes evitar la transferencia de datos entre el cliente y las máquinas de TPU para tareas de procesamiento intensivo, como la carga de datos y la creación de puntos de control. Para configurar tu clúster de Pathways de modo que ejecute la API de JAX de Python ubicada en el mismo lugar, compila y envía una imagen de contenedor sidecar con el Dockerfile sidecar de muestra.

Carga de datos

Durante el entrenamiento, cargamos lotes de un conjunto de datos de forma repetida para alimentar el modelo. Es importante tener un cargador de datos asíncrono y eficiente que divida el lote en fragmentos entre los hosts para evitar que los aceleradores se queden sin trabajo. Cuando se ejecuta el entrenamiento con Pathways, el cargador de datos se ejecuta en una VM de CPU (a diferencia de una VM de TPU que se usa en configuraciones de varios controladores) y envía datos a las VMs de TPU. Esto genera una latencia más alta en la lectura de datos, pero se mitiga parcialmente leyendo con anticipación una cantidad X de lotes en el host de la CPU y enviando los datos leídos de forma asíncrona a las TPU. Esta solución es suficiente cuando se ejecuta a pequeña o mediana escala.

Cuando ejecutas un cargador de datos de varios procesos (como Grain) en el contenedor de carga de trabajo principal en GKE, el contenedor comparte el nodo principal con los sidecars pathways-proxy y pathways-rm. Debido a esta arquitectura compartida, debes configurar lo siguiente:

  • Establece un recuento de trabajadores explícito: Dentro de un contenedor, os.cpu_count() informa las CPU virtuales totales de la VM host en lugar del límite de CPU del contenedor. Pasa un recuento de trabajadores explícito (no más que el límite de CPU del contenedor) en lugar de derivarlo de os.cpu_count().
  • Activa un volumen /dev/shm respaldado por memoria: Los cargadores de datos de varios procesos usan la memoria compartida de POSIX (/dev/shm) para la comunicación entre procesos. Asegúrate de que un volumen emptyDir (medium: Memory) respaldado por memoria esté activado en /dev/shm para que los lotes recuperados previamente no excedan el límite de memoria compartida del contenedor predeterminado.

Para obtener un rendimiento óptimo a gran escala, te recomendamos que ubiques tu canalización de datos de entrada con Python ubicado en el mismo lugar para ejecutar tu canalización de datos directamente en los aceleradores. Esto elimina el cuello de botella de la CPU y aprovecha las interconexiones rápidas de la TPU para la transferencia de datos.

Puedes encontrar una implementación de referencia para migrar una canalización de entrada basada en TFDS en la implementación de RemoteIterator en multihost_dataloading.py. Esta implementación funciona en JAX y Pathways de varios controladores de manera distribuida con la API de JAX de Python ubicada en el mismo lugar.

Control de versiones de Jax

Las versiones de Pathways están estrechamente vinculadas con las versiones de JAX para garantizar la compatibilidad y la estabilidad. Para evitar posibles problemas, verifica que tus artefactos de Pathways y tu versión de JAX estén alineados. Cada versión de Pathways especifica con claridad las versiones de JAX compatibles a través de una etiqueta con el formato jax-<version>.

Caché de compilación

La caché de compilación persistente de Pathways es una función que permite que los servidores de Pathways almacenen ejecutables de XLA compilados en una ubicación persistente, como Cloud Storage, para evitar la compilación redundante. Esta función está habilitada de forma predeterminada. La ubicación de la caché se pasa como una marca --gcs_scratch_location a los contenedores de trabajadores de Pathways y del administrador de recursos. Para mantener al mínimo los costos de almacenamiento asociados, la caché adjunta una política de ciclo de vida a la ubicación de Cloud Storage. Existe un límite de 50 políticas por bucket de Cloud Storage. Por lo tanto, recomendamos usar una ubicación común de Cloud Storage en todas las cargas de trabajo.

Si omites la marca --gcs_scratch_location o la dejas vacía, Pathways recurre automáticamente a una caché de compilación en la memoria. Esta caché no conserva los ejecutables de XLA compilados en las implementaciones de Pathways.

Esta caché es similar a la caché de compilación de JAX, que pathwaysutils.initialize() inhabilita para las cargas de trabajo de Pathways.

Se requieren los siguientes permisos de Cloud Storage para la caché de compilación:

  • storage.buckets.get: Para recuperar los metadatos del bucket.
  • storage.buckets.update: Esencial para que Pathways configure políticas de ciclo de vida de objetos y aplique el TTL para la expulsión de la caché.
  • storage.objects.list: Para enumerar los objetos de caché existentes dentro del bucket.
  • storage.objects.create: Para escribir nuevos ejecutables compilados en la caché.
  • storage.objects.get: Para leer los ejecutables almacenados en caché desde el bucket.

Genera perfiles

Puedes usar el generador de perfiles de JAX para generar seguimientos de un programa de JAX. Existen dos formas comunes que se admiten con Pathways:

  • Programática
    • Captura perfiles de forma programática desde tu código de JAX
  • Manual
    • Captura perfiles bajo demanda después de iniciar el servidor del generador de perfiles desde tu código de JAX

En ambos casos, los perfiles se escriben en un bucket de Cloud Storage. Se crearán varios archivos de seguimiento en el bucket de Cloud Storage, posiblemente en diferentes carpetas de marca de tiempo, por ejemplo:

  • Es el proceso principal de Python que invocó el seguimiento (por lo general, la VM de tu notebook): <jax-client-vm-name>.xplane.pb
  • Proxy de IFRT de Pathways: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Administrador de recursos de Pathways: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Trabajadores de Pathways: server.*<tpu-node-name>.xplane.pb

Estos archivos de seguimiento se pueden analizar con TensorBoard ejecutando el siguiente comando. Para obtener más información sobre TensorBoard y todas sus herramientas de generación de perfiles, consulta Optimiza el rendimiento de TensorFlow con el generador de perfiles.

# verify trace files are present
gcloud storage ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --load_fast=false --logdir=gs://BUCKET/PREFIX

Reemplaza lo siguiente:

  • BUCKET : un bucket de Cloud Storage para almacenar los archivos de seguimiento
  • PREFIX: Es una ruta de acceso dentro de tu bucket de Cloud Storage para almacenar los archivos de seguimiento.

Captura de perfiles programática

Captura un perfil desde tu código. Los perfiles se guardan en gs://<bucket>/<prefix>, en un directorio con marca de tiempo.

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

Captura manual de perfiles

Para capturar la información de perfil de forma manual, debes iniciar el servidor del generador de perfiles desde tu código de Python:

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_server(JAX_PROFILER_PORT)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

Mientras se ejecuta el servidor del generador de perfiles, puedes capturar un perfil y exportar los datos a la ubicación de destino de Cloud Storage:

export DURATION_IN_SECS=6
export OUTPUT_DIR="gs://BUCKET/PREFIX/"

curl -X POST http://localhost:JAX_PROFILER_PORT/profiling \
     -H "Content-Type: application/json" \
     -d "{\"duration_ms\": $((DURATION_IN_SECS * 1000)), \"repository_path\": \"${OUTPUT_DIR}\"}"

Puedes encontrar información de tiempo para los métodos de cliente proxy de IFRT, como Compile y Execute, en el seguimiento de tu programa. Estos eventos, que detallan las interacciones con el servidor gRPC del proxy de IFRT durante la compilación y la ejecución, aparecen en el subproceso llamado GrpcClientSessionUserFuturesWorkQueue. Si examinas este subproceso en tu registro, puedes obtener estadísticas sobre el rendimiento de estas operaciones.

Marcas de XLA

Cuando usas Pathways, debes configurar las marcas de XLA en el contenedor pathways-proxy.

Cuando uses Cluster Toolkit (el comando gcluster job submit), establece las marcas de XLA con --pathways-proxy-args:

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

Volcado de HLO

Para inspeccionar las entradas del optimizador de alto nivel (HLO) que se proporcionan al compilador de XLA, puedes configurar Pathways para que vuelque el HLO en una ubicación de Cloud Storage especificada con --pathways-proxy-args:

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

¿Qué sigue?