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 grandes de TPU, que pueden abarcar miles de chips de TPU.

Pathways se usa internamente en Google para entrenar modelos grandes como Gemini. Pathways on Cloud ofrece los mismos beneficios a los Google Cloud clientes.

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 de Pathways en Google Kubernetes Engine (GKE) para cargas de trabajo interactivas, en tiempo real y por lotes. Se supone que ya estás familiarizado con el uso de TPU con GKE incluidas las TPU de una y varias porciones en Google Kubernetes Engine, así como 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 (Pathways)	Varios controladores (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 sistema único y 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 el sistema aparezca como 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 cálculo más complejos más allá de SPMD, incluido el paralelismo de canalización asimétrica 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 del acelerador y es responsable de coordinar la asignación de aceleradores para los trabajos del usuario. Supervisa el estado de los trabajadores y controla la programación, la pausa y la reanudación de los trabajos. Funciona como un único punto de contacto 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 interino (IFRT) que sirve como punto de entrada al sistema 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 la 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 (VM 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 OSS de la API del entorno de ejecución del framework interino (IFRT) que desacopla 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 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 OSS 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 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 para reducir la latencia de transferencia de datos del controlador a los aceleradores. El servidor sidecar interactúa con el trabajador de Pathways a través de un protocolo con versiones personalizadas en el transporte gRPC.

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

API de PathwaysJob

La API de PathwaysJob es una API nativa de Kubernetes de OSS 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 recursos personalizados (CRD) te brinda una interfaz de alto nivel para definir tus cargas de trabajo de Pathways, 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 recursos personalizados (CRD) de `PathwaysJob`: pathways-job.pathways.domain/v1.
`kind`	Identifica el objeto de Kubernetes como `PathwaysJob`.
`metadata.name`	El nombre del objeto `PathwaysJob` en Kubernetes, que suele seguir el patrón pathways-.
`spec`	Define el estado y la configuración deseados para `PathwaysJob`.
`spec.maxRestarts`	La cantidad máxima de veces que el sistema puede reiniciar automáticamente `PathwaysJob` si encuentra fallas.
`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`	Un array que define la configuración del grupo de trabajadores de `PathwaysJob`, que suele utilizar recursos de TPU.
`spec.workers[].type`	El tipo de máquina de TPU que se usará para los nodos trabajadores (por ejemplo, $TPU_MACHINE_TYPE podría ser ct6e-standard-4t).
`spec.workers[].topology`	La topología de las porciones de TPU asignadas a los trabajadores (por ejemplo, $TOPOLOGY podría ser 2x2, 4x4, 2x2x2).
`spec.workers[].numSlices`	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`	(Opcional) La cantidad máxima de veces que se puede reiniciar un trabajador individual dentro de una porción si falla.
`spec.customComponents`	(Opcional) Un array que te permite definir e 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`	La imagen de Docker que se usará para el contenedor de este componente personalizado.
`spec.customComponents[].customFlags`	Un array de marcas personalizadas de línea de comandos que se pasarán al contenedor cuando se inicie.
`spec.customComponents[].customEnv`	Un array de variables de entorno personalizadas que se establecerán dentro del contenedor. Cada elemento tiene un nombre y un valor.
`spec.pathwaysDir`	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 Pathways, 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) La cantidad máxima de porciones de TPU que pueden dejar de estar disponibles durante la ejecución del trabajo antes de que se considere en mal estado.
`spec.controller.template`	(Opcional) Define la plantilla de Pod para el trabajo del usuario. Esto es necesario para las cargas de trabajo por lotes, pero no para las cargas de trabajo interactivas.
`spec.controller.template.spec`	La especificación del Pod para el trabajo del usuario.
`spec.controller.template.spec.containers`	Un array que define los contenedores que se ejecutarán dentro del trabajo del usuario.
`spec.controller.template.spec.containers[].name`	El nombre del contenedor dentro del trabajo del usuario (en este ejemplo, es main).
`spec.controller.template.spec.containers[].image`	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`	El comando que se ejecutará cuando se inicie el contenedor principal. En este ejemplo, instala `pathwaysutils`, inicializa Pathways y muestra 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 contenedores y 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/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: El tipo de nodo. Se debe establecer en "resource_manager" para el administrador de recursos de Pathways y no es necesario para los otros contenedores.
--instance_count: La cantidad de porciones de TPU.
--instance_type: El tipo de TPU y la topología de la porción. En el formato de tpu{TPU type}:{TPU topology}, por ejemplo, tpuv5e:4x4.
--gcs_scratch_location: Un bucket de Cloud Storage que se usa para 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: 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: Un bucket de Cloud Storage que se usa para 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: 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 --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: Un bucket de Cloud Storage que se usa para archivos temporales.

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

Introducción a las rutas de aprendizaje en Cloud Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.

Antes de comenzar

Controlador único y varios controladores

Componentes de Pathways

Administrador de recursos de Pathways

Cliente de Pathways

Trabajador de Pathways

Cliente proxy de IFRT

Servidor proxy de IFRT

Servidor sidecar

API de PathwaysJob

Componentes de Pathways en GKE

Administrador de recursos de Pathways

Servidor proxy de IFRT

Trabajador de Pathways

¿Qué sigue?

Introducción a las rutas de aprendizaje en Cloud