Addestramento su pod di TPU

Panoramica

Le TPU sono state progettate per essere scalate in un pod di TPU. Un pod di TPU è un insieme di dispositivi TPU collegati da interfacce di rete ad alta velocità dedicate. Un pod di TPU consente di distribuire il carico di elaborazione su più TPU. Ogni scheda TPU è collegata a una macchina host basata su CPU ad alte prestazioni per operazioni come il caricamento e l'elaborazione preliminare dei dati. Per sfruttare al meglio un numero maggiore di TPU, devi ottimizzare diversi parametri delle attività di addestramento.

La configurazione per l'addestramento con i pod di TPU è diversa per ogni framework. Utilizza i seguenti link per visualizzare informazioni dettagliate sull'addestramento su Pod con ciascun framework:

Le sezioni seguenti spiegano alcuni problemi comuni, le modifiche da apportare ai modelli e le best practice per ridurre o evitare gli errori dei pod.

Scalare la dimensione del batch e i passaggi di addestramento

Per ottenere una scalabilità lineare su tipi di TPU più grandi, mantieni invariata la dimensione del batch per core.

Ad esempio, se utilizzi un batch di 1024 su una versione 2-8, utilizza un batch di 4096 (4 * 1024) su una versione 2-32. In questo modo viene utilizzato completamente l'hardware TPU. Puoi utilizzare lotti di dimensioni inferiori, ma l'addestramento non verrà scalato in modo lineare.

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

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

Batch di dimensioni maggiori possono modificare il comportamento di convergenza del modello, quindi potresti anche ottimizzare alcuni iperparametri, come il tasso di apprendimento.

Utilizzo di bucket Google Cloud Storage regionali nella stessa regione del pod di TPU

In generale, la best practice per l'addestramento TPU è utilizzare sempre le risorse nella stessa regione. La regione delle risorse è particolarmente importante quando utilizzi i pod TPU perché le velocità di trasferimento dei dati sono più elevate quando il bucket Google Cloud Storage e la TPU si trovano nella stessa regione.

Assicurati di utilizzare un bucket Google Cloud Storage regionale nella stessa regione della TPU per i set di dati di addestramento e i checkpoint.

Best practice per i flussi di lavoro per lo sviluppo su pod TPU

Quando sviluppi un nuovo carico di lavoro TPU, spesso è ottimale iniziare lo sviluppo sulle TPU più piccole e procedere gradualmente a dimensioni TPU più grandi. Inizia utilizzando una versione di TPU piccola (ad esempio v2-8 o v3-8).

  • Verifica la funzionalità del carico di lavoro
  • Testare e convalidare le prestazioni utilizzando gli strumenti di rendimento

Una volta che il carico di lavoro è funzionale e raggiunge i tuoi target di rendimento, esegui l'upgrade a un tipo di TPU più grande, ad esempio v2-32 o v3-32. Aumenta gradualmente e in modo iterativo le dimensioni della TPU, convalidando la scalabilità (funzionalità e prestazioni) fino a raggiungere le dimensioni della TPU desiderate.