Introducción a las rutas de aprendizaje en Cloud

Pathways es un sistema diseñado para permitir la creación de sistemas de aprendizaje automático (AA) a gran escala, de múltiples tareas y activados de forma dispersa. Te permite usar miles o decenas de miles de aceleradores, con la capacidad de asignar de forma dinámica cantidades variables de procesamiento para diferentes tareas según sus requisitos de procesamiento.

Pathways simplifica los cálculos de AA a gran escala, ya que permite que un solo cliente de JAX organice las cargas de trabajo en porciones de TPU grandes, que pueden abarcar miles de chips TPU.

Google usa Pathways de forma interna para entrenar modelos grandes, como Gemini. Pathways on Cloud ofrece los mismos beneficios a los clientes de Google Cloud .

Antes de comenzar

Antes de comenzar

Puedes ejecutar Pathways en Google Kubernetes Engine (GKE) o en un clúster de Cluster Director con Slurm. Selecciona la pestaña de tu plataforma.

GKE

Completa los siguientes requisitos previos:

En este documento, se proporciona una descripción general de cómo usar las TPU administradas por Pathways en GKE para cargas de trabajo por lotes, en tiempo real y de forma interactiva. Se supone que ya estás familiarizado con el uso de las TPU en GKE, incluidas las TPU de porción única y de porciones múltiples en GKE, y que tienes experiencia general con las TPU de porciones múltiples.

Cluster Director

Completa los siguientes requisitos previos:

En este documento, se proporciona una descripción general de cómo usar Pathways con TPU en un clúster de Cluster Director para cargas de trabajo por lotes e interactivas. Se supone que ya conoces Cluster Director, las máquinas de TPU en Cluster Director y la ejecución de trabajos con Slurm, y que tienes experiencia general con las TPU de Multislice.

Para crear un clúster, consulta Crea un clúster de Cluster Director con Pathways.

Controlador único y varios controladores

Existen dos formas principales de administrar y organizar los cálculos en varios dispositivos:

Función Controlador único (Pathways) Controlador múltiple (predeterminado de JAX)

Control

Punto de control único: Un solo programa cliente actúa como el controlador central.

Control distribuido: Participan varios procesos, cada uno con su propia instancia de intérprete de Python.

Ver

Vista unificada: El cliente ve todos los dispositivos como un único sistema unificado.

Vista localizada: Cada proceso de Python solo ve los dispositivos conectados a él.

Programación

Programación simplificada: Interactúas con un solo cliente, lo que hace que el sistema aparezca como una sola máquina grande con muchos aceleradores locales.

Programa único, datos múltiples (SPMD): Usa principalmente el paradigma de SPMD, lo que requiere que todos los dispositivos ejecuten el mismo programa.

Flexibilidad

Admite patrones de procesamiento complejos más allá de SPMD, incluido el paralelismo de canalización asimétrico y la dispersión computacional.

Puede ser menos flexible en la administración de recursos, en especial en diferentes porciones de TPU.

Componentes de Pathways

En las siguientes secciones, se describen los componentes principales de la arquitectura de Pathways.

Administrador de recursos de Pathways

Este componente es el plano de control central del sistema de Pathways. Administra todos los recursos de aceleradores y coordina la asignación de aceleradores para tus trabajos. Supervisa el estado de los trabajadores y controla la programación, la pausa y la reanudación de los trabajos. Sirve como un punto de contacto único para los errores y el estado del sistema. Este componente solo requiere recursos de CPU.

Cliente de Pathways

Este componente es una implementación del entorno de ejecución del framework interino (IFRT) que sirve como punto de entrada al sistema de Pathways. Recibe operaciones de alto nivel (HLO) de tu programa. El cliente de Pathways se coordina con el administrador de recursos de Pathways para determinar dónde colocar los programas compilados para la ejecución según tu código. Presenta una vista unificada del sistema a un cliente de JAX determinado. Este componente solo requiere recursos de CPU.

Trabajador de Pathways

Los procesos de los trabajadores de Pathways se ejecutan en las máquinas aceleradoras (VMs de TPU). Reciben ejecutables compilados de tu programa desde el servidor proxy de IFRT y realizan los cálculos en las TPU. Los trabajadores de Pathways envían datos de vuelta a tu programa a través del servidor proxy de IFRT. Este componente requiere recursos de acelerador.

Cliente proxy de IFRT

El cliente proxy de IFRT es una implementación de código abierto de la API de Interim Framework Runtime (IFRT) que desvincula el código del usuario del entorno de ejecución subyacente y mejora la portabilidad y la transparencia del código. JAX usa esta implementación como alternativa a su entorno de ejecución de varios controladores predeterminado. El cliente proxy de IFRT actúa como un puente de comunicación entre tu programa y los componentes de Pathways. Envía solicitudes al servidor proxy de IFRT y recibe resultados de él. Este componente solo requiere recursos de CPU.

Servidor proxy de IFRT

El servidor proxy de IFRT es un servidor gRPC que recibe solicitudes del cliente proxy de IFRT y las reenvía al cliente de Pathways, que controla la distribución del trabajo. Este componente solo requiere recursos de CPU.

Servidor sidecar

El servidor de sidecar es un servidor de gRPC que se ubica junto con el trabajador de Pathways en la VM del acelerador. Ejecuta el código de Python especificado por el usuario directamente en la VM del acelerador para reducir la latencia de transferencia de datos entre el controlador y los aceleradores. El servidor de sidecar interactúa con el trabajador de Pathways a través de un protocolo con versiones personalizado en el transporte de gRPC.

Diagrama de arquitectura que muestra cómo interactúan el cliente de JAX, el proxy de IFRT, el administrador de recursos de Pathways y los trabajadores de Pathways.
Componentes de Pathways

Componentes de Pathways en GKE

En esta sección, se asignan los componentes de Pathways a los componentes de GKE, como los contenedores y los Pods.

Puedes encontrar las imágenes de contenedor de Pathways en las siguientes ubicaciones:

Tipo de contenedor Ubicación

Servidor proxy de IFRT

us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:jax-JAX_VERSION

Administrador o trabajador de recursos de Pathways

us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:jax-JAX_VERSION

Administrador de recursos de Pathways

Después de crear un clúster de GKE, puedes usar el siguiente containerSpec para implementar el administrador de recursos de Pathways:

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

El contenedor usa los siguientes argumentos:

  • --server_port: Es el puerto que usa el administrador de recursos de Pathways para comunicarse con otros componentes de Pathways.
  • --node_type: Es el tipo de nodo. Establece este argumento en resource_manager para el administrador de recursos de Pathways; no es necesario para los otros contenedores.
  • --instance_count: Es la cantidad de porciones de TPU.
  • --instance_type: Es el tipo de TPU y la topología de la porción, en el formato tpuTPU_TYPE:TOPOLOGY, por ejemplo, tpuv5e:4x4.
  • --gcs_scratch_location: Es un bucket de Cloud Storage que se usa para los archivos temporales.

Servidor proxy de IFRT

Puedes usar el siguiente containerSpec para implementar un servidor proxy de IFRT:

  - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

El contenedor usa los siguientes argumentos:

  • --resource_manager_address: Es el nombre de host y el puerto que usa el servidor proxy para comunicarse con el administrador de recursos de Pathways. El puerto debe coincidir con el valor de --server_port que se usa para el contenedor del administrador de recursos de Pathways.
  • --server_port: Es el puerto que usa el servidor proxy de IFRT para comunicarse con el cliente proxy de IFRT.
  • --gcs_scratch_location: Es un bucket de Cloud Storage que se usa para los archivos temporales.

Trabajador de Pathways

Puedes usar el siguiente containerSpec para implementar trabajadores de Pathways:

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

El contenedor usa los siguientes argumentos:

  • --resource_manager_address: Es el nombre de host y el puerto que usan los trabajadores de TPU para comunicarse con el administrador de recursos de Pathways. El puerto debe coincidir con el valor de --server_port que se usa para el contenedor del administrador de recursos de Pathways.
  • --server_port: Es el puerto que usan los trabajadores para comunicarse con el servidor proxy y el administrador de recursos de Pathways.
  • --gcs_scratch_location: Es un bucket de Cloud Storage que se usa para los archivos temporales.

El administrador de recursos de Pathways, el servidor proxy de IFRT y los trabajadores de Pathways pueden tener puertos diferentes, pero, en este ejemplo, el administrador de recursos de Pathways y el trabajador de Pathways comparten el mismo puerto.

Componentes de Pathways en Cluster Director

En esta sección, se asignan los componentes de Pathways a los trabajos de Slurm en un clúster de Cluster Director.

Las secuencias de comandos de Pathways ejecutan los componentes de Pathways con Enroot. La secuencia de comandos setup_env.sh usa las siguientes imágenes de contenedor, que coinciden con la versión de JAX (0.11.1) del entorno de Python que compila:

Tipo de contenedor Ubicación

Servidor proxy de IFRT

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20261103-jax_0.11.1

Administrador o trabajador de recursos de Pathways

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20261103-jax_0.11.1

Para iniciar los componentes de Pathways, envía la secuencia de comandos start_pathways.sh como el trabajo del servicio de Pathways. No escribes los siguientes comandos por tu cuenta: las secuencias de comandos derivan los argumentos de las variables de entorno que configuras cuando envías start_pathways.sh, como SLICE_SHAPE, MIN_SLICES, MAX_SLICES y GCS_BUCKET.

Administrador de recursos de Pathways

El trabajo del servicio de Pathways ejecuta el administrador de recursos de Pathways en un nodo de CPU. Ejecuta el siguiente comando en el contenedor del servidor de Pathways:

/pathways/server \
    --server_port=29001 \
    --node_type=resource_manager \
    --instance_type=tpuTPU_TYPE:TOPOLOGY \
    --instance_count=MAX_SLICES \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

El comando usa los siguientes argumentos:

  • --server_port: Es el puerto en el que escucha el administrador de recursos de Pathways. El servidor proxy de IFRT y los trabajadores de Pathways se conectan a este puerto.
  • --node_type: el tipo de nodo, resource_manager. Solo el administrador de recursos de Pathways establece este argumento.
  • --instance_type: Es el tipo de TPU y la topología de cada porción, de SLICE_SHAPE. Por ejemplo, SLICE_SHAPE=2x4 establece tpuv6e:2x4.
  • --instance_count: Es la cantidad de porciones de TPU, de MAX_SLICES.
  • --gcs_scratch_location: Es una ruta de acceso en tu bucket de Cloud Storage, desde GCS_BUCKET, que se usa para archivos temporales, como la caché de compilación.
  • --enable_gcs_olm: indica si Pathways crea reglas de ciclo de vida de objetos de Cloud Storage para limpiar las entradas de caché de compilación antiguas. La secuencia de comandos desactiva esta opción porque requiere permisos adicionales en el bucket.

Servidor proxy de IFRT

El trabajo del servicio de Pathways ejecuta el servidor proxy de IFRT en el mismo nodo de CPU que el administrador de recursos de Pathways. Ejecuta el siguiente comando en el contenedor del servidor proxy de IFRT:

/pathways/proxy_server \
    --server_port=29000 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --num_elastic_instances=MAX_SLICES

El comando usa los siguientes argumentos:

  • --server_port: Es el puerto en el que escucha el servidor proxy de IFRT. Tu programa de JAX se conecta a este puerto.
  • --resource_manager_address: Es la dirección del administrador de recursos de Pathways. RM_IP es la dirección IP interna del nodo de CPU.
  • --gcs_scratch_location: Es la misma ruta de acceso de Cloud Storage que el administrador de recursos de Pathways.
  • --num_elastic_instances: Es la cantidad de porciones que pueden no estar disponibles sin detener la sesión del proxy de IFRT, de MAX_SLICES. Tu carga de trabajo aplica MIN_SLICES.

Cuando se ejecutan al menos MIN_SLICES porciones, el trabajo de servicio escribe la dirección del servidor proxy de IFRT en ~/pathways/cluster.env. Tu programa de JAX, el cliente proxy de IFRT, lee este archivo para conectarse a Pathways. El programa de JAX se ejecuta como tu carga de trabajo, por ejemplo, run_maxtext.sh, o en un notebook de Jupyter.

Trabajador de Pathways

El trabajo de servicio de Pathways envía un array de trabajos de Slurm (run_worker_slice.sh) con una tarea de array para cada porción de TPU. Cada tarea de array ejecuta el siguiente comando en el contenedor del servidor de Pathways en cada VM de TPU de su porción:

/pathways/server \
    --server_port=29001 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

El comando usa los siguientes argumentos:

  • --server_port: Es el puerto en el que escucha el trabajador de Pathways. El administrador de recursos de Pathways y el servidor proxy de IFRT se conectan a este puerto.
  • --resource_manager_address: Es la dirección del administrador de recursos de Pathways, la misma que para el servidor proxy de IFRT.
  • --gcs_scratch_location, --enable_gcs_olm: Es lo mismo que para el administrador de recursos de Pathways.

La secuencia de comandos también establece la topología de TPU y las variables de entorno MEGASCALE_* para cada trabajador a partir de los nodos de la porción y el índice de tareas del array.

Si se interrumpe o falla una porción de TPU, Slurm vuelve a poner en cola su tarea de array. Si tienes SSH sin contraseñasudo en los nodos del clúster, el trabajo de servicio también apaga los nodos de la porción para que Slurm pueda volver a crearlos. El administrador de recursos de Pathways y el servidor proxy de IFRT siguen ejecutándose, por lo que tu programa de JAX no necesita conectarse a una dirección nueva.

¿Qué sigue?