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 de forma dinámica 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 organice las cargas de trabajo en varias porciones de TPU grandes, lo que puede abarcar miles de chips TPU.

Pathways se usa de forma interna en Google para entrenar modelos grandes, como Gemini. Pathways on Cloud ofrece 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 por lotes, interactivas y en tiempo real. Se supone que ya estás familiarizado con el uso de TPU con GKE, incluidas las TPU de una sola porción y las TPU de varias porciones en Google Kubernetes Engine, así como con la 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

Controlador único (rutas de aprendizaje)

Multi-controller (JAX Default)

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 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 aparezca como una sola máquina grande con muchos aceleradores locales.

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

Flexibilidad

Admite patrones de procesamiento más 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 la siguiente sección, se describen los componentes principales de la arquitectura de Pathways.

Administrador de recursos de Pathways

Este es el plano de control central del sistema Pathways. Administra todos los recursos de aceleradores 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 punto de contacto único para los errores y el estado del sistema. Este componente solo requiere recursos de CPU.

Cliente de Pathways

Esta es una implementación del entorno de ejecución del framework provisional (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 coordinar 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 de acelerador.

Cliente proxy de IFRT

Esta es una implementación de OSS 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. Es una implementación de OSS de la API de IFRT. Este componente solo requiere recursos de CPU.

Servidor proxy de IFRT

Este servidor de 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 sidecar

Este servidor gRPC se ubica junto con el trabajador de Pathways en la VM del acelerador para ejecutar el código de Python especificado por el usuario en la VM del acelerador directamente y 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 Pathways.
Componentes de Pathways

Componentes de Pathways en GKE

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

Puedes encontrar 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/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 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/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 debe establecerse en "resource_manager" para el administrador de recursos de Pathways y 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 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 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 --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 TPU para comunicarse con el administrador de recursos de Pathways. El puerto debe ser el mismo que el valor --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?