Addestramento su pod di TPU

Panoramica

Le TPU sono state progettate per essere scalate a un pod TPU. Un pod di TPU è una raccolta di dispositivi TPU connessi 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 la pre-elaborazione dei dati. Per sfruttare al meglio un numero maggiore di TPU, devi ottimizzare diversi parametri dell'attività di addestramento.

La configurazione per l'addestramento con pod TPU è diversa a seconda del framework. Utilizza i seguenti link per visualizzare informazioni dettagliate sull'addestramento dei pod con ogni framework:

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

Scalabilità delle dimensioni del batch e dell'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 del batch di 1024 su una v2-8, utilizza una dimensione del batch di 4096 (4 * 1024) su una v2-32. Viene utilizzato completamente l'hardware TPU. È possibile utilizzare batch di dimensioni 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 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 effettiva del batch di 4000, devi modificare il numero di passaggi in 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 maggiori possono modificare il comportamento di convergenza del modello, quindi potresti regolare alcuni iperparametri, come la percentuale di apprendimento.

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

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

Assicurati di utilizzare un bucket Google Cloud Storage a livello di regione nella stessa area geografica della TPU per i set di dati e i punti di controllo per l'addestramento.

Best practice per il flusso di lavoro dello sviluppo sui pod TPU

Quando si sviluppa un nuovo carico di lavoro TPU, è spesso ottimale iniziare a sviluppare le TPU più piccole ed eseguire progressivamente l'iterazione in dimensioni più grandi. Inizia utilizzando una piccola versione di TPU (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 carico di lavoro è funzionale e raggiunge gli obiettivi di prestazioni, fai lo scale up fino 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 la dimensione TPU desiderata.