Guida alle prestazioni di Cloud TPU
Il primo passaggio per la risoluzione dei problemi relativi alle prestazioni della TPU è creare un profilo per il modello. Per ulteriori informazioni sull'acquisizione di un profilo delle prestazioni, vedi Profilazione del modello su Cloud TPU.
Prestazioni del modello TPU
In questa sezione vengono descritti i problemi generali che possono ridurre le prestazioni del modello e come risolverli.
Il modello è associato all'input
Le TPU possono eseguire calcoli molto rapidamente. Per assicurare che la TPU non sia inattiva, è importante assicurarsi che venga caricato un flusso stabile di dati sulla TPU. L'avanzamento dipende dal modo in cui carichi ed esegui il pre-elaborazione del set di dati. Ad esempio, puoi leggere i file di dati in parallelo utilizzando tf.data.TFRecordset() e il parametro
num_parallel_reads.Le dimensioni del batch sono troppo piccole a causa del partizionamento orizzontale (suddivisione dei batch tra i core)
Il runtime TPU suddivide un batch tra 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 le dimensioni batch più elevate possibili in base alla memoria TPU. Ogni core TPU utilizza registri registrati 8 x 128 bidimensionali per l'elaborazione di moltiplicazioni delle matrici. In generale, le dimensioni batch dovrebbero essere divisibili uniformemente di 8 o 128.
Ottimizzazioni del compilatore XLA
XLA è un compilatore per machine learning che può produrre programmi binari per TPU, CPU, GPU e altre piattaforme. XLA fa parte del codebase standard di TensorFlow, ma può essere utilizzato anche sui modelli PyTorch e JAX. I modelli di Cloud TPU vengono tradotti in un grafico XLA, che poi compila in base a un file eseguibile di TPU. Per ulteriori informazioni su XLA, consulta il documento XLA: Optimizing Compiler for Machine Learning.
Spaziatura interna
Per utilizzare la memoria in modo efficiente, struttura i dati in modo che possano essere inseriti in riquadri con blocchi di 128 x 8. Quando i dati per il calcolo di una matrice non riempiono un intero blocco di 128 x 8, il compilatore XLA registra i tensori. La spaziatura interna presenta due svantaggi:
- I tensori imbottiti sottoutilizzano il core TPU.
- La spaziatura interna aumenta la quantità di spazio di archiviazione su chip richiesta per un tensore e può causare un errore di memoria insufficiente.
La spaziatura interna viene eseguita automaticamente dal compilatore XLA, se necessario, ma puoi determinare la quantità di spaziatura interna eseguita utilizzando lo strumento di visualizzazione della memoria. Puoi evitare la spaziatura interna scegliendo le dimensioni del tensore adatte per la TPU.
Dimensioni Tensor
Il compilatore XLA arrotonda le dimensioni dei tensori archiviati nella memoria HBM di TPU per eseguire i calcoli in modo più efficiente. Questa spaziatura interna avviene in modo trasparente a livello di hardware e non influisce sui risultati. Tuttavia, in alcuni casi la spaziatura interna può comportare un aumento significativo dell'utilizzo della memoria e dei tempi di esecuzione.
Il runtime TPU definisce i tensori in memoria per massimizzare l'efficienza computazionale e ridurre al minimo la spaziatura interna. Per ridurre al minimo l'overhead di memoria e massimizzare l'efficienza di calcolo, uno dei seguenti elementi deve essere vero:
La dimensione totale del batch deve essere un multiplo di 64 (8 per core TPU) e le dimensioni della 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 della funzionalità devono essere un multiplo di 8.
L'utilizzo di una dimensione batch di 1024 e di dimensioni delle caratteristiche che sono un multiplo di 128 genera la migliore efficienza, 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, prendi in considerazione la seguente serie di operazioni:
tmp = tf.add(x, y)
result = tf.multiply(tmp, z)
Questo codice è più o meno equivalente al seguente pseudo codice:
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 si verificano contemporaneamente:
for (i = 0; i < element_count; i++) {
result = (x[i] + y[i]) * z[i];
}
In questo esempio, il numero di corse di andata e ritorno in memoria è ridotto e XLA non deve allocare spazio per 'tmp'.
Fusion è un'ottimizzazione fondamentale e offre vantaggi a Cloud TPU in diversi modi:
- Riduce il trasferimento della memoria eliminando la necessità di archiviare i risultati intermedi nella memoria principale, che è lenta.
- Permette un maggiore utilizzo delle unità hardware che altrimenti sarebbero inutilizzate.
- Può ridurre l'utilizzo della memoria di un modello, poiché è necessario eseguire contemporaneamente un numero inferiore di buffer.
Trasmetto
La trasmissione implicitamente si verifica quando due tensori con forme diverse, ma compatibili, vengono combinate.
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 ulteriori dettagli, consulta la guida alla trasmissione array.
Spesso le trasmissioni possono essere fuse con i consumatori, ma forzare la trasmissione può ridurre le prestazioni e aumentare l'utilizzo della memoria.
Nell'esempio seguente, la trasmissione implicita nell'aggiunta di un vettore e di una matrice non può essere unita con l'argmax risultante in una trasmissione materializzata:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`