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 maggiori 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 le prestazioni del modello e come puoi risolverli.
Il modello è vincolato dall'input
Le TPU eseguono i calcoli molto rapidamente. Per assicurarti che la TPU non sia inattiva, è importante verificare che venga caricato un flusso costante di dati sulla TPU. Il modo in cui avviene questa operazione dipende dal modo in cui carichi e pre-elabora il tuo 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, usa la dimensione del batch più grande compatibile con TPU la memoria. Ogni core TPU utilizza registri vettoriali bidimensionali 8 X 128 per elaborare le moltiplicazioni matriciali. In generale, la dimensione del batch deve essere uniformemente 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 di un calcolo matriciale non riempiono un intero Con una dimensione di 128 x 8, il compilatore XLA padroneggia i tensori. Esistono due svantaggi del padding:
- I tensori con spaziatura interna sottoutilizzano il core TPU.
- Il padding aumenta la quantità di memoria on-chip richiesta per un tensore e può causare un errore di memoria insufficiente.
Mentre la spaziatura interna viene eseguita automaticamente dal compilatore XLA quando necessario, può determinare la quantità di spaziatura interna eseguita utilizzando lo strumento Visualizzatore memoria. Puoi evita il padding scegliendo dimensioni del tensore adatte alle TPU.
Dimensioni del tensore
Il compilatore XLA arrotonda le dimensioni dei tensori archiviati nella memoria HBM della TPU per eseguire 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 caratteristiche 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
Fusione è 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 round trip di memoria è ridotto e XLA non devi allocare spazio per "tmp".
Fusion è un'ottimizzazione fondamentale e offre vantaggi per Cloud TPU in diversi modi:
- Riduce i trasferimenti di memoria eliminando la necessità di archiviare 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 a
la forma della matrice. Il risultato dell'operazione ha la stessa forma della
matrice. Per ulteriori dettagli, consulta la guida
matrici di broadcasting.
Anche se spesso le trasmissioni possono essere integrate con i consumatori, potrebbe causare prestazioni scadenti e un aumento dell'utilizzo della memoria.
Nell'esempio seguente, la trasmissione implicita nell'aggiunta di un vettore e la matrice non può essere unita con l'argmax generando una trasmissione materializzata:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`