Portare i carichi di lavoro JAX su Pathways

A causa della natura distribuita di JAX con Pathways, alcune operazioni potrebbero non scalare bene a causa dell'overhead di comunicazione. Sebbage Pathways riduca al minimo questi overhead con funzionalità come l'invio asincrono, ci sono alcune cose di cui devi tenere conto quando esegui il porting dei workload JAX su Pathways o quando scali un workload JAX con Pathways a un numero elevato di acceleratori.

Prima di iniziare

Assicurati di avere:

Indice del processo

JAX con Pathways tratta tutti i dispositivi del cluster Pathways come locali. In questo modo, la gestione dei dispositivi viene semplificata e JAX può utilizzare tutte le risorse disponibili. In pratica, questo significa che:

  • jax.process_index() è sempre 0 per tutti i dispositivi.
  • jax.devices() e jax.local_devices() restituiscono tutti i dispositivi TPU nell'intero job.

Tipo di hardware e colocation

Per prestazioni ottimali, posiziona tutti i componenti Pathways e il job utente nella stessa zona cloud Google Cloud . Utilizza una CPU di grandi dimensioni come il proxy IFRT e il gestore delle risorse. Consigliamo almeno un n2-standard-64 dedicato, dotato di 64 vCPU e 256 GB di memoria.

PathwaysUtils

Pathways-utils è un repository GitHub basato su Python che fornisce utilità e strumenti essenziali che ti consentono di semplificare il deployment e l'esecuzione dei workload JAX sull'architettura Pathways on Cloud. Questo pacchetto gestisce gli adattamenti necessari per l'ambiente cloud, consentendo agli sviluppatori JAX di concentrarsi sui flussi di lavoro di machine learning principali con una configurazione minima specifica per la piattaforma. Nello specifico, offre:

  • Un backend JAX "proxy": questo backend personalizzato consente alla tua applicazione JAX di utilizzare l'infrastruttura Pathways impostando la variabile di ambiente JAX_PLATFORMS=proxy.
  • Utilità di profilazione integrate: funzionalità di profilazione che ti consentono di comprendere le prestazioni della tua applicazione. Utilizzando le API di profilazione JAX standard come jax.profiler.start_trace e jax.profiler.start_server, puoi profilare non solo il codice JAX, ma anche i componenti Pathways sottostanti, fornendo una visione olistica dell'esecuzione all'interno dell'ambiente cloud.
  • Distributed Checkpointing con Orbax: un gestore di checkpoint Orbax personalizzato che ti consente di utilizzare i checkpoint distribuiti e di ripristinarli quando utilizzi la libreria Orbax nell'ambiente Pathways. Questa integrazione mira a funzionare senza richiedere modifiche al codice di checkpointing Orbax esistente, a condizione che importi pathwaysutils.
  • Primitive di addestramento elastico: fornisce primitive di addestramento elastico di base che puoi utilizzare per creare workflow di addestramento solidi e scalabili utilizzando Pathways. Queste primitive consentono ai job di addestramento di adattarsi dinamicamente alle modifiche delle risorse disponibili, migliorando l'efficienza e la resilienza negli ambienti cloud.

Checkpointing

Orbax è stato testato accuratamente con Pathways per il checkpointing e il ripristino distribuiti con Cloud Storage. Quando imposti la variabile di ambiente ENABLE_PATHWAYS_PERSISTENCE=1 e chiami import pathwaysutils; pathwaysutils.initialize() nel tuo train.py, viene registrato un ArrayHandler personalizzato che gestisce in modo efficiente le operazioni di checkpoint tramite il proxy IFRT, consentendo ai worker Pathways sugli acceleratori di salvare e ripristinare direttamente i dati.

Python collocato

Python in co-locazione è un'API JAX open source che ti consente di eseguire il codice Python specificato dall'utente direttamente sugli host TPU o GPU, il che è più semplice in JAX multi-controller. Utilizzando questo approccio, puoi evitare il trasferimento di dati tra il client e le macchine TPU per attività ad alta intensità di calcolo, come il caricamento dei dati e il checkpoint. Per configurare il cluster Pathways in modo che esegua l'API JAX Python collocata, crea ed esegui il push di un'immagine container sidecar utilizzando il Dockerfile sidecar di esempio.

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente i batch da un set di dati per inserirli nel modello. Avere un caricatore di dati asincrono ed efficiente che suddivide il batch tra gli host è importante per evitare di privare gli acceleratori di lavoro. Quando esegui l'addestramento con Pathways, il caricatore di dati viene eseguito su una VM CPU (a differenza di una VM TPU utilizzata nelle configurazioni multi-controller) e invia i dati alle VM TPU. Ciò comporta una latenza maggiore nella lettura dei dati, che viene però parzialmente mitigata dalla lettura anticipata di un numero X di batch sull'host della CPU e dall'invio asincrono dei dati letti alle TPU. Questa soluzione è sufficiente per l'esecuzione su scala medio-piccola.

Quando esegui un caricatore di dati multiprocesso (come Grain) nel container del workload head su GKE, il container condivide il nodo head con i sidecar pathways-proxy e pathways-rm. A causa di questa architettura condivisa, devi configurare quanto segue:

  • Imposta un numero esplicito di worker: all'interno di un container, os.cpu_count() restituisce le vCPU totali della VM host anziché il limite di CPU del container. Passa un numero esplicito di worker (non superiore al limite di CPU del container) anziché derivarlo da os.cpu_count().
  • Monta un volume /dev/shm basato sulla memoria: i caricatori di dati multiprocesso utilizzano la memoria condivisa POSIX (/dev/shm) per la comunicazione tra processi. Assicurati che un volume emptyDir (medium: Memory) basato sulla memoria sia montato su /dev/shm in modo che i batch precaricati non superino il limite di memoria condivisa del container predefinito.

Per prestazioni ottimali su larga scala, ti consigliamo vivamente di collocare la pipeline di dati di input utilizzando Python collocato per eseguire la pipeline di dati direttamente sugli acceleratori. In questo modo, si elimina il collo di bottiglia della CPU e si sfruttano le interconnessioni rapide della TPU per il trasferimento dei dati.

Puoi trovare un'implementazione di riferimento della migrazione di una pipeline di input basata su TFDS nell'implementazione RemoteIterator in multihost_dataloading.py. Questa implementazione funziona sia su JAX multi-controller sia su Pathways in modo distribuito utilizzando l'API Python JAX collocata.

Controllo delle versioni di Jax

Le release di Pathways sono strettamente accoppiate alle versioni di JAX per garantire compatibilità e stabilità. Per evitare potenziali problemi, verifica che gli artefatti Pathways e la versione di JAX siano allineati. Ogni release di Pathways specifica chiaramente le versioni JAX compatibili tramite un tag nel formato jax-<version>.

Cache di compilazione

La cache di compilazione persistente di Pathways è una funzionalità che consente ai server Pathways di archiviare gli eseguibili XLA compilati in una posizione persistente, come Cloud Storage, per evitare la compilazione ridondante. Questa funzionalità è attivata per impostazione predefinita. La posizione della cache viene passata come flag --gcs_scratch_location al gestore delle risorse e ai container worker di Pathways. Per ridurre al minimo i costi di archiviazione associati, la cache associa una policy del ciclo di vita alla località di Cloud Storage. Esiste un limite di 50 policy per bucket Cloud Storage. Pertanto, ti consigliamo di utilizzare una posizione Cloud Storage comune per tutti i workload.

Se ometti il flag --gcs_scratch_location o lo lasci vuoto, Pathways utilizza automaticamente una cache di compilazione in memoria. Questa cache non conserva gli eseguibili XLA compilati tra i deployment di Pathways.

Questa cache è simile alla cache di compilazione JAX, che viene disabilitata da pathwaysutils.initialize() per i workload Pathways.

Per la cache di compilazione sono necessarie le seguenti autorizzazioni di Cloud Storage:

  • storage.buckets.get: per recuperare i metadati del bucket.
  • storage.buckets.update: essenziale per Pathways per configurare le policy del ciclo di vita degli oggetti per applicare il TTL per l'eliminazione della cache.
  • storage.objects.list: per elencare gli oggetti della cache esistenti nel bucket.
  • storage.objects.create: per scrivere nuovi eseguibili compilati nella cache.
  • storage.objects.get: per leggere gli eseguibili memorizzati nella cache dal bucket.

Profilazione

Puoi utilizzare il profiler JAX per generare tracce di un programma JAX. Esistono due modi comuni supportati da Pathways:

  • Pubblicità programmatica
    • Acquisisci i profili in modo programmatico dal codice JAX
  • Manuale
    • Acquisizione di profili on demand dopo l'avvio del server del profiler dal codice JAX

In entrambi i casi, i profili vengono scritti in un bucket Cloud Storage. Nel bucket Cloud Storage verranno creati più file di traccia, potenzialmente in cartelle con timestamp diversi, ad esempio:

  • Processo Python principale che ha richiamato la traccia (in genere la VM del blocco note): <jax-client-vm-name>.xplane.pb
  • Proxy IFRT Pathways: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Pathways Resource manager: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Worker Pathways: server.*<tpu-node-name>.xplane.pb

Questi file di traccia possono essere analizzati con TensorBoard eseguendo il seguente comando. Per saperne di più su TensorBoard e su tutti i suoi strumenti di profilazione, consulta Ottimizza le prestazioni di TensorFlow utilizzando Profiler.

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

Sostituisci quanto segue:

  • BUCKET : un bucket Cloud Storage per archiviare i file di traccia
  • PREFIX: un percorso all'interno del bucket Cloud Storage in cui archiviare i file di traccia

Acquisizione del profilo di programmazione

Acquisisci un profilo dall'interno del codice. I profili vengono salvati all'interno di gs://<bucket>/<prefix> in una directory con timestamp

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

Acquisizione manuale del profilo

Per acquisire manualmente le informazioni del profilo, devi avviare il server del profiler dal codice Python:

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

Mentre il server del profiler è in esecuzione, puoi acquisire un profilo ed esportare i dati nella località Cloud Storage di destinazione:

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

Puoi trovare informazioni sui tempi per i metodi client del proxy IFRT come Compile e Execute nella traccia del programma. Questi eventi, che descrivono in dettaglio le interazioni con il server gRPC del proxy IFRT durante la compilazione e l'esecuzione, vengono visualizzati nel thread denominato GrpcClientSessionUserFuturesWorkQueue. Esaminando questo thread nella traccia, puoi ottenere informazioni sul rendimento di queste operazioni.

Flag XLA

Quando utilizzi Pathways, devi impostare i flag XLA nel container pathways-proxy.

Quando utilizzi Cluster Toolkit (il comando gcluster job submit), imposta i flag XLA utilizzando --pathways-proxy-args:

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

Dump HLO

Per ispezionare gli input dell'ottimizzatore di alto livello (HLO) forniti al compilatore XLA, puoi configurare Pathways in modo che esegua il dump dell'HLO in una posizione Cloud Storage specificata utilizzando --pathways-proxy-args:

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

Passaggi successivi