Introduzione ai percorsi su Cloud

Pathways è un sistema progettato per consentire la creazione di sistemi di machine learning su larga scala, multi-task e attivati in modo sparso. Consente l'utilizzo di migliaia o decine di migliaia di acceleratori, con la possibilità di allocare dinamicamente quantità variabili di risorse di calcolo per attività diverse in base ai relativi requisiti di elaborazione.

Pathways semplifica i calcoli di machine learning su larga scala consentendo a un singolo client JAX di orchestrare i carichi di lavoro su più slice TPU di grandi dimensioni, potenzialmente su migliaia di chip TPU.

Pathways viene utilizzato internamente in Google per addestrare modelli di grandi dimensioni come Gemini. Pathways su Cloud offre gli stessi vantaggi ai Google Cloud clienti.

Prima di iniziare

Assicurati di avere:

• Installato gli strumenti Kubernetes
• Installato gcloud CLI
• Abilitato l'API TPU
• Abilitato l'API Google Kubernetes Engine

Questo documento fornisce una panoramica su come utilizzare le TPU gestite da Pathways su Google Kubernetes Engine (GKE) per carichi di lavoro batch, in tempo reale e interattivi. Presuppone che tu abbia già familiarità con l'utilizzo delle TPU con GKE incluse le TPU a slice singola e multipla su Google Kubernetes Engine, nonché un'esperienza generale con le TPU a slice multipla.

Controller singolo e multi-controller

Esistono principalmente due modi diversi per gestire e orchestrare i calcoli su più dispositivi:

Funzionalità	Controller singolo (Pathways)	Multi-controller (JAX predefinito)
Controllo	Punto di controllo singolo: un singolo programma client funge da controller centrale.	Controllo distribuito: partecipano più processi, ognuno con la propria istanza dell'interprete Python.
Visualizza	Visualizzazione unificata: il client vede tutti i dispositivi come un unico sistema unificato.	Visualizzazione localizzata: ogni processo Python vede solo i dispositivi a cui è connesso.
Programmazione	Programmazione semplificata: gli utenti interagiscono con un singolo client, facendo apparire il sistema come una singola macchina di grandi dimensioni con molti acceleratori locali.	SPMD: utilizza principalmente il paradigma SPMD, che richiede che tutti i dispositivi eseguano lo stesso programma.
Flessibilità	Supporta pattern di calcolo più complessi oltre a SPMD, tra cui il parallelismo della pipeline asimmetrica e la sparsità computazionale.	Può essere meno flessibile nella gestione delle risorse, soprattutto tra diverse slice TPU.

Componenti di Pathways

La sezione seguente descrive i componenti principali dell'architettura di Pathways.

Gestore risorse di Pathways

Questo è il piano di controllo centrale del sistema Pathways. Gestisce tutte le risorse dell'acceleratore ed è responsabile del coordinamento dell'allocazione degli acceleratori per i job utente. Monitora l'integrità dei worker e gestisce la pianificazione, la sospensione e la ripresa dei job. Funge da punto di contatto unico per gli errori e lo stato del sistema. Questo componente richiede solo risorse CPU.

Client Pathways

Si tratta di un'implementazione di Interim Framework Runtime (IFRT) che funge da punto di accesso al sistema Pathways. Riceve operazioni di alto livello (HLO) dal tuo programma. Il client Pathways è responsabile del coordinamento con il gestore risorse Pathways per determinare dove inserire i programmi compilati per l'esecuzione in base al codice utente. Presenta una visualizzazione unificata del sistema a un determinato client JAX. Questo componente richiede solo risorse CPU.

Worker Pathways

Questi sono i processi eseguiti sulle macchine dell'acceleratore (VM TPU). Ricevono eseguibili compilati del tuo programma dal server proxy IFRT ed eseguono i calcoli sulle TPU. I worker Pathways inviano i dati al tuo programma tramite il server proxy IFRT. Questo componente richiede risorse dell'acceleratore.

Client proxy IFRT

Si tratta di un'implementazione OSS dell'API Interim Framework Runtime (IFRT) che disaccoppia il codice utente dal runtime sottostante e migliora la portabilità e la trasparenza del codice. JAX utilizza questa implementazione come alternativa al runtime multi-controller predefinito. Il client proxy IFRT funge da ponte di comunicazione tra il tuo programma e i componenti Pathways. Invia richieste al server proxy IFRT e riceve i risultati. Si tratta di un'implementazione OSS dell'API IFRT. Questo componente richiede solo risorse CPU.

Server proxy IFRT

Questo server gRPC riceve le richieste dal client proxy IFRT e le inoltra al client Pathways, che gestisce la distribuzione effettiva del lavoro. Questo componente richiede solo risorse CPU.

Server sidecar

Questo server gRPC si trova nella stessa posizione del worker Pathways sulla VM dell'acceleratore per eseguire direttamente il codice Python specificato dall'utente sulla VM dell'acceleratore al fine di ridurre la latenza di trasferimento dei dati dal controller agli acceleratori. Il server sidecar interagisce con il worker Pathways tramite un protocollo con controllo delle versioni personalizzato sul trasporto gRPC.

Componenti di Pathways

API PathwaysJob

L'API PathwaysJob è un'API OSS nativa di Kubernetes che utilizzi per eseguire il deployment dei carichi di lavoro di addestramento ML e di inferenza batch. Il controller per PathwaysJob utilizza l'API JobSet per gestire il ciclo di vita e il coordinamento di tutti i componenti Pathways. Questa definizione di risorsa personalizzata (CRD) fornisce un'interfaccia di alto livello per definire i carichi di lavoro Pathways, eliminando la necessità di gestire direttamente le singole specifiche dei pod per gli scenari comuni. Per un elenco completo di tutti i parametri e dei relativi significati specifici, consulta la documentazione dell'API PathwaysJob su 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())'

La tabella seguente descrive le impostazioni dell'API PathwaysJob:

Attributo	Descrizione
apiVersion	Specifica la versione dell'API per la definizione di risorsa personalizzata (CRD) PathwaysJob: pathways-job.pathways.domain/v1.
kind	Identifica l'oggetto Kubernetes come PathwaysJob.
metadata.name	Il nome dell'oggetto PathwaysJob in Kubernetes, in genere seguendo il pattern pathways-.
spec	Definisce lo stato e la configurazione desiderati per PathwaysJob.
spec.maxRestarts	Il numero massimo di volte in cui PathwaysJob può essere riavviato automaticamente dal sistema in caso di errori.
spec.pathwaysVersion	(Facoltativo) Specifica la versione desiderata del framework JAX da utilizzare nell'ambiente Pathways per questo job (ad esempio, jax-0.5.3).
spec.workers	Un array che definisce la configurazione del pool di worker di PathwaysJob, in genere utilizzando le risorse TPU.
spec.workers[].type	Il tipo di macchina TPU da utilizzare per i nodi worker (ad esempio, $TPU_MACHINE_TYPE potrebbe essere ct6e-standard-4t)
spec.workers[].topology	La topologia delle slice TPU allocate ai worker (ad esempio, $TOPOLOGY potrebbe essere 2x2, 4x4, 2x2x2).
spec.workers[].numSlices	Il numero di slice TPU da eseguire il provisioning per il pool di worker (ad esempio, $WORKLOAD_NODEPOOL_COUNT potrebbe essere 2).
spec.workers[].maxSliceRestarts	(Facoltativo) Il numero massimo di volte in cui un singolo worker all'interno di una slice può essere riavviato in caso di errore.
spec.customComponents	(Facoltativo) Un array che consente di definire ed eseguire il deployment di componenti personalizzati (come server proxy, server Pathways o worker aggiuntivi) insieme al job principale.
spec.customComponents[].componentType	Specifica il tipo di componente personalizzato in fase di definizione (ad esempio, proxy_server, pathways_server, worker, colocated_python_sidecar).
spec.customComponents[].image	L'immagine Docker da utilizzare per il container di questo componente personalizzato.
spec.customComponents[].customFlags	Un array di flag della riga di comando personalizzati che verranno passati al container all'avvio.
spec.customComponents[].customEnv	Un array di variabili di ambiente personalizzate da impostare all'interno del container. Ogni elemento ha un nome e un valore.
spec.pathwaysDir	Il bucket Cloud Storage utilizzato da PathwaysJob per archiviare gli artefatti di compilazione e altri dati temporanei. Questo bucket deve essere creato prima di eseguire il carico di lavoro.
spec.controller	Impostazioni di configurazione del controller Pathways, che gestisce l'esecuzione complessiva del job.
spec.controller.deploymentMode	Specifica la modalità di deployment delle risorse CPU del controller Pathways (gestore risorse Pathways e server proxy). La modalità predefinita le esegue il deployment su un nodo CPU dedicato mentre colocate_head_with_workers le esegue il deployment insieme a un worker TPU.
spec.controller.elasticSlices	(Facoltativo) Il numero massimo di slice TPU che possono diventare non disponibili durante l'esecuzione del job prima che venga considerato non integro.
spec.controller.template	(Facoltativo) Definisce il modello di pod per il job utente. Questo è obbligatorio per i carichi di lavoro batch, ma non per i carichi di lavoro interattivi.
spec.controller.template.spec	La specifica del pod per il job utente.
spec.controller.template.spec.containers	Un array che definisce i container che verranno eseguiti all'interno del job utente.
spec.controller.template.spec.containers[].name	Il nome del container all'interno del job utente (in questo esempio, è main).
spec.controller.template.spec.containers[].image	L'immagine Docker da utilizzare per il container nel container principale (in questo esempio, è python:3.11).
spec.controller.template.spec.containers[].command	Il comando da eseguire all'avvio del container principale. In questo esempio, installa `pathwaysutils`, inizializza Pathways e stampa i dispositivi JAX.

Componenti di Pathways su GKE

Questa sezione mappa i componenti di Pathways ai componenti di Google Kubernetes Engine, come container e pod.

Nota: le specifiche dei container descritte in questa sezione sono esempi e vengono utilizzate solo a scopo dimostrativo. I container sottostanti vengono astratti e gestiti tramite l'API della definizione di risorsa personalizzata (CRD) PathwaysJob. Per configurare e gestire questi container, consulta l' API PathwaysJob. La manipolazione dei singoli flag dei container, delle variabili di ambiente e delle posizioni delle immagini è supportata tramite la specifica customComponents.

Puoi trovare le immagini container di Pathways nelle seguenti posizioni.

Tipo di container	Località
Server proxy IFRT	us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>
Gestore risorse/worker Pathways	us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

Gestore risorse di Pathways

Dopo aver creato un cluster GKE, puoi utilizzare il seguente containerSpec per eseguire il deployment del gestore risorse 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

Descrizioni degli argomenti:

• --server_port: il gestore risorse Pathways utilizza questa porta per comunicare con altri componenti Pathways.
• --node_type: il tipo di nodo. Per il gestore risorse Pathways, questo valore deve essere impostato su "resource_manager" e non è necessario per gli altri container.
• --instance_count: il numero di slice TPU.
• --instance_type: il tipo di TPU e la topologia della slice. Nel formato tpu{TPU type}:{TPU topology}, ad esempio tpuv5e:4x4.
• --gcs_scratch_location: un bucket Cloud Storage utilizzato per i file temporanei.

Server proxy IFRT

Puoi utilizzare il seguente containerSpec per eseguire il deployment di un server proxy 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

Descrizioni degli argomenti:

• --resource_manager_address: il nome host e la porta utilizzati dal server proxy per comunicare con il gestore risorse Pathways. La porta deve essere uguale al valore --server_port utilizzato per il container del gestore risorse Pathways.
• --server_port: il server proxy IFRT utilizza questa porta per comunicare con il client proxy IFRT.
• --gcs_scratch_location: un bucket Cloud Storage utilizzato per i file temporanei.

Worker Pathways

Puoi utilizzare il seguente containerSpec per eseguire il deployment dei worker 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"

Descrizioni degli argomenti:

• --resource_manager_address: il nome host e la porta utilizzati dai worker TPU per comunicare con il gestore risorse Pathways. La porta deve essere uguale al valore --server_port utilizzato per il container del gestore risorse Pathways.
• --server_port: i worker utilizzano questa porta per comunicare con il server proxy e il gestore risorse Pathways.
• --gcs_scratch_location: un bucket Cloud Storage utilizzato per i file temporanei.

Il gestore risorse Pathways, il server proxy IFRT e i worker Pathways possono avere porte diverse, ma in questo esempio il gestore risorse Pathways e il worker Pathways condividono la stessa porta.

Passaggi successivi

• Crea un cluster GKE con Pathways
• Esegui un carico di lavoro batch con Pathways
• Esegui l'inferenza multihost utilizzando Pathways
• Esegui un carico di lavoro interattivo con Pathways
• Addestramento resiliente con Pathways
• Porta i carichi di lavoro JAX a Pathways
• Risolvi i problemi di Pathways su Cloud