Guida alle prestazioni di Cloud TPU

Il primo passo per risolvere i problemi delle prestazioni della TPU è profilare il modello. Per saperne di più 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.

  1. Il modello è associato all'input

    Le TPU eseguono i calcoli molto rapidamente. Per garantire che la TPU non sia inattiva, è importante assicurare che il caricamento di un flusso costante di dati sulla TPU sia stabile. La modalità dipende da come carichi e pre-elabora il set di dati. Ad esempio, puoi leggere i file di dati in parallelo utilizzando tf.data.TFRecordset() e il parametro num_parallel_reads.

  2. La dimensione del batch è troppo ridotta a causa dello sharding (divisione 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 batch globale pari a 128, ogni core riceve una dimensione batch pari a 16 (128 / 8).

    Per un utilizzo ottimale della memoria, utilizza la dimensione batch più grande adatta alla memoria TPU. Ogni core TPU utilizza registri vettoriali bidimensionali 8 X 128 per l'elaborazione delle moltiplicazioni delle matrici. In generale, le dimensioni del batch devono essere divisibili uniformemente 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. XLA fa parte del codebase standard di TensorFlow, ma può essere utilizzato anche sui modelli PyTorch e JAX. I modelli per Cloud TPU vengono convertiti in un grafico XLA, che viene poi compilato da XLA in una TPU eseguibile. Per saperne di più su XLA, consulta XLA: Optimizing Compiler for Machine Learning.

Spaziatura interna

Per utilizzare la memoria TPU in modo efficiente, struttura i dati in modo che possano essere suddivisi in blocchi di 128 x 8. Quando i dati per un calcolo matriciale non riempiono un intero blocco 128 x 8, il compilatore XLA padroneggia i tensori. La spaziatura interna presenta due svantaggi:

  1. I tensori imbottiti sottoutilizzano il core TPU.
  2. La spaziatura interna aumenta la quantità di spazio di archiviazione su chip richiesta per un tensore e può causare un errore di esaurimento della memoria.

Mentre la spaziatura interna viene eseguita automaticamente dal compilatore XLA quando necessario, puoi determinare la quantità di spaziatura interna eseguita utilizzando lo strumento di visualizzazione della memoria. Puoi evitare la spaziatura interna scegliendo le dimensioni dei tensori adatte alla TPU.

Dimensioni Tensor

Il compilatore XLA arrotonda le dimensioni dei tensori archiviati nella memoria HBM 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 dispone di tensori in memoria per massimizzare l'efficienza di calcolo e ridurre al minimo il riempimento. Per ridurre al minimo l'overhead di memoria e massimizzare l'efficienza di calcolo, deve essere vera una delle seguenti condizioni:

  1. Le dimensioni totali del batch devono essere un multiplo di 64 (8 per core TPU), mentre le dimensioni delle funzionalità devono essere un multiplo di 128.

  2. 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'utilizzo di dimensioni batch 1024 e dimensioni della funzionalità che sono un multiplo di 128 consentono di ottenere la migliore efficienza, anche se questo potrebbe non essere possibile per tutti i modelli.

Fusione

Fusion è una tecnica generale utilizzata dal compilatore XLA per ottimizzare i programmi. Un'operazione fondata è la combinazione di più operazioni costitutive da eseguire in combinazione.

Prendi in considerazione, ad esempio, la seguente serie di operazioni:

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Questo codice è all'incirca 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 round trip di memoria è ridotto e XLA non deve allocare spazio per "tmp".

Fusion è un'ottimizzazione fondamentale e offre vantaggi alla Cloud TPU in diversi modi:

  • Riduce i trasferimenti di memoria eliminando la necessità di archiviare i risultati intermedi nella memoria principale, che è lenta.
  • Consente un utilizzo maggiore di unità hardware che altrimenti non verrebbero utilizzate.
  • Può ridurre l'utilizzo della memoria da parte di un modello, in quanto devono essere attivi contemporaneamente un numero inferiore di buffer.

Trasmissione

La trasmissione si verifica implicitamente quando vengono combinati due tensori con forme diverse, ma compatibili.

Ad esempio, tf.add(vector, matrix) richiede che il vettore venga trasmesso nella forma della matrice. Il risultato dell'operazione ha la stessa forma della matrice. Per ulteriori dettagli, consulta la guida alla trasmissione di array.

Sebbene le trasmissioni possano spesso essere fonde con i consumatori, l'applicazione forzata di una trasmissione potrebbe comportare prestazioni scadenti e un maggiore utilizzo della memoria.

Nell'esempio seguente, la trasmissione implicita nell'aggiunta di un vettore e una matrice non può essere fusa con l'argmax, generando una trasmissione materializzata:

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`