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 a gran escala, de múltiples tareas y activados de forma dispersa. Permite el uso de miles o decenas de miles de aceleradores, con la capacidad de asignar dinámicamente cantidades variables de procesamiento para diferentes tareas según sus requisitos de procesamiento.

Pathways simplifica los cálculos de aprendizaje automático a gran escala, ya que permite que un solo cliente de JAX coordine cargas de trabajo en varias porciones de TPU grandes, que pueden abarcar miles de chips de TPU.

En Google, Pathways se usa de forma interna para entrenar modelos grandes, como Gemini. Las rutas de aprendizaje en Cloud brindan los mismos beneficios a los clientes de Google Cloud .

Antes de comenzar

Asegúrate de tener lo siguiente:

En este documento, se proporciona una descripción general de cómo usar las TPU administradas por Pathways en Google Kubernetes Engine (GKE) para cargas de trabajo interactivas, en tiempo real y por lotes. Se supone que ya sabes cómo usar las TPU con GKE, incluidas las TPU de una sola porción y las de varias porciones en Google Kubernetes Engine, así como que tienes experiencia general con las TPU de varias porciones.

Controlador único y varios controladores

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

Función

Un solo controlador (Pathways)

Multi-controller (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 del intérprete de Python.

Ver

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

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

Programación

Programación simplificada: Los usuarios interactúan con un solo cliente, lo que hace que el sistema parezca una sola máquina grande con muchos aceleradores locales.

SPMD: Usa principalmente el paradigma SPMD, que requiere que todos los dispositivos ejecuten el mismo programa.

Flexibilidad

Admite patrones de procesamiento más complejos que SPMD, incluidos 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 segmentos de TPU.

Componentes de rutas de aprendizaje

En la siguiente sección, se describen los componentes principales de la arquitectura de Pathways.

Administrador de recursos de rutas de aprendizaje

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

Cliente de Pathways

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

Trabajador de Pathways

Estos son los procesos que 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 del acelerador.

Cliente proxy de IFRT

Esta es una implementación de código abierto de la API de Interim Framework Runtime (IFRT), que desacopla el código del usuario del tiempo 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 predeterminado de varios controladores. 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. Es una implementación de código abierto de la API de IFRT. Este componente solo requiere recursos de CPU.

Servidor proxy de IFRT

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

Servidor de sidecar

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

Muestra la relación de los componentes de Rutas de aprendizaje.
Componentes de rutas de aprendizaje

Componentes de Pathways en GKE

En esta sección, se asignan los componentes de Pathways a los componentes de Google Kubernetes Engine, como los contenedores y los pods.

Puedes encontrar imágenes de contenedores de Pathways en las siguientes ubicaciones.

Tipo de contenedor

Ubicación

Servidor proxy de IFRT

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>

Administrador o trabajador de recursos de Pathways

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

Administrador de recursos de rutas de aprendizaje

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

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/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

Descripciones de los argumentos:

  • --server_port: El administrador de recursos de Pathways usa este puerto para comunicarse con otros componentes de Pathways.
  • --node_type: Es el tipo de nodo. Este campo debe establecerse en "resource_manager" para el administrador de recursos de Pathways y no es necesario para los demás 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 tpu{TPU type}:{TPU 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 objeto containerSpec para implementar un servidor proxy de IFRT:

 - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/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

Descripciones de los 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 ser el mismo que el valor de --server_port que se usa para el contenedor del administrador de recursos de Pathways.
  • --server_port: El servidor proxy de IFRT usa este puerto 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/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"

Descripciones de los 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 ser el mismo que el valor de --server_port que se usa para el contenedor del administrador de recursos de Pathways.
  • --server_port: Los trabajadores usan este puerto 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.

¿Qué sigue?