Formazione resiliente con Pathways

Pathways offre vantaggi in termini di resilienza nei seguenti modi:

  • Sospensione-Ripresa: consente al workload di tollerare interruzioni pianificate, come le notifiche di prerilascio, senza richiedere un codice di gestione del prerilascio personalizzato.
  • Addestramento elastico: consente al workload di tollerare errori hardware non pianificati senza causare l'arresto anomalo del client utilizzando la logica di recupero nel codice di addestramento.

Prima di iniziare

Prima di iniziare, completa le seguenti attività:

Sospensione-ripristino

Google Kubernetes Engine (GKE) in genere invia una notifica di prerilascio a un pod acceleratore prima che il pod venga prerilasciato. La tolleranza di prerilascio di Pathways è abilitata per impostazione predefinita su tutti i deployment Cloud Deploy e i job dell'acceleratore Pathways ascoltano queste notifiche.

Quando arriva una notifica di prerilascio, Pathways determina innanzitutto se può salvare e ripristinare il workload in modo trasparente. In questo caso, tenta di sospendere in modo trasparente il workload ML scrivendo lo stato attuale nello spazio di archiviazione permanente, ad esempio Cloud Storage, prima che GKE rimuova i job dell'acceleratore. Quando GKE ripianifica i job in un secondo momento, Pathways riprende il workload ML leggendo lo stato persistente.

Se il workload non è ripristinabile, Pathways arresta il job dell'acceleratore e inoltra l'errore al tuo job se è configurato l'addestramento elastico. Se l'addestramento elastico non è configurato, GKE riavvia l'intero workload in base alla policy di riavvio del JobSet.

I workload ML tipici definiti utilizzando JAX si basano su componenti Pathways XLA stateless che possono essere ripristinati utilizzando uno snapshot della memoria a larghezza di banda elevata (HBM). Alcuni workload ML, come quelli definiti utilizzando l'API Python collocata di JAX, si basano su componenti Pathways stateful; questi workload non sono ripristinabili.

Addestramento elastico

L'addestramento elastico consente al job di addestramento di continuare quando si verificano errori hardware o prerilasci di sezioni, senza riavviare l'intero workload. In un job di addestramento multi-slice, Pathways mantiene attivi il processo client del controller singolo e le sezioni TPU integre quando una singola sezione non è disponibile.

L'addestramento elastico combina le funzionalità del sistema Pathways con la logica di recupero nel codice di addestramento:

  • Rilevamento degli errori: quando si verifica un errore hardware o una prerilascio su una sezione TPU, Pathways annulla l'assegnazione della sezione interessata e genera un jax.errors.JaxRuntimeError la volta successiva che il workload accede ai dati su quella sezione. Poiché il processo client rimane attivo, il codice può rilevare questo errore, riconfigurare le risorse e continuare l'addestramento.
  • Modalità di recovery: la classe pathwaysutils.elastic.manager.Manager e il relativo decoratore elastic_retry supportano due strategie di recovery:
    • Metti in pausa e riprendi: quando minimum_slice_count è None o uguale al numero totale di sezioni nel job, l'addestramento viene messo in pausa quando una sezione viene persa e attende che tutte le sezioni siano di nuovo integre prima di riprendere. Poiché il numero di sezioni non cambia tra i tentativi, il workload può riutilizzare il mesh di dispositivi esistente e gli eseguibili compilati.
    • Ridimensionamento delle repliche: quando minimum_slice_count è inferiore al numero totale di slice, l'addestramento continua immediatamente sugli slice attivi rimanenti ricostruendo il mesh del dispositivo e ricompilando per la topologia più piccola. Un thread di monitoraggio in background controlla periodicamente gli slice inattivi e registra i nuovi slice integri in elastic_manager.available_inactive_slices. Il tuo loop di addestramento può quindi ricostruire il mesh del dispositivo con gli slice sostitutivi, reinizializzare per la topologia espansa e posizionare lo stato attuale sul nuovo mesh.
  • Recupero dello stato: per riprendere l'addestramento dopo un evento elastico, il workload ripristina i parametri del modello, lo stato dell'ottimizzatore e lo stato dell'iteratore dei dati utilizzando snapshot in memoria. La classe Snapshotter in orbax.checkpoint copia in modo asincrono lo stato replicato dalla memoria HBM della TPU alla memoria della CPU pinned_host su ogni sezione. Durante il recupero, Snapshotter ricostruisce lo stato dalle repliche host sopravvissute e lo ripartisce direttamente nella mesh TPU attiva senza leggere dallo spazio di archiviazione remoto.

Abilita l'elasticità sul server proxy Pathways

Per abilitare i set di dispositivi elastici su Pathways in GKE, aggiungi il flag --num_elastic_instances al campo args del container pathways-proxy nel manifest JobSet e specifica il numero di sezioni TPU elastiche nel tuo workload:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

Sostituisci quanto segue:

  • BUCKET_NAME: il bucket Cloud Storage utilizzato per i file temporanei di Pathways.
  • NUM_SLICES: il numero di sezioni TPU nel workload.

Filtra i dispositivi attivi e crea la rete di dispositivi

In un workload Pathways, jax.devices() restituisce tutti i dispositivi nell'allocazione iniziale per la durata del processo client. Quando utilizzi l'addestramento elastico, utilizza elastic_manager.active_devices in modo che la mesh del dispositivo e le dimensioni del batch utilizzino solo slice attivi e integri:

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

Quando il numero di sezioni attive cambia in modalità di ridimensionamento delle repliche, utilizza un caricatore di dati elastico come grain.experimental.ElasticIterator. Gli iteratori di dati standard per host salvano lo stato per host e non possono eseguire il ripristino su un numero diverso di host, mentre ElasticIterator monitora una singola posizione del set di dati globale che può essere ripristinata su un numero qualsiasi di sezioni attive.

Loop di addestramento elastico

Un ciclo di addestramento elastico coordina il monitoraggio delle sezioni, il recupero degli errori e la creazione di snapshot in memoria, come mostrato nell'esempio di codice seguente:

  • La funzione train crea una singola istanza manager.Manager() e un Orbax Snapshotter all'esterno del wrapper di ripetizione, in modo che il monitoraggio degli slice e gli snapshot della memoria host vengano mantenuti tra i vari tentativi.
  • La decorazione della funzione interna run_training_loop con @elastic_manager.elastic_retry attende gli slice integri, imposta jax.default_device su uno slice attivo, mantiene attivo il thread del monitor in background tra i passaggi e riprova il loop quando si verifica un jax.errors.JaxRuntimeError.
  • Quando esegui il recupero da un errore di slice all'inizio di un tentativo, il codice deve ricostruire la mesh del dispositivo da elastic_manager.active_devices e chiamare snapshotter.load per ripristinare l'ultimo snapshot sugli sharding della mesh aggiornati, lasciando che i buffer del dispositivo del tentativo non riuscito escano dall'ambito.
  • Durante ogni passaggio del loop, il codice chiama periodicamente snapshotter.save per eseguire il backup dello stato di addestramento e controlla elastic_manager.available_inactive_slices per rilevare slice di sostituzione integre. Quando sono disponibili sezioni sostitutive, il codice le aggiunge a elastic_manager.active_slice_indices, ricrea la mesh, reinizializza il loop di addestramento e posiziona lo stato sulla mesh espansa.

L'esempio seguente mostra come implementare un loop di addestramento elastico combinando manager.Manager, elastic_retry e Snapshotter in memoria di Orbax:

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

Configura il gestore elastico e il comportamento di ripetizione

Crea un'istanza di manager.Manager(devices=None) una volta per processo. Se ometti devices, il gestore utilizza jax.devices() e raggruppa i dispositivi in base al loro attributo slice_index.

L'istanza Manager espone le seguenti proprietà per ispezionare la topologia del cluster durante l'addestramento:

  • slice_to_devices: una mappatura di ogni slice_index alla sequenza di oggetti jax.Device appartenenti a quella sezione.
  • all_slice_indices: un frozenset di tutti gli indici delle sezioni nell'allocazione.
  • active_slice_indices: un frozenset di indici di sezioni che partecipano alla mesh di calcolo attiva.
  • active_devices: un elenco di oggetti jax.Device appartenenti alle sezioni attive, ordinati per slice_index.
  • inactive_slice_indices: un frozenset di indici di slice nell'allocazione che non partecipano alla mesh di calcolo.
  • available_inactive_slices: un frozenset di indici di sezioni inattive che il thread del monitor in background ha rilevato come integri e pronti per essere reintegrati.
  • total_slice_count e active_slice_count: il numero totale di sezioni allocate e il numero di sezioni attive.
  • default_device: dispositivo 0 di uno slice attivo, che elastic_retry imposta automaticamente come dispositivo JAX predefinito per ogni tentativo.

Puoi personalizzare il comportamento di ripetizione e scalabilità passando i seguenti argomenti a @elastic_manager.elastic_retry:

  • minimum_slice_count: il numero minimo di sezioni integre richieste prima di avviare l'esecuzione iniziale o un tentativo di nuovo tentativo. Se impostato su None (il valore predefinito), elastic_retry attende tutte le sezioni nell'allocazione (total_slice_count), operando in modalità di pausa e ripresa. Se impostato su un numero intero più piccolo (ad esempio 1), elastic_retry opera in modalità di ridimensionamento delle repliche.
  • poll_interval: l'intervallo in secondi tra i controlli di integrità delle sezioni eseguiti dall'elastic training manager (Manager) durante l'addestramento. Il valore predefinito è 10.
  • timeout: il numero massimo di secondi di attesa prima che gli slice minimum_slice_count diventino attivi prima di generare un TimeoutError. Il valore predefinito è None (nessun timeout).
  • pre_callback: un callback facoltativo che viene eseguito all'inizio di ogni tentativo dopo l'aggiornamento di active_slice_indices e prima dell'esecuzione della funzione del loop di addestramento con wrapping.
  • on_elastic_event_callback: un callback facoltativo che viene eseguito dopo che viene rilevato un errore di slice-down jax.errors.JaxRuntimeError, prima dell'inizio del tentativo di nuovo tentativo successivo. Puoi utilizzare questo callback per ripulire i checkpoint incompleti o reimpostare gli iteratori di dati.
  • retry_policy: un (attempt: int, error: Exception) -> bool richiamabile che determina se continuare i tentativi dopo un evento elastico. Utilizza manager.ElasticRetryLimit(max_attempts) per limitare il numero totale di tentativi. Se omesso, elastic_retry riprova a tempo indeterminato.

In un evento elastico, JAX libera i buffer del dispositivo dal tentativo non riuscito una volta che lo stack frame si srotola e nessuna variabile fa riferimento a questi array. Evita di conservare riferimenti ad array di dispositivi di un tentativo non riuscito all'interno di variabili globali o callback di lunga durata, in modo che la memoria venga rilasciata prima dell'inizio del tentativo successivo.

Snapshot in memoria con Orbax

La classe ocp.v1.training.pathways.Snapshotter (disponibile in orbax-checkpoint versione 0.12.6 e successive) mantiene un backup in memoria dello stato di addestramento nella memoria pinned_host tra le sezioni:

  • Snapshotter(replica_axis_index=0): inizializza lo snapshotter e specifica quale asse della mesh rappresenta le repliche di slice parallele ai dati. Tutti gli array JAX nello stato salvato devono essere replicati lungo questo asse della mesh in modo che gli slice sopravvissuti conservino una copia completa del modello e dello stato dell'ottimizzatore quando uno slice non funziona. Se lo stato contiene array di singoli dispositivi, replicali nella mesh prima di chiamare save.
  • save(step, state): trasferisce in modo asincrono state alla memoria pinned_host e lo suddivide in alberi per replica lungo replica_axis_index. Se un trasferimento di snapshot precedente è ancora in corso, save salta la nuova richiesta per evitare di bloccare il loop di addestramento.
  • load(abstract_state, reset_snapshot_state=True): esegue il probe delle repliche host salvate per identificare le sezioni sopravvissute, concatena le repliche integre e ripartisce lo stato negli sharding di destinazione specificati da abstract_state.
  • join(): blocca fino al completamento di qualsiasi operazione di snapshot in coda.
  • latest: restituisce i metadati (incluso step) per lo snapshot completato più di recente o None se non è stato completato alcuno snapshot.

Riduci lo sharding in caso di errore di una sezione

Quando una sezione TPU non va a buon fine o viene prerilasciata durante l'esecuzione, Pathways genera un'eccezione jax.errors.JaxRuntimeError (ad esempio DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND o INTERNAL). Il decoratore elastic_retry rileva questa eccezione e controlla elastic.is_error_due_to_slice_down(error):

  • Se l'errore non è causato da un evento di riduzione, elastic_retry genera nuovamente l'eccezione immediatamente.
  • Se l'errore è causato da un evento di riduzione, elastic_retry interrompe qualsiasi traccia del profiler JAX attiva, esegue on_elastic_event_callback e valuta retry_policy(attempt, error).
  • Prima di iniziare il tentativo successivo, elastic_retry chiama elastic.wait_for_slices una volta, che esegue il polling dello stato degli slice finché non sono attivi almeno minimum_slice_count slice, aggiorna elastic_manager.active_slice_indices ed esegue la funzione del loop di addestramento sugli slice sopravvissuti.

Aumenta lo sharding quando le sezioni si uniscono di nuovo

Quando il workload viene eseguito su un numero ridotto di sezioni in modalità di ridimensionamento delle repliche, il thread del monitor in background esegue periodicamente controlli di integrità su elastic_manager.inactive_slice_indices. Quando le sezioni riparate o ripianificate superano il controllo di integrità, il thread di monitoraggio aggiunge i relativi indici di sezione a elastic_manager.available_inactive_slices.

Quando elastic_manager.available_inactive_slices non è vuoto alla fine di un passaggio di addestramento, lo scale up del workload viene eseguito in loco senza perdere i progressi dell'addestramento:

  • L'aggiunta degli indici delle sezioni in elastic_manager.available_inactive_slices a elastic_manager.active_slice_indices incorpora le sezioni riunite in elastic_manager.active_devices e le cancella da available_inactive_slices.
  • La ricostruzione della mesh del dispositivo da elastic_manager.active_devices e la reinizializzazione del passaggio di addestramento e dell'iteratore dei dati adattano il calcolo alla topologia espansa.
  • La chiamata a jax.device_put posiziona lo stato di addestramento corrente sui nuovi shard della mesh senza eseguire il ripristino da uno snapshot.

Hot-swap

Hot-swap si riferisce a una funzionalità dell'API JobSet di GKE in cui un job con priorità più alta può assumere rapidamente il controllo delle risorse di un job con priorità più bassa, riducendo al minimo i tempi di inattività e garantendo un recupero più rapido.

Quando viene creato un JobSet, GKE pianifica il workload su più sezioni, come specificato nella configurazione del JobSet. Se si verifica un errore hardware su una o più sezioni, i pod interessati vengono contrassegnati come non riusciti. Se il tuo cluster GKE ha una sezione di riserva che esegue un job con priorità inferiore, il sistema JobSet rimappa il workload della sezione con errore dal job con priorità più alta a quella sezione di riserva. In genere, questa nuova mappatura richiede meno di un minuto.

Quando un JobSet viene riavviato, l'hot swap può verificarsi nelle seguenti situazioni:

  • Modalità predefinita: se all'interno dello stesso cluster sono disponibili sezioni TPU di riserva inattive, lo scheduler Kubernetes dà la priorità alla pianificazione dei job riavviati su queste sezioni, anziché attendere la riparazione delle sezioni con errori. Questo comportamento di pianificazione consente un recupero più rapido.
  • Workload eterogenei: nei cluster che eseguono più workload con una risorsa Kubernetes PriorityClass configurata, un JobSet riavviato può attivare un hot swap. Se l'affinità del job riavviato corrisponde alle risorse di un job con priorità inferiore, Kubernetes prerilascia il job con priorità inferiore in modo che il job con priorità superiore possa avviarsi senza attendere la riparazione dello slice. Ad esempio, puoi configurare i pod worker di Pathways con priorità diverse utilizzando una risorsa PriorityClass.

Per utilizzare le priorità nel cluster, definisci un manifest PriorityClass:

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

Applica questo YAML al tuo cluster GKE:

kubectl apply -f high-prior-job.yaml

Successivamente, associa la nuova classe di priorità al job worker di Pathways aggiungendo il seguente testo alla specifica del pod pathways-worker.

priorityClassName: high-prior-job

Passaggi successivi