Guida al rendimento di Cloud TPU

Il primo passaggio per risolvere i problemi di prestazioni della TPU è profilare il modello. Per saperne di più sull'acquisizione di un profilo di rendimento, consulta la sezione 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.

  1. Il modello è vincolato all'input

    Le TPU eseguono i calcoli molto velocemente. Per assicurarti che la TPU non sia inattiva, è importante assicurarsi che venga caricato un flusso costante di dati sulla TPU. La modalità dipende da come 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.

  2. La dimensione del batch è troppo piccola 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 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ù grandi che rientrano nella memoria TPU. Ogni core TPU utilizza registri vettoriali bidimensionali 8 x 128 per l'elaborazione delle moltiplicazioni di matrici. In generale, la dimensione del batch deve essere divisibile per 8 o 128.

  3. Ottimizzazione della gestione della memoria

    Puoi utilizzare le variabili di ambiente TPU_PREMAPPED_BUFFER_SIZE per ottimizzare i comportamenti di runtime di basso livello.

  • Descrizione: TPU_PREMAPPED_BUFFER_SIZE imposta le dimensioni del buffer di memoria dell'host (in byte) pre-mappato e bloccato per l'utilizzo da parte del runtime TPU per i trasferimenti di dati (ad esempio, DMA). Il valore predefinito è 4294967296 byte. Il valore deve essere un multiplo di 2^12 (4 KB = 4 * 1024 byte = 4096 = 2^12).

    I seguenti esempi sono valori validi di TPU_PRE_MAPPED_BUFFER_SIZE.

        17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
    40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).

  • Impatto:l'aumento di questa dimensione può potenzialmente migliorare le prestazioni di trasferimento dei dati tra l'host e il dispositivo TPU, in particolare per i carichi di lavoro con tensori di grandi dimensioni o comunicazioni host-dispositivo frequenti. Tuttavia, aumenta anche la quantità di memoria host bloccata, riducendo la memoria disponibile per altri processi.

    Dimensione del buffer

    Se la regione buffer pre-mappata non è abbastanza grande per allocare memoria durante l'esecuzione del programma, il workload non andrà a buon fine e restituirà un errore RESOURCE_EXHAUSTED simile a:

    "Allocating buffer from premmaped region failed with: RESOURCE_EXHAUSTED: Attempting to allocate allocation_size. Non è stato possibile. Ci sono available_size posti liberi."

    Se il buffer è eccessivamente grande, l'inizializzazione della TPU può richiedere molto più tempo (potenzialmente più di 15 secondi), facendo sembrare che la TPU sia bloccata.

    Per diagnosticare il problema, esamina i log di runtime della TPU. Questi log descrivono in dettaglio le operazioni eseguite, inclusa la pre-mappatura dei buffer. Puoi trovare i log in /tmp/tpu_logs/tpu_driver.INFO o stamparli direttamente nella console impostando la variabile di ambiente TPU_STDERR_LOG_LEVEL=0. Questa impostazione genererà un output simile al seguente:

     I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
 I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
 I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
 I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
 ```

This output will tell you how long it took to initialize the TPU and
the size of the premapped buffer.

  • Utilizzo:se il buffer premappato è troppo piccolo o troppo grande, puoi impostare manualmente le dimensioni del buffer utilizzando le seguenti variabili di ambiente.

    TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
pre-mapped buffer region.
TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
a single buffer that can be allocated from the pre-mapped region.

    Ad esempio, puoi:

     export TPU_PREMAPPED_BUFFER_SIZE=4294967296

    per impostare le dimensioni del buffer e:

     export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
 ```
 to enable it.

 This export sets the size to the default.

  • Indicazioni:modifica il valore di TPU_PREMAPPED_BUFFER_SIZE se sospetti che il trasferimento di dati tra host e dispositivo sia un collo di bottiglia. Monitora l'utilizzo della memoria dell'host e le prestazioni del modello per trovare un equilibrio ottimale. Il valore predefinito è in genere sufficiente per la maggior parte dei casi d'uso.

Ottimizzazioni del compilatore XLA

XLA è un compilatore per il machine learning che può produrre file binari per TPU, CPU, GPU e altre piattaforme. Sebbene XLA faccia parte della base di codice TensorFlow standard, può essere utilizzato anche su modelli PyTorch e JAX. I modelli per Cloud TPU vengono tradotti in un grafico XLA, che XLA compila in un eseguibile TPU. 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 della matrice non riempiono un intero blocco 128 x 8, il compilatore XLA esegue il padding dei tensori. L'imbottitura presenta due svantaggi:

  1. I tensori con padding utilizzano in modo insufficiente il core TPU.
  2. Il padding aumenta la quantità di spazio di archiviazione della memoria on-chip necessaria per un tensore e può causare un errore di memoria insufficiente.

Il padding viene eseguito automaticamente dal compilatore XLA quando necessario, ma puoi determinare la quantità di padding eseguita utilizzando lo strumento di visualizzazione della memoria. Puoi evitare il padding scegliendo dimensioni del tensore adatte alla TPU.

Dimensioni tensore

Per ottenere il massimo numero di FLOP, le dimensioni della moltiplicazione della matrice devono essere maggiori della dimensione MXU per la versione TPU che stai utilizzando. La dimensione MXU è 256 x 256 per v6e e 128 x 128 per le versioni precedenti alla v6e. Per maggiori informazioni, consulta la sezione Architettura di sistema di Cloud TPU.

Dimensione del batch

Il compilatore XLA arrotonda per eccesso le dimensioni dei tensori archiviati nella memoria HBM della TPU per eseguire i calcoli in modo più efficiente. Questo riempimento avviene in modo trasparente a livello 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 ridurre al minimo il padding. Per ridurre al minimo l'overhead della memoria e massimizzare l'efficienza computazionale, una delle seguenti condizioni deve essere vera:

  1. 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.

  2. La dimensione totale del batch deve essere un multiplo di 1024 (128 per core TPU) e le dimensioni delle funzionalità devono essere un multiplo di 8.

L'utilizzo di una dimensione batch di 1024 e di dimensioni delle funzionalità che siano un multiplo di 128 consente di ottenere la massima efficienza, anche se ciò 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[i] = tmp[i] * z[i];
    }

Con la fusione, gli accessi all'array avvengono contemporaneamente:

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

In questo esempio, il numero di round trip della memoria viene ridotto e XLA non deve allocare spazio per "tmp".

La fusione è un'ottimizzazione fondamentale e offre diversi vantaggi a Cloud TPU:

  • Riduce i trasferimenti di memoria eliminando la necessità di archiviare i risultati intermedi nella memoria principale, che è lenta.
  • Consente un maggiore utilizzo delle unità hardware che altrimenti non verrebbero utilizzate.
  • Può ridurre l'utilizzo della memoria di un modello, poiché è necessario che siano attivi meno buffer contemporaneamente.

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 alla forma della matrice. Il risultato dell'operazione ha la stessa forma della matrice. Per maggiori dettagli, consulta la guida agli array di trasmissione.

Anche se le trasmissioni possono spesso essere unite ai loro consumatori, forzarne una può comportare prestazioni scarse e un maggiore utilizzo della memoria.

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

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