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 calcolo per diverse attività in base ai loro 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 da Google per addestrare modelli di grandi dimensioni come Gemini. I percorsi su Cloud offrono gli stessi vantaggi ai clienti di Google Cloud .

Prima di iniziare

Assicurati di avere:

Questo documento fornisce una panoramica di come utilizzare le TPU gestite 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 single-slice e multi-slice su Google Kubernetes Engine, nonché esperienza generale con le TPU multi-slice.

Singolo controller e più controller

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

Funzionalità	Singolo controller (percorsi)	Multi-controller (JAX Default)
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 a cui è connesso.
Programmazione	Programmazione semplificata: gli utenti interagiscono con un singolo client, facendo apparire il sistema come un'unica macchina di grandi dimensioni con molti acceleratori locali.	SPMD: utilizza principalmente il paradigma SPMD, che richiede l'esecuzione dello stesso programma su tutti i dispositivi.
Flessibilità	Supporta pattern di calcolo più complessi oltre a SPMD, tra cui parallelismo della pipeline asimmetrico e sparsità computazionale.	Può essere meno flessibile nella gestione delle risorse, soprattutto tra diverse sezioni TPU.

Componenti di Pathways

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

Resource Manager di Pathways

Si tratta del control plane 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 unico punto di contatto per errori e 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 delle 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.

Operatore di Pathways

Si tratta dei processi eseguiti sulle macchine acceleratrici (VM TPU). Ricevono eseguibili compilati del tuo programma dal server proxy IFRT ed eseguono i calcoli sulle TPU. I worker di Pathways inviano i dati al tuo programma tramite il server proxy IFRT. Questo componente richiede risorse di accelerazione.

Client proxy IFRT

Si tratta di un'implementazione OSS dell'API Interim Framework Runtime (IFRT) che separa 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 di Pathways. Invia richieste al server proxy IFRT e riceve i risultati. È 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 insieme al worker Pathways sulla VM dell'acceleratore per eseguire il codice Python specificato dall'utente direttamente sulla VM dell'acceleratore, in modo da 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.

Mostra la relazione tra i componenti di Pathways. — Componenti dei percorsi

API PathwaysJob

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

Attributo	Descrizione
`apiVersion`	Specifica la versione 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 percorsi-.
`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 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 sezioni 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 sezione può essere riavviato in caso di errore.
`spec.customComponents`	(Facoltativo) Un array che ti 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 workload.
`spec.controller`	Impostazioni di configurazione del controller dei percorsi, che gestisce l'esecuzione complessiva del job.
`spec.controller.deploymentMode`	Specifica come vengono implementate le risorse CPU del controller Pathways (Pathways resource manager e server proxy). La modalità predefinita li esegue il deployment su un nodo CPU dedicato, mentre colocate_head_with_workers li 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 dell'utente. Questo è necessario per i workload batch, ma non per quelli 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 contenitore 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.

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/pathways/proxy_server:jax-<jax-version>`
Resource Manager/worker di Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>`

Resource Manager di Pathways

Dopo aver creato un cluster GKE, puoi utilizzare il seguente containerSpec per eseguire il deployment di Resource Manager per i percorsi:

  - 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: Resource Manager di Pathways utilizza questa porta per comunicare con altri componenti di Pathways.
--node_type: il tipo di nodo. Deve essere impostato su "resource_manager" per il gestore delle risorse di Pathways e non è necessario per gli altri contenitori.
--instance_count: il numero di sezioni TPU.
--instance_type: il tipo di TPU e la topologia della sezione. 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: l'hostname e la porta utilizzati dal server proxy per comunicare con Pathways Resource Manager. La porta deve essere la stessa del valore --server_port utilizzato per il contenitore di Pathways Resource Manager.
--server_port: il server proxy IFRT utilizza questa porta per comunicare con il client proxy IFRT.
--gcs_scratch_location: un bucket GCS utilizzato per i file temporanei.

Operatore di Pathways

Puoi utilizzare i seguenti containerSpec per eseguire il deployment dei worker di 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: l'hostname e la porta utilizzati dai worker TPU per comunicare con il gestore delle risorse Pathways. La porta deve essere la stessa del valore --server_port utilizzato per il contenitore di Pathways Resource Manager
--server_port: i worker utilizzano questa porta per comunicare con il server proxy e il gestore delle risorse di Pathways.
--gcs_scratch_location: un bucket Cloud Storage utilizzato per i file temporanei.

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

Introduzione ai percorsi su Cloud Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

Prima di iniziare

Singolo controller e più controller

Componenti di Pathways

Resource Manager di Pathways

Client Pathways

Operatore di Pathways

Client proxy IFRT

Server proxy IFRT

Server sidecar

API PathwaysJob

Componenti di Pathways su GKE

Resource Manager di Pathways

Server proxy IFRT

Operatore di Pathways

Passaggi successivi

Introduzione ai percorsi su Cloud