Addestramento su slice di TPU
Le TPU sono progettate per essere scalate orizzontalmente a 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 su più TPU. Ogni scheda TPU è collegata a una macchina host basata su CPU ad alte prestazioni per attività come il caricamento e la preelaborazione dei dati. Per sfruttare appieno 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 sui pod con ogni 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 sui tipi di TPU più grandi, mantieni invariata la dimensione del batch per core.
Ad esempio, se utilizzi una dimensione del batch di 1024 su una v6e-8, utilizza una dimensione del batch di 4096 (4 * 1024) su una v6e-32. In questo modo, l'hardware TPU viene utilizzato completamente. Puoi utilizzare dimensioni del batch più piccole, 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 del numero di passaggi di addestramento in modo che il numero totale di esempi di addestramento rimanga lo stesso.
Ad esempio, se hai una dimensione del batch di 1000 per 100 passaggi, durante l'addestramento vengono elaborati 100.000 esempi. Se ora hai 4 worker e una dimensione del batch effettiva di 4000, dovrai modificare il numero di passaggi a 25 per elaborare gli stessi 100.000 esempi. Se il modello utilizza un flag epochs, non devi scalare il numero di passaggi.
Le dimensioni del batch più grandi possono modificare il comportamento di convergenza del modello, pertanto potresti anche ottimizzare alcuni iperparametri, come il tasso di apprendimento.
Utilizzare bucket Cloud Storage regionali nella stessa regione del pod di TPU
In generale, la best practice per l'addestramento delle TPU è utilizzare sempre le risorse nella stessa regione. La regione delle risorse è particolarmente importante quando si utilizzano i pod di TPU perché le velocità di trasferimento dei dati sono più elevate quando il bucket Cloud Storage e la TPU si trovano nella stessa regione.
Assicurati di utilizzare un bucket Cloud Storage regionale nella stessa regione della TPU per i set di dati di addestramento e i checkpoint.
Best practice per il flusso di lavoro per lo sviluppo sui pod di TPU
Quando sviluppi un nuovo carico di lavoro TPU, spesso è ottimale iniziare lo sviluppo sulle TPU più piccole e passare progressivamente a dimensioni di TPU più grandi. Inizia utilizzando una versione TPU di piccole dimensioni (ad esempio, v6e-8).
- Testa il carico di lavoro per verificare il comportamento previsto
- Testa e convalida le prestazioni utilizzando gli strumenti per le prestazioni
Una volta che il carico di lavoro è funzionale e raggiunge i target di prestazioni, fai lo scale up a un tipo di TPU più grande, ad esempio v6e-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.