Guida alle prestazioni di Cloud TPU
Il primo passaggio per risolvere i problemi di prestazioni della TPU consiste nel creare il profilo del modello. Per ulteriori informazioni sull'acquisizione di un profilo delle prestazioni, consulta Profilazione del modello su Cloud TPU.
Prestazioni del modello TPU
Questa sezione descrive i problemi generali che possono ridurre il rendimento del modello e come risolverli.
Il modello è vincolato dall'input
Le TPU eseguono calcoli molto rapidamente. Per assicurarti che la TPU non sia inattiva, è importante verificare che venga caricato un flusso costante di dati sulla TPU. La modalità di esecuzione dipende da come carichi e pre-elabori il set di dati. Ad esempio, puoi leggere i file di dati in parallelo utilizzando tf.data.TFRecordset() e il parametro
num_parallel_reads.La dimensione del batch è troppo piccola a causa dello sharding (suddivisione dei batch tra i core)
Il runtime TPU suddivide un batch su tutti gli 8 core di un dispositivo TPU (ad esempio v2-8 o v3-8). Se specifichi una dimensione del batch globale di 128, ogni core riceve una dimensione del batch di 16 (128 / 8).
Per un utilizzo ottimale della memoria, utilizza la dimensione del batch più grande che rientra nella memoria TPU. Ogni core TPU utilizza registri vettoriali 8 x 128 bidimensionali per elaborare le moltiplicazioni di matrici. In generale, la dimensione del batch deve essere perfettamente divisibile per 8 o 128.
Ottimizzazioni del compilatore XLA
XLA è un compilatore per il machine learning che può produrre binari per TPU, CPU, GPU e altre piattaforme. Sebbene XLA faccia parte della base di codice standard di TensorFlow, può essere utilizzato anche sui modelli PyTorch e JAX. I modelli per Cloud TPU vengono tradotti in un grafo XLA, che XLA compila in un file eseguibile TPU. Per saperne di più su XLA, consulta XLA: compilatore ottimizzato per il machine learning.
Spaziatura interna
Per utilizzare la memoria TPU in modo efficiente, struttura i dati in modo che possano essere suddivisi in schemi di 128 x 8. Quando i dati per un calcolo della matrice non riempiono un intero blocco 128 x 8, il compilatore XLA aggiunge spazi ai tensori. Il padding presenta due svantaggi:
- I tensori con spaziatura interna sottoutilizzano il core TPU.
- Il padding aumenta la quantità di memoria on-chip richiesta per un tensore e può portare a un errore di memoria insufficiente.
Sebbene il padding venga eseguito automaticamente dal compilatore XLA, se necessario, puoi determinare la quantità di padding eseguita utilizzando lo strumento di visualizzazione della memoria. Puoi evitare il padding scegliendo dimensioni dei tensori adatte per le TPU.
Dimensioni del tensore
Il compilatore XLA arrotonda per eccesso le dimensioni dei tensori memorizzati nella memoria HBM della TPU per eseguire i calcoli in modo più efficiente. Questo riempimento avviene in modo trasparente a livello di hardware e non influisce sui risultati. Tuttavia, in alcuni casi il padding può comportare un aumento significativo dell'utilizzo della memoria e del tempo di esecuzione.
Il runtime TPU dispone i tensori in memoria per massimizzare l'efficienza di calcolo e minimizzare il padding. Per ridurre al minimo il sovraccarico di memoria e massimizzare l'efficienza computazionale, deve essere vera una delle seguenti condizioni:
La dimensione totale del batch deve essere un multiplo di 64 (8 per core TPU) e le dimensioni delle dimensioni delle funzionalità devono essere un multiplo di 128.
La dimensione totale del batch deve essere un multiplo di 1024 (128 per core TPU) e le dimensioni delle dimensioni delle funzionalità devono essere un multiplo di 8.
L'utilizzo di un batch di 1024 e dimensioni delle funzionalità che sono un multiplo di 128 consente di ottenere l'efficienza migliore, anche se potrebbe non essere possibile per tutti i modelli.
Fusione
Fusion è una tecnica generale utilizzata dal compilatore XLA per ottimizzare i programmi. Un'operazione fusa è la combinazione di più operazioni costituenti che devono essere eseguite in combinazione.
Ad esempio, considera la seguente serie di operazioni:
tmp = tf.add(x, y)
result = tf.multiply(tmp, z)
Questo codice è approssimativamente equivalente al seguente pseudocodice:
for (i = 0; i < element_count; i++) {
tmp[i] = x[i] + y[i];
}
for (i = 0; i < element_count; i++) {
result = tmp[i] * z[i];
}
Con la fusione, gli accessi all'array avvengono contemporaneamente:
for (i = 0; i < element_count; i++) {
result = (x[i] + y[i]) * z[i];
}
In questo esempio, il numero di viaggi di andata e ritorno della memoria è ridotto e XLA non deve allocare spazio per "tmp".
La fusione è un'ottimizzazione fondamentale e offre vantaggi a Cloud TPU in diversi modi:
- Riduce i trasferimenti di memoria eliminando la necessità di memorizzare i risultati intermedi nella memoria principale, che è lenta.
- Consente un maggiore utilizzo delle unità hardware che altrimenti rimarrebbero inutilizzate.
- Può ridurre l'utilizzo della memoria di un modello in quanto è necessario attivare meno buffer contemporaneamente.
Trasmissione
La trasmissione avviene implicitamente quando vengono combinati due tensori con forme diverse, ma compatibili.
Ad esempio, tf.add(vector, matrix) richiede che il vettore venga trasmesso alla forma della matrice. Il risultato dell'operazione ha la stessa forma della
matrice. Per maggiori dettagli, consulta la guida agli array di trasmissione.
Sebbene le trasmissioni possano spesso essere fuse con i relativi consumatori, forzare una trasmissione potrebbe comportare prestazioni scadenti e un aumento dell'utilizzo della memoria.
Nell'esempio seguente, la distribuzione implicita nell'aggiunta di un vettore e di una matrice non può essere combinata con l'argmax, dando luogo a una distribuzione materializzata:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`