Addestramento su sezioni TPU

Le TPU sono progettate per essere scalate in un pod TPU. Un pod TPU è un insieme di dispositivi TPU connessi da interfacce di rete ad alta velocità dedicate. Un pod 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 sulla formazione sui 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 una dimensione del batch pari a 1024 su una versione v6e-8, utilizza una dimensione del batch pari a 4096 (4 * 1024) su una versione v6e-32. In questo modo viene utilizzato completamente l'hardware TPU. Puoi utilizzare lotti di dimensioni inferiori, 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, riduci 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 Cloud Storage regionali nella stessa regione del pod 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 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 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 TPU di piccole dimensioni (ad esempio v6e-8).

  • Verificare il comportamento previsto del carico di lavoro
  • Testare e convalidare il rendimento utilizzando gli strumenti di rendimento

Una volta che il carico di lavoro è funzionale e raggiunge i tuoi target di prestazioni, esegui l'upgrade a un tipo di TPU più grande, ad esempio una 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 che preferisci.