Addestramento su pod di TPU

Panoramica

Le TPU sono state progettate per essere sottoposte a scale out verso un pod di TPU. Un pod di TPU è una raccolta di dispositivi TPU connessi tramite interfacce di rete dedicate ad alta velocità. Un pod di TPU può avere fino a 2048 core TPU, che ti consentono di distribuire il carico di elaborazione tra più TPU. Ogni scheda TPU è collegata a una macchina host basata su CPU ad alte prestazioni, ad esempio per il caricamento e il pre-elaborazione dei dati. Per sfruttare appieno un numero maggiore di TPU, devi regolare diversi parametri dell'attività di addestramento. Questo documento spiega alcuni problemi comuni, le modifiche da apportare ai modelli e le best practice per ridurre o evitare errori dei pod.

Ridimensionare le dimensioni del batch e i passaggi dell'addestramento

Per ottenere una scalabilità lineare per i tipi di TPU più grandi, mantieni le stesse dimensioni in batch per core.

Ad esempio, se utilizzi una dimensione batch di 1024 su una versione v2-8, utilizza una dimensione batch di 4096 (4 * 1024) su una versione v2-32. Tutto questo sfrutta completamente l'hardware TPU. Puoi utilizzare batch di dimensioni più piccole, ma in questo caso l'addestramento non verrà scalato in modo lineare.

Molti modelli includono un flag train_steps in cui un passaggio corrisponde all'elaborazione di un singolo batch di dati. Quando aumenti le dimensioni del batch, fai lo scale down del numero di passaggi di addestramento in modo che il numero totale di esempi di addestramento rimanga invariato.

Ad esempio, se la dimensione del batch è 1000 per 100 passaggi, durante l'addestramento vengono elaborati 100.000 esempi. Se ora hai 4 worker e una dimensione del batch effettiva pari a 4000, devi regolare il numero di passaggi su 25 per elaborare gli stessi 100.000 esempi. Se il modello utilizza un flag epochs, non è necessario scalare il numero di passaggi.

I batch di dimensioni più grandi possono cambiare il comportamento di convergenza del modello, quindi puoi anche regolare alcuni iperparametri, ad esempio la frequenza di apprendimento.

Utilizzo dei bucket Google Cloud Storage a livello di area geografica nella stessa area geografica del pod TPU

In generale, la best practice per l'addestramento di TPU consiste nell'utilizzare sempre le risorse nella stessa area geografica. L'area geografica delle risorse è particolarmente importante quando si utilizzano i pod TPU, perché la velocità di trasferimento dei dati è più elevata quando il bucket Google Cloud Storage e la TPU sono nella stessa area geografica. Assicurati di utilizzare un bucket Google Cloud Storage a livello di area geografica nella stessa area geografica della TPU per l'addestramento di set di dati e punti di controllo.

Utilizzo di una NFS per l'archiviazione dei dati

Quando crei un pod TPU, viene creata una VM separata per ogni nodo TPU. Per impostazione predefinita, a ogni VM TPU viene assegnato un ID utente diverso (UID). Questo causa problemi quando provi ad accedere alla stessa directory NFS da più nodi: la stessa directory avrebbe proprietari diversi su nodi diversi e le autorizzazioni Linux standard non si applicano a tutti i nodi. Ad esempio, un processo in un nodo non può scrivere nella directory del log creata da un altro nodo,

Per aggirare il problema, puoi attivare OS Login. OS Login può essere configurato per un'istanza VM specifica o per un progetto. Per ulteriori informazioni, consulta Configurare OS Login.

Best practice per il flusso di lavoro per lo sviluppo su pod di TPU

Durante lo sviluppo di un nuovo carico di lavoro TPU, spesso è preferibile iniziare a sviluppare sulle TPU più piccole e passare progressivamente alle dimensioni TPU più grandi. Inizia utilizzando la TPU più piccola (ad esempio, v2-8 o v3-8).

  • Testa la funzionalità del tuo carico di lavoro
  • Testa e convalida il rendimento utilizzando gli strumenti per il rendimento

Quando il tuo carico di lavoro è funzionale e raggiunge i tuoi obiettivi di rendimento, fai lo scale up fino a v2-32 o v3-32. Aumenta in modo graduale e graduale le dimensioni della TPU convalidando la scalabilità (funzionalità e prestazioni) fino a raggiungere la dimensione della TPU desiderata.