Introducción a las rutas de aprendizaje en Cloud

Versión preliminar

Este producto o función está con sujeción a las "Condiciones de las Ofertas de la Fase Previa a la DG" de la sección de Términos Generales del Servicio de los Términos Específicos del Servicio. Los productos y las funciones de la fase previa a la DG están disponibles sin modificaciones y pueden estar sujetos a asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

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 ofrecen los mismos beneficios a los clientes de Google Cloud .

Antes de comenzar

Asegúrate de tener lo siguiente:

• Herramientas de Kubernetes instaladas
• Instalaste la CLI de gcloud
• Habilitaste la API de TPU
• Habilitaste la API de Google Kubernetes Engine

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 (rutas)	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 funciona 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 de aceleración.

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.

Componentes de rutas de aprendizaje

API de PathwaysJob

La API de PathwaysJob es una API de OSS nativa de Kubernetes que se usa para implementar cargas de trabajo de entrenamiento de AA y de inferencia por lotes. El controlador de PathwaysJob aprovecha la API de JobSet para administrar el ciclo de vida y la coordinación de todos los componentes de Pathways. Esta Definición de Recurso Personalizado (CRD) te brinda una interfaz de alto nivel para definir tus cargas de trabajo de Rutas, lo que abstrae la necesidad de administrar directamente las especificaciones de Pod individuales para situaciones comunes. Para obtener una lista completa de todos los parámetros y sus significados específicos, consulta la documentación de la API de PathwaysJob en GitHub.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

En la siguiente tabla, se describen los parámetros de configuración de la API de PathwaysJob:

Atributo	Descripción
apiVersion	Especifica la versión de la API para la Definición de recurso personalizado (CRD) PathwaysJob: pathways-job.pathways.domain/v1.
kind	Identifica el objeto de Kubernetes como un PathwaysJob.
metadata.name	Nombre del objeto PathwaysJob en Kubernetes, que suele seguir el patrón pathways-.
spec	Define el estado y la configuración deseados para el PathwaysJob.
spec.maxRestarts	Es la cantidad máxima de veces que el sistema puede reiniciar automáticamente el PathwaysJob si se producen errores.
spec.pathwaysVersion	(Opcional) Especifica la versión deseada del framework de JAX que se usará en el entorno de Pathways para este trabajo (por ejemplo, jax-0.5.3).
spec.workers	Es un array que define la configuración del grupo de trabajadores de PathwaysJob, que suele utilizar recursos de TPU.
spec.workers[].type	Es el tipo de máquina TPU que se usará para los nodos trabajadores (por ejemplo, $TPU_MACHINE_TYPE podría ser ct6e-standard-4t).
spec.workers[].topology	Es la topología de las porciones de TPU asignadas a los trabajadores (por ejemplo, $TOPOLOGY podría ser 2x2, 4x4 o 2x2x2).
spec.workers[].numSlices	Es la cantidad de porciones de TPU que se aprovisionarán para el grupo de trabajadores (por ejemplo, $WORKLOAD_NODEPOOL_COUNT podría ser 2).
spec.workers[].maxSliceRestarts	Es la cantidad máxima de veces que se puede reiniciar un trabajador individual dentro de una división si falla (opcional).
spec.customComponents	(Opcional) Es un array que te permite definir y, luego, implementar componentes personalizados (como servidores proxy, servidores de Pathways o trabajadores adicionales) junto con el trabajo principal.
spec.customComponents[].componentType	Especifica el tipo de componente personalizado que se define (por ejemplo, proxy_server, pathways_server, worker, colocated_python_sidecar).
spec.customComponents[].image	Es la imagen de Docker que se usará para el contenedor de este componente personalizado.
spec.customComponents[].customFlags	Es un array de marcas de línea de comandos personalizadas que se pasarán al contenedor cuando se inicie.
spec.customComponents[].customEnv	Es un array de variables de entorno personalizadas que se configurarán dentro del contenedor. Cada elemento tiene un nombre y un valor.
spec.pathwaysDir	Es el bucket de Cloud Storage que usa PathwaysJob para almacenar artefactos de compilación y otros datos temporales. Este bucket debe crearse antes de ejecutar tu carga de trabajo.
spec.controller	Parámetros de configuración del controlador de rutas, que administra la ejecución general del trabajo.
spec.controller.deploymentMode	Especifica cómo se implementan los recursos de CPU del controlador de Pathways (administrador de recursos de Pathways y servidor proxy). El modo predeterminado los implementa en un nodo de CPU dedicado, mientras que colocate_head_with_workers los implementa junto con un trabajador de TPU.
spec.controller.elasticSlices	(Opcional) Es la cantidad máxima de segmentos de TPU que pueden dejar de estar disponibles durante la ejecución del trabajo antes de que se considere que no está en buen estado.
spec.controller.template	(Opcional) Define la plantilla de Pod para el trabajo del usuario. Esto es obligatorio para las cargas de trabajo por lotes, pero no para las cargas de trabajo interactivas.
spec.controller.template.spec	Es la especificación del Pod para el trabajo del usuario.
spec.controller.template.spec.containers	Es un array que define los contenedores que se ejecutarán dentro del trabajo del usuario.
spec.controller.template.spec.containers[].name	Nombre del contenedor dentro del trabajo del usuario (en este ejemplo, es main).
spec.controller.template.spec.containers[].image	Es la imagen de Docker que se usará para el contenedor en el contenedor principal (en este ejemplo, es python:3.11).
spec.controller.template.spec.containers[].command	Comando que se ejecutará cuando se inicie el contenedor principal. En este ejemplo, se instala `pathwaysutils`, se inicializa Pathways y se imprimen los dispositivos JAX.

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.

Nota: Las especificaciones del contenedor que se describen en esta sección son ejemplos y se usan solo con fines ilustrativos. Los contenedores subyacentes se abstraen y administran a través de la API de la Definición de Recurso Personalizado (CRD) de PathwaysJob. Consulta la API de PathwaysJob para configurar y administrar estos contenedores. La manipulación de marcas de contenedores individuales, variables de entorno y ubicaciones de imágenes se admite a través de la especificación customComponents.

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 de acceso:

  - 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 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 GCS 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?

• Crea un clúster de GKE con Pathways
• Ejecuta una carga de trabajo por lotes con Pathways
• Infiere varios hosts con rutas de aprendizaje
• Ejecuta una carga de trabajo interactiva con Pathways
• Capacitación resiliente con Pathways
• Cómo portar cargas de trabajo de JAX a Pathways
• Soluciona problemas de Pathways on Cloud