Addestramento su pod di TPU

Panoramica

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

La configurazione per l'addestramento con i pod TPU è diversa per ogni framework. Usa i seguenti link per visualizzare informazioni dettagliate sull'addestramento sui pod con ciascun framework:

• TensorFlow

• PyTorch

• JAX

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

Scalabilità delle dimensioni del batch e passaggi di addestramento

Per ottenere la scalabilità lineare su tipi di TPU più grandi, mantieni le stesse dimensioni del 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) in una v2-32. Questa funzionalità sfrutta completamente l'hardware TPU. Puoi utilizzare batch di dimensioni ridotte, ma in questo caso 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 lo stesso.

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 batch effettiva di 4000, dovresti modificare 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.

Dimensioni batch più grandi possono modificare il comportamento di convergenza del modello, quindi potresti ottimizzare alcuni iperparametri, come il tasso di apprendimento.

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

In generale, la best practice per l'addestramento di TPU è utilizzare sempre risorse nella stessa regione. La regione delle risorse è particolarmente importante quando si utilizzano 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 l'addestramento di set di dati e checkpoint.

Best practice del flusso di lavoro per lo sviluppo di pod TPU

Durante lo sviluppo di un nuovo carico di lavoro TPU, spesso è ottimale iniziare lo sviluppo sulle TPU più piccole e ripetere l'iterazione progressivamente fino a dimensioni TPU più grandi. Inizia utilizzando una versione di TPU piccola (ad esempio v2-8 o v3-8).

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

Quando il carico di lavoro è funzionale e raggiunge i target di prestazioni, fai lo scale up a un tipo di TPU più grande, come v2-32 o v3-32. Aumenta gradualmente e in modo graduale le dimensioni della TPU convalidando al contempo la scalabilità (funzionalità e prestazioni) fino a raggiungere la dimensione TPU desiderata.