Gestire gli eventi di manutenzione con i pod di Cloud TPU

Panoramica

I nodi TPU e le VM TPU sono istanze di VM di Compute Engine con hardware TPU collegato. Le VM di Compute Engine sono soggette a eventi di manutenzione delle VM di Compute Engine. Ogni TPU è connessa a una VM di Compute Engine, utilizzando quindi più TPU (ad esempio in un pod TPU) per aumentare la probabilità che una delle tue VM riscontri un evento di manutenzione.

In questo documento vengono illustrati i vari approcci per gestire gli eventi di manutenzione per i job di addestramento a lunga esecuzione sulle Cloud TPU.

Utilizzo di punti di controllo per il recupero rapido in seguito agli eventi di manutenzione

I punti di accesso sono fondamentali per i rapidi ripristini degli eventi di manutenzione e devono essere salvati di frequente: una buona regola generale è salvare i punti di controllo circa ogni ora. Mancare il controllo di un numero sufficiente di rischi rischia di perdere molti progressi di addestramento a causa di eventi di manutenzione o altre interruzioni della formazione.

In genere, i punti di controllo si riferiscono a tutti i parametri salvati utilizzati nell'addestramento, come le ponderazioni del modello. Il tempo necessario per salvare un posto di blocco può variare dall'ordine dei secondi all'ordine dei minuti.

Anche se la maggior parte degli eventi di manutenzione viene recuperata automaticamente e i job di addestramento continuano senza interventi manuali, potrebbero verificarsi casi limite in cui il job non viene riavviato e continua automaticamente. In questo caso, devi eliminare e ricreare le risorse TPU, quindi riavviare il job di addestramento da un checkpoint salvato. Per informazioni su come rilevare e ripristinare da errori di ripristino automatico, consulta la sezione Rilevare e recuperare da errori di TPU di seguito.

I meccanismi utilizzati per salvare e caricare i punti di controllo sono diversi per ogni framework ML. In genere, i modelli di Cloud TPU supportati hanno il checkpoint integrato. Per ulteriori informazioni sul checkpoint, consulta TensorFlow 2.x, PyTorch o JAX/flax.

Rilevamento degli eventi di manutenzione

Puoi rilevare se e quando si verifica un evento di manutenzione sulla TPU utilizzando il seguente comando gcloud describe:

VM TPU

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

Nodi TPU

$ gcloud compute tpus describe tpu-name --zone=zone | grep 'health'

L'output di questo comando visualizza lo stato attuale della TPU e una descrizione dell'evento di manutenzione più recente. L'output dovrebbe avere il seguente aspetto:

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

Gestire gli eventi di manutenzione

Esistono diversi modi per ridurre le interruzioni degli eventi di manutenzione.

  1. Salvare periodicamente i punti di controllo

    Nello scenario ideale, quando si verifica un "evento di interruzione", l'addestramento riprende semplicemente dal punto di controllo più recente.

  2. Tentativi script di addestramento

    Lo script di addestramento potrebbe interrompersi a causa di un "evento di interruzione". Puoi utilizzare uno script bash per riprovare continuamente lo script di addestramento fino al completamento dell'addestramento. Ogni nuovo tentativo dovrebbe continuare dal punto di controllo più recente, quindi gli script per i nuovi tentativi dovrebbero essere sempre utilizzati insieme ai punti di controllo.

    Le pipeline di addestramento pronto per la produzione devono utilizzare GKE o altri sistemi di gestione delle risorse. In caso contrario, puoi implementare uno script bash per riprovare continuamente lo script di addestramento fino al completamento. Ad esempio:

    Con nodo TPU:

    (Dalla tua VM) while ! python3 [training command]; do sleep 1; done

    Con VM TPU: while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done

    Tieni presente che devi eseguire il comando della VM TPU da GCP o da un terminale, non dalla VM TPU.

  3. Rileva e recupera da errori della TPU

    Quando una TPU non recupera da un evento di manutenzione, puoi utilizzare uno script di ripristino per rilevare lo stato della TPU ed eliminare e ricreare la TPU. Per un esempio di questo script, fai clic qui. Per informazioni dettagliate sull'eliminazione e sulla creazione manuale delle TPU, consulta la sezione Gestire le TPU.

    Quando crei o ricrei una VM TPU, puoi specificare uno script di avvio con il parametro --metadata startup-script. Uno script di avvio viene eseguito ogni volta che viene creata una VM TPU. Per ulteriori informazioni, consulta l'articolo Eseguire script di installazione standard.