Introduzione ai percorsi su Cloud

Pathways è un sistema progettato per consentire la creazione di sistemi di machine learning (ML) su larga scala, multi-task e attivati in modo sparso. Ti consente di utilizzare migliaia o decine di migliaia di acceleratori, con la possibilità di allocare dinamicamente quantità variabili di computing per diverse attività in base ai requisiti di elaborazione.

Pathways semplifica i calcoli di ML su larga scala consentendo a un singolo client JAX di orchestrare i workload su grandi sezioni TPU, che possono potenzialmente comprendere migliaia di chip TPU.

Google utilizza Pathways internamente per addestrare modelli di grandi dimensioni come Gemini. Pathways on Cloud offre gli stessi vantaggi ai clienti di Google Cloud .

Prima di iniziare

Prima di iniziare

Puoi eseguire Pathways su Google Kubernetes Engine (GKE) o su un cluster Cluster Director con Slurm. Seleziona la scheda relativa alla tua piattaforma.

GKE

Completa i seguenti prerequisiti:

Questo documento fornisce una panoramica di come utilizzare le TPU gestite da Pathways su GKE per workload batch, in tempo reale e interattivi. Si presuppone che tu abbia già familiarità con l'utilizzo delle TPU in GKE, incluse le TPU a sezione singola e multisezione su GKE, e che tu abbia un'esperienza generale con le TPU multisezione.

Cluster Director

Completa i seguenti prerequisiti:

Questo documento fornisce una panoramica su come utilizzare Pathways con le TPU su un cluster Cluster Director per workload batch e interattivi. Si presuppone che tu abbia già familiarità con Cluster Director, con le macchine TPU in Cluster Director e con l'esecuzione di job con Slurm, oltre a un'esperienza generale con le TPU multislice.

Per creare un cluster, consulta Crea un cluster Cluster Director con Pathways.

Controller singolo e multicontroller

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

Funzionalità Controller singolo (Pathways) Multi-controller (impostazione predefinita di JAX)

Controllo

Punto di controllo unico: 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 connessi.

Programmazione

Programmazione semplificata: interagisci con un singolo client, facendo apparire il sistema come un'unica grande macchina con molti acceleratori locali.

Single Program, Multiple Data (SPMD): utilizza principalmente il paradigma SPMD, che richiede che tutti i dispositivi eseguano lo stesso programma.

Flessibilità

Supporta pattern di calcolo complessi oltre a SPMD, tra cui il parallelismo di pipeline asimmetrico e la scarsità computazionale.

Può essere meno flessibile nella gestione delle risorse, soprattutto tra diverse sezioni TPU.

Componenti di Pathways

Le sezioni seguenti descrivono i componenti principali dell'architettura Pathways.

Pathways resource manager

Questo componente è il control plane centrale del sistema Pathways. Gestisce tutte le risorse degli acceleratori e coordina l'allocazione degli acceleratori per i tuoi job. Monitora l'integrità dei worker e gestisce la pianificazione, la messa in pausa 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 di Pathways

Questo componente è 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 si coordina con il gestore delle risorse Pathways per determinare dove posizionare i programmi compilati per l'esecuzione in base al tuo codice. Presenta una visualizzazione unificata del sistema a un determinato client JAX. Questo componente richiede solo risorse CPU.

Lavoratore Pathways

I processi worker di Pathways vengono eseguiti sulle macchine acceleratrici (VM TPU). Ricevono i file 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 acceleratore.

Client proxy IFRT

Il client proxy IFRT è un'implementazione open source 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. Questo componente richiede solo risorse CPU.

Server proxy IFRT

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

Server sidecar

Il server sidecar è un server gRPC che si trova insieme al worker Pathways sulla VM dell'acceleratore. Esegue il codice Python specificato dall'utente direttamente sulla VM dell'acceleratore per ridurre la latenza del trasferimento di dati tra il controller e gli acceleratori. Il server sidecar interagisce con il worker Pathways tramite un protocollo con versione personalizzato sul trasporto gRPC.

Diagramma dell'architettura che mostra l'interazione tra il client JAX, il proxy IFRT, il gestore delle risorse Pathways e i worker Pathways.
Componenti di Pathways

Componenti Pathways su GKE

Questa sezione mappa i componenti di Pathways ai componenti di GKE, come container e pod.

Puoi trovare le immagini container di Pathways nelle seguenti posizioni:

Tipo di contenitore Località

Server proxy IFRT

us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:jax-JAX_VERSION

Pathways resource manager/worker

us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:jax-JAX_VERSION

Pathways resource manager

Dopo aver creato un cluster GKE, puoi utilizzare il seguente containerSpec per eseguire il deployment del resource manager di Pathways:

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

Il container utilizza i seguenti argomenti:

  • --server_port: la porta utilizzata dal gestore delle risorse Pathways per comunicare con altri componenti Pathways.
  • --node_type: il tipo di nodo. Imposta questo argomento su resource_manager per il gestore delle risorse Pathways; non è necessario per gli altri container.
  • --instance_count: il numero di sezioni TPU.
  • --instance_type: il tipo di TPU e la topologia dello slice, nel formato tpuTPU_TYPE: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/cloud-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

Il container utilizza i seguenti argomenti:

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

Lavoratore Pathways

Puoi utilizzare i seguenti containerSpec per eseguire il deployment dei worker Pathways:

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

Il container utilizza i seguenti argomenti:

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

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

Componenti Pathways su Cluster Director

Questa sezione mappa i componenti di Pathways ai job Slurm su un cluster Cluster Director.

Gli script Pathways eseguono i componenti Pathways con Enroot. Lo script setup_env.sh utilizza le seguenti immagini container, che corrispondono alla versione JAX (0.11.1) dell'ambiente Python che crea:

Tipo di contenitore Località

Server proxy IFRT

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20261103-jax_0.11.1

Pathways resource manager/worker

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20261103-jax_0.11.1

Per avviare i componenti Pathways, invia lo script start_pathways.sh come job del servizio Pathways. Non scrivi tu i seguenti comandi: gli script derivano gli argomenti dalle variabili di ambiente che imposti quando invii start_pathways.sh, ad esempio SLICE_SHAPE, MIN_SLICES, MAX_SLICES e GCS_BUCKET.

Pathways resource manager

Il job del servizio Pathways esegue il gestore delle risorse Pathways su un nodo CPU. Esegue il seguente comando nel contenitore del server Pathways:

/pathways/server \
    --server_port=29001 \
    --node_type=resource_manager \
    --instance_type=tpuTPU_TYPE:TOPOLOGY \
    --instance_count=MAX_SLICES \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

Il comando utilizza i seguenti argomenti:

  • --server_port: la porta su cui ascolta il gestore delle risorse Pathways. Il server proxy IFRT e i worker Pathways si connettono a questa porta.
  • --node_type: il tipo di nodo, resource_manager. Solo il resource manager di Pathways imposta questo argomento.
  • --instance_type: il tipo di TPU e la topologia di ogni sezione, da SLICE_SHAPE. Ad esempio, SLICE_SHAPE=2x4 imposta tpuv6e:2x4.
  • --instance_count: il numero di sezioni TPU, da MAX_SLICES.
  • --gcs_scratch_location: un percorso nel bucket Cloud Storage, da GCS_BUCKET, utilizzato per i file temporanei, ad esempio la cache di compilazione.
  • --enable_gcs_olm: indica se Pathways crea regole del ciclo di vita degli oggetti Cloud Storage per ripulire le voci della cache di compilazione precedenti. Lo script disattiva questa opzione perché richiede autorizzazioni aggiuntive sul bucket.

Server proxy IFRT

Il job del servizio Pathways esegue il server proxy IFRT sullo stesso nodo CPU del gestore delle risorse Pathways. Esegue il seguente comando nel contenitore del server proxy IFRT:

/pathways/proxy_server \
    --server_port=29000 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --num_elastic_instances=MAX_SLICES

Il comando utilizza i seguenti argomenti:

  • --server_port: la porta su cui il server proxy IFRT rimane in ascolto. Il tuo programma JAX si connette a questa porta.
  • --resource_manager_address: l'indirizzo del gestore delle risorse Pathways. RM_IP è l'indirizzo IP interno del nodo CPU.
  • --gcs_scratch_location: lo stesso percorso Cloud Storage del gestore delle risorse Pathways.
  • --num_elastic_instances: il numero di sezioni che possono non essere disponibili senza interrompere la sessione del proxy IFRT, da MAX_SLICES. Il tuo workload applica MIN_SLICES.

Quando sono in esecuzione almeno MIN_SLICES slice, il job del servizio scrive l'indirizzo del server proxy IFRT in ~/pathways/cluster.env. Il tuo programma JAX, il client proxy IFRT, legge questo file per connettersi a Pathways. Il programma JAX viene eseguito come workload, ad esempio run_maxtext.sh, o in un notebook Jupyter.

Lavoratore Pathways

Il job del servizio Pathways invia un array di job Slurm (run_worker_slice.sh) con un'attività di array per ogni sezione TPU. Ogni attività dell'array esegue il seguente comando nel contenitore del server Pathways su ogni VM TPU nella sua sezione:

/pathways/server \
    --server_port=29001 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

Il comando utilizza i seguenti argomenti:

  • --server_port: la porta su cui il worker Pathways è in ascolto. Il gestore delle risorse Pathways e il server proxy IFRT si connettono a questa porta.
  • --resource_manager_address: l'indirizzo del gestore delle risorse Pathways, lo stesso del server proxy IFRT.
  • --gcs_scratch_location, --enable_gcs_olm: come per il gestore delle risorse Pathways.

Lo script imposta anche la topologia TPU e le variabili di ambiente MEGASCALE_* per ogni worker dai nodi della sezione e dall'indice delle attività dell'array.

Se una sezione TPU viene prerilasciata o non funziona, Slurm mette nuovamente in coda il relativo job di array. Se hai sudo senza password sui nodi del cluster, il job di servizio spegne anche i nodi dello slice in modo che Slurm possa ricrearli. Il gestore delle risorse Pathways e il server proxy IFRT continuano a funzionare, quindi il tuo programma JAX non deve connettersi a un nuovo indirizzo.

Passaggi successivi