Guida alle prestazioni di Cloud TPU
Il primo passo per la risoluzione dei problemi relativi alle prestazioni delle TPU è profilare il modello. Per ulteriori informazioni sull'acquisizione di un profilo di prestazioni, consulta la sezione 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 le modalità per risolverli.
Il modello è associato all'input
Le TPU eseguono i calcoli molto rapidamente. Per assicurarti che il TPU non sia inattivo, è importante assicurarti che venga caricato un flusso di dati costante sulla TPU. Ciò dipende dal modo in cui carichi e preelabora 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.Le dimensioni del batch sono troppo piccole a causa dello sharding (suddividendo i 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 batch globale di 128, ogni core riceve una dimensione del batch di 16 (128 / 8).
Per un utilizzo ottimale della memoria, utilizza le dimensioni del batch più grandi che si adattano alla memoria della TPU. Ogni core TPU utilizza registri vettore 2 x 128 bidimensionali per l'elaborazione delle moltiplicazioni delle matrici. In generale, le dimensioni del batch devono essere uniformemente divisibili per 8 o 128.
Ottimizzazioni del compilatore XLA
XLA è un compilatore per il machine learning in grado di produrre programmi binari per TPU, CPU, GPU e altre piattaforme. Sebbene XLA faccia parte del codebase standard di TensorFlow, può essere utilizzato anche sui modelli PyTorch e JAX. I modelli per Cloud TPU vengono tradotti in un grafico XLA, che XLA successivamente compila con un eseguibile TPU. Per ulteriori informazioni sulla XLA, consulta l'articolo XLA: ottimizzazione del compilatore per il machine learning.
Spaziatura interna
Per utilizzare la memoria TPU in modo efficiente, struttura i dati in modo che possano essere suddivisi in blocchi da 128 x 8. Quando i dati di un calcolo a matrice non occupano un intero blocco di 128 x 8, il compilatore XLA esegue il pad dei tensori. La spaziatura interna presenta due svantaggi:
- I tensori imbottiti sottoutilizzano il nucleo TPU.
- La spaziatura interna aumenta la quantità di memoria interna al chip richiesta per un tensore e può causare un errore di memoria esaurita.
Mentre la spaziatura interna viene eseguita automaticamente dal compilatore XLA, quando necessario, puoi determinare la quantità di spaziatura interna eseguita utilizzando lo strumento visualizzatore di memoria. Puoi evitare la spaziatura interna scegliendo dimensioni tensore adatte a TPU.
Dimensioni Tensor
Il compilatore XLA arrotonda le dimensioni dei tensori memorizzati nella memoria HBM della 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. In alcuni casi, tuttavia, la spaziatura interna può comportare un aumento significativo dell'utilizzo e dei tempi di esecuzione della memoria.
Il runtime TPU dispone dei tensori in memoria per massimizzare l'efficienza di calcolo e ridurre al minimo la spaziatura interna. Per ridurre al minimo l'overhead di memoria e massimizzare l'efficienza di calcolo, deve essere soddisfatta una delle seguenti condizioni:
Le dimensioni totali del batch devono essere un multiplo di 64 (8 per core TPU) e le dimensioni delle dimensioni devono essere un multiplo di 128.
Le dimensioni totali del batch devono essere un multiplo di 1024 (128 per core TPU) e le dimensioni delle funzionalità devono essere un multiplo di 8.
L'uso 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 costitutive da eseguire in combinazione.
Ad esempio, considera la seguente serie di operazioni:
tmp = tf.add(x, y)
result = tf.multiply(tmp, z)
Questo codice è più o meno 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 agli array avvengono contemporaneamente:
for (i = 0; i < element_count; i++) {
result = (x[i] + y[i]) * z[i];
}
In questo esempio, il numero di viaggi andata e ritorno in memoria è ridotto e la XLA non deve allocare alcun spazio per "tmp".
Fusion è un'ottimizzazione fondamentale e offre vantaggi a Cloud TPU in diversi modi:
- Riduce i trasferimenti di memoria eliminando la necessità di archiviare risultati intermedi nella memoria principale, che è lenta.
- Consente un maggiore utilizzo delle unità hardware che altrimenti resterebbero inutilizzate.
- Consente di ridurre l'utilizzo della memoria di un modello, in quanto è necessario creare contemporaneamente meno buffer.
Trasmissione in corso...
La trasmissione implicita avviene 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 agli array di trasmissione.
Sebbene le trasmissioni possano essere spesso fuse dai consumatori, forzarne la trasmissione potrebbero comportare prestazioni scarse e un maggiore utilizzo della memoria.
Nell'esempio seguente, la trasmissione implicita nell'aggiunta di un vettore e della matrice non può essere unita al valore argmax che determina la trasmissione materializzata:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`