Formazione resiliente con Pathways

Questa pagina spiega come configurare le funzionalità di resilienza per i workload Pathways su Google Kubernetes Engine (GKE) in modo che i job di addestramento su larga scala possano riprendersi da interruzioni pianificate e guasti hardware non pianificati.

Pathways offre vantaggi in termini di resilienza nei seguenti modi:

  • Sospensione-ripristino: 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 arrestare in modo anomalo il client utilizzando la logica di recupero nel codice di addestramento.
  • Checkpointing multilivello: consente al workload di salvare checkpoint locali e replicati tra peer ad alta frequenza tra i worker TPU utilizzando Python in co-locazione, con backup asincroni di Cloud Storage per ridurre i tempi di ripristino.
  • Hot-swap: consente a un workload JobSet GKE con priorità più alta di assumere rapidamente il controllo delle sezioni TPU di riserva o con priorità più bassa quando si verifica un guasto hardware.

Prima di iniziare

Prima di iniziare, completa le seguenti attività:

Se utilizzi l'addestramento elastico, installa i seguenti pacchetti nell'immagine container del workload:

  • pathwaysutils
  • Orbax (orbax-checkpoint): >= 0.12.6 (obbligatorio per gli snapshot in memoria con ocp.v1.training.pathways.Snapshotter)
  • Granularità (grain): obbligatoria se utilizzi grain.experimental.ElasticIterator

Se utilizzi il checkpoint multilivello, completa anche i seguenti prerequisiti:

  • Configura il cluster GKE per il checkpointing multilivello seguendo le istruzioni riportate in Addestra modelli di machine learning su larga scala su GKE con il checkpointing multilivello:
    • Abilita i componenti aggiuntivi Multi-Tier Checkpointing e Cloud Storage FUSE (GcsFuseCsiDriver), insieme a Workload Identity Federation for GKE, sul tuo cluster.
    • Crea un bucket Cloud Storage con lo spazio dei nomi gerarchico abilitato.
    • Concedi i ruoli Identity and Access Management (IAM) roles/storage.objectUser e roles/storage.insightsCollectorService al account di servizio Kubernetes di checkpoint multi-livello.
    • Esegui il deployment di una risorsa personalizzata CheckpointConfiguration specificando il tuo bucket Cloud Storage, il selettore di nodi TPU e le dimensioni del volume in memoria (inMemoryVolumeSize).
  • Crea un'immagine container sidecar Python in co-localizzazione ed eseguine il push, ad esempio utilizzando il Dockerfile sidecar di esempio. Nelle immagini del workload e del sidecar, utilizza versioni corrispondenti di jax e jaxlib che siano in linea con la tua release di Pathways e soddisfino i seguenti requisiti minimi di versione:
    • Versione Orbax: >= 0.12.4
    • Versione JAX (incluso Python in co-locazione): >= 0.10.1

Sospensione-ripristino

In genere, GKE invia una notifica di prerilascio a un pod acceleratore prima di prerilasciarlo. Pathways attiva la tolleranza di prelazione per impostazione predefinita su tutti i deployment cloud 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. Se possibile, Pathways 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 workload 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 tipici workload ML 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 JAX collocata, si basano su componenti Pathways stateful e 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: se imposti minimum_slice_count su None o sul numero totale di sezioni, l'addestramento viene messo in pausa quando una sezione non va a buon fine e riprende quando tutte le sezioni sono di nuovo integre. 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: se imposti minimum_slice_count su un valore inferiore al numero totale di slice, l'addestramento continua sugli slice attivi rimanenti quando uno slice non va a buon fine, ricostruendo la mesh del dispositivo e ricompilando per la topologia più piccola. Un thread di monitoraggio in background controlla periodicamente gli slice inattivi e registra gli slice appena tornati integri in elastic_manager.available_inactive_slices. Il loop di addestramento può quindi ricostruire la mesh dei dispositivi con gli slice sostitutivi, reinizializzare per la topologia espansa e posizionare lo stato attuale sulla nuova 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/cloud-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 essere ripristinati su un numero diverso di host. Al contrario, ElasticIterator monitora una singola posizione del set di dati globale che puoi ripristinare in 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.
  • L'applicazione di @elastic_manager.elastic_retry a run_training_loop configura la funzione in modo che attenda gli slice integri, imposti jax.default_device su uno slice attivo, mantenga attivo il thread del monitor in background e riprovi su un jax.errors.JaxRuntimeError.
  • All'inizio di un tentativo di recupero, 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. Lascia che i buffer del dispositivo del tentativo non riuscito escano dall'ambito in modo che JAX possa liberare la memoria.
  • 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 che il gestore aggiorna active_slice_indices e prima dell'esecuzione della funzione del ciclo di addestramento con wrapping.
  • on_elastic_event_callback: un callback facoltativo che viene eseguito dopo che elastic_retry rileva un errore di riduzione dello slice 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 dei 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 riesce o viene prerilasciata durante l'esecuzione, Pathways genera un'eccezione jax.errors.JaxRuntimeError (ad esempio DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND o INTERNAL). In una funzione decorata con @elastic_manager.elastic_retry, 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 per eseguire il polling dello stato dello slice finché non sono attivi almeno minimum_slice_count slice. Quindi aggiorna elastic_manager.active_slice_indices ed esegue la funzione del ciclo 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.

Checkpointing multilivello

I job di addestramento su larga scala subiscono frequenti interruzioni e la scrittura o il ripristino dei checkpoint direttamente da e verso Cloud Storage può creare colli di bottiglia nell'avanzamento dell'addestramento. Il checkpointing multilivello riduce l'overhead del checkpointing e il tempo di recupero coordinando tre livelli di archiviazione su GKE:

  • Disco RAM locale: i worker TPU scrivono checkpoint ad alta frequenza in un volume CSI in memoria (multitier-checkpoint.csi.storage.gke.io) montato su ogni nodo worker. Questo volume persiste tra i riavvii dei pod sullo stesso nodo.
  • Replica peer in-cluster: il servizio di replica GKE replica gli shard dei checkpoint locali tra i worker TPU peer nel cluster.
  • Backup di Cloud Storage: il servizio di replica esegue il backup asincrono dei checkpoint locali in un bucket Cloud Storage con spazio dei nomi gerarchico abilitato, senza bloccare i passaggi di addestramento della TPU.

A differenza di JAX multi-controller, in cui ogni host TPU esegue il proprio script di addestramento, Pathways esegue un singolo controller sul pod pathways-head e invia il lavoro ai pod pathways-worker. Per leggere e scrivere shard di checkpoint direttamente sul disco RAM locale di ogni host TPU senza instradare i dati dei checkpoint tramite il singolo controller, Pathways utilizza Python collocato. Ogni pod worker esegue un container colocated-python-sidecar insieme al container pathways-worker ed entrambi i container montano il volume CSI di checkpointing multilivello e una directory di memoria condivisa (/tmp/sidecar). Durante l'inizializzazione e l'addestramento, Orbax invia la configurazione, i salvataggi e i ripristini dei checkpoint dal controller Pathways ai sidecar worker su ogni host TPU. Il checkpointing multilivello su Pathways supporta anche l'addestramento elastico quando si fa lo scale down o si fa lo scale up del numero di slice TPU attive.

Pathways supporta il checkpointing multilivello con MaxText. Per le istruzioni di configurazione, consulta la guida al checkpointing multilivello di MaxText.

Configura il workload Pathways

Puoi configurare un workload Pathways per il checkpoint multilivello utilizzando Cluster Toolkit (gcluster job submit) o configurando un manifest JobSet Pathways personalizzato.

Utilizza Cluster Toolkit

Quando invii un job con gcluster job submit --pathways, passa i flag --gke-mtc-enabled, --gke-mtc-ramdisk-dir e --pathways-colocated-python-sidecar-image:

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

Sostituisci SIDECAR_IMAGE con l'URI dell'immagine container sidecar Python collocato.

Quando imposti questi flag, gcluster configura automaticamente quanto segue:

  • Passa --sidecar_name=external al container pathways-proxy.
  • Esegue il container colocated-python-sidecar in ogni pod worker e monta una directory /tmp/sidecar condivisa (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) sia su pathways-worker che su colocated-python-sidecar.
  • Monta il volume CSI multitier-checkpoint.csi.storage.gke.io nel percorso specificato da --gke-mtc-ramdisk-dir sia in pathways-worker sia in colocated-python-sidecar sui pod worker, ma non sul pod pathways-head.

Utilizza un manifest JobSet di Pathways personalizzato

Se gestisci il tuo manifest JobSet di Pathways anziché utilizzare gcluster, passa --sidecar_name=external a pathways-proxy. Nella specifica del pod worker, monta sia il volume di memoria sidecar condiviso sia il volume CSI multitier-checkpoint.csi.storage.gke.io sui container pathways-worker e colocated-python-sidecar:

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Utilizza il checkpointing multilivello Orbax in un workload JAX personalizzato

In uno script di addestramento JAX personalizzato in esecuzione su Pathways, utilizza il pacchetto orbax.checkpoint.experimental.emergency.multi_tier_checkpointing e imposta use_colocated_python=True sia in initialize_multi_tier_checkpointing sia in ReplicatorCheckpointManagerOptions:

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

Hot-swap

La sostituzione a caldo è una funzionalità dell'API JobSet di GKE che consente a un job con priorità più alta di assumere rapidamente il controllo delle risorse di un job con priorità più bassa, riducendo al minimo i tempi di inattività e accelerando il recupero.

Quando crei 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, GKE contrassegna i pod interessati come non riusciti. Se il cluster GKE ha una sezione di riserva che esegue un job con priorità inferiore, il controller JobSet rimappa il workload della sezione con errore dal job con priorità più alta alla sezione di riserva. In genere, questa rimappatura 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 e salvalo come high-prior-job.yaml:

apiVersion: scheduling.k8s.io/v1
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, collega la nuova classe di priorità al workload Pathways. Se invii il tuo workload utilizzando Cluster Toolkit, passa --priority=high-prior-job al comando gcluster job submit, che imposta priorityClassName sia sui pod pathways-head sia sui pod worker. Se scrivi il tuo manifest JobSet, imposta priorityClassName nel template del pod del job replicato worker, che esegue il container pathways-worker:

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

Se hai creato il cluster utilizzando un progetto TPU di Cluster Toolkit con Pathways abilitato, il cluster include già le classi di priorità very-low, low, medium, high e very-high, quindi puoi passarne una a --priority anziché crearne una tua.

Passaggi successivi