Guida al rendimento di Cloud TPU

Il primo passo per risolvere i problemi di rendimento della TPU è profilare il modello. Per ulteriori informazioni sull'acquisizione di un profilo di rendimento, consulta Profilare il modello.

Rendimento del modello TPU

Questa sezione descrive i problemi generali che possono ridurre il rendimento del modello e come risolverli.

Modelli con vincoli di input

Le TPU eseguono i calcoli molto rapidamente. Per assicurarti che la TPU non sia inattiva, è importante che venga caricato un flusso costante di dati sulla TPU. La procedura dipende da come carichi ed esegui la pre-elaborazione del set di dati. Ad esempio, puoi leggere i file di dati in parallelo utilizzando tf.data.TFRecordset() e il num_parallel_reads parametro.

Dimensione del batch ridotta a causa dello sharding

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 la dimensione del batch più grande che rientra 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.

Ottimizzazione della gestione della memoria

Puoi utilizzare le variabili di ambiente correlate alla memoria per ottimizzare i comportamenti di runtime di basso livello.

TPU_PREMAPPED_BUFFER_SIZE

TPU_PREMAPPED_BUFFER_SIZE imposta la dimensione del buffer di memoria host (in byte) che viene 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 TPU_PRE_MAPPED_BUFFER_SIZE validi.

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

L'aumento di questa dimensione può potenzialmente migliorare il rendimento del trasferimento dei dati tra l'host e il dispositivo TPU, in particolare per i workload 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.

Risolvere i problemi di memoria

Se la regione del buffer pre-mappato non è abbastanza grande da allocare la memoria durante il runtime del programma, il workload non riuscirà e restituirà un errore RESOURCE_EXHAUSTED simile a:

"Allocating buffer from premmaped region failed with: RESOURCE_EXHAUSTED: Attempting to allocate allocation_size. That was not possible. There are available_size free."

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

Per diagnosticare il problema, esamina i log del runtime 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 a:

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

Questo output indica il tempo impiegato per inizializzare la TPU e la dimensione del buffer pre-mappato.

Impostare la dimensione del buffer

Se il buffer pre-mappato è troppo piccolo o troppo grande, puoi impostare manualmente la dimensione del buffer utilizzando le seguenti variabili di ambiente.

• TPU_PREMAPPED_BUFFER_SIZE: imposta la dimensione totale (in byte) della regione del buffer pre-mappato.

• TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: imposta la dimensione massima di un singolo buffer che può essere allocato dalla regione pre-mappata.

Ad esempio, puoi:

export TPU_PREMAPPED_BUFFER_SIZE=4294967296

per impostare la dimensione del buffer e:

export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES

per attivarlo.

Questa esportazione imposta la dimensione sul valore predefinito.

Modifica il valore di TPU_PREMAPPED_BUFFER_SIZE se sospetti che il trasferimento dei dati host-dispositivo sia un collo di bottiglia. Monitora l'utilizzo della memoria host e il rendimento del modello per trovare un equilibrio ottimale. Il valore predefinito è in genere sufficiente per la maggior parte dei casi d'uso.

Configurazione di tcmalloc

La libreria tcmalloc viene utilizzata per impostazione predefinita sulle VM Cloud TPU per migliorare il rendimento dei modelli con allocazioni di memoria di dimensioni considerevoli e frequenti. Questa impostazione viene configurata tramite la variabile di ambiente LD_PRELOAD.

Tuttavia, per alcuni workload (ad esempio, DLRM con allocazioni di tabelle di incorporamento molto grandi), tcmalloc può causare un rallentamento. In questi casi, puoi ripristinare la funzione malloc standard annullando l'impostazione della variabile LD_PRELOAD nella sessione della shell prima di eseguire lo script di addestramento:

unset LD_PRELOAD

Ottimizzazioni del rendimento della rete

Le sezioni seguenti descrivono come ottimizzare il rendimento della rete configurando l'unità massima di trasmissione (MTU) e utilizzando più NIC per gli ambienti Multislice.

Configurare l'MTU

Per un rendimento ottimale della rete, utilizza una rete con MTU (unità massima di trasmissione) di 8896.

Per impostazione predefinita, un Virtual Private Cloud (VPC) fornisce solo un MTU di 1460 byte, che offre un rendimento della rete non ottimale. Puoi impostare l'MTU di una rete VPC su qualsiasi valore compreso tra 1300 e 8896 byte (inclusi). Le dimensioni MTU personalizzate comuni sono 1500 byte (Ethernet standard) o 8896 byte (il massimo possibile). Per ulteriori informazioni, consulta Dimensioni MTU della rete VPC valide.

Per ulteriori informazioni sulla modifica dell'impostazione MTU per una rete esistente o predefinita, consulta Modificare l'impostazione MTU di una rete VPC network.

Utilizzare l'opzione multi-NIC per Multislice

Quando si addestrano modelli di grandi dimensioni in ambienti Multislice costituiti da migliaia di chip TPU, la comunicazione tra le slice sulla rete del data center (DCN) può rappresentare un collo di bottiglia. Per migliorare la larghezza di banda di rete per i workload con vincoli di rete, puoi utilizzare multi-NIC per aumentare il numero di interfacce di rete sulle VM TPU. Quando utilizzi multi-NIC, a ogni VM TPU vengono allocate interfacce di rete aggiuntive, ognuna connessa a una rete VPC univoca, aumentando la velocità effettiva complessiva della rete. Le NIC aggiuntive devono trovarsi in intervalli IP reciprocamente esclusivi.

Per ulteriori informazioni sull'attivazione di più reti quando utilizzi Google Kubernetes Engine (GKE), consulta Migliorare il rendimento della rete senza hostNetwork su TPU Trillium o Ironwood (TPU7x). Per un esempio di utilizzo di multi-NIC con XPK, consulta Creare un cluster con supporto multi-NIC utilizzando XPK.

Ottimizzazioni del compilatore XLA

XLA è un compilatore per il machine learning che può produrre file binari per TPU, CPU, GPU e altre piattaforme. I modelli scritti in PyTorch, JAX o TensorFlow generano un grafico di operazioni di alto livello (HLO) XLA, che XLA compila in un eseguibile TPU. Per ulteriori informazioni su XLA, consulta XLA: compilatore di ottimizzazione 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 128 x 8. Quando i dati per un calcolo di matrici non riempiono un intero blocco 128 x 8, il compilatore XLA aggiunge la spaziatura interna ai tensori. La spaziatura interna presenta due svantaggi:

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

Sebbene la spaziatura interna venga 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 dimensioni dei tensori adatte alla TPU.

Dimensioni dei tensori

Per ottenere il massimo numero di FLOP, le dimensioni della moltiplicazione di matrici devono essere maggiori della dimensione MXU per la versione TPU in uso. La dimensione MXU è 256 x 256 per v6e e TPU7x (Ironwood) e 128 x 128 per le versioni precedenti a v6e. Per ulteriori informazioni, consulta Architettura di sistema di Cloud TPU.

Dimensione del batch

Il compilatore XLA arrotonda le dimensioni dei tensori archiviati nella memoria HBM della TPU per eseguire i calcoli in modo più efficiente. Questa spaziatura interna avviene in modo trasparente a livello hardware e non influisce sui risultati. Tuttavia, in alcuni casi la spaziatura interna 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 computazionale e ridurre al minimo la spaziatura interna. Per ridurre al minimo l'overhead di memoria e massimizzare l'efficienza computazionale, deve essere vera una delle seguenti condizioni:

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 dimensioni delle funzionalità devono essere un multiplo di 8.

L'utilizzo di una dimensione del batch di 1024 e di dimensioni delle funzionalità che sono un multiplo di 128 comporta la massima efficienza, anche se ciò potrebbe non essere possibile per tutti i modelli.

Nota: Dimensione della funzionalità si riferisce alla dimensione nascosta di un livello completamente connesso o al numero di canali di output in una convoluzione. Non tutti i livelli possono rispettare questa regola, in particolare il primo e l'ultimo livello della rete. Non è un problema, la maggior parte dei modelli richiede una certa quantità di spaziatura interna.

Fusione

La fusione è una tecnica generale utilizzata dal compilatore XLA per ottimizzare i programmi. Un'operazione unita è 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 agli 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 di 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 un numero inferiore di buffer sia attivo 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 alla 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 unite ai relativi consumatori, forzare una trasmissione può comportare un rendimento scadente e un aumento della memoria utilizzata.

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

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

Consigli sul rendimento per l'architettura dual-chiplet di Ironwood

Il modello di programmazione Ironwood consente di accedere a due dispositivi TPU anziché alla singola architettura di core logico (nota anche come MegaCore) utilizzata nelle generazioni precedenti (TPU v4 e v5p). Questa modifica migliora il rapporto costo-efficacia e l'efficienza della produzione del chip. Sebbene si tratti di un cambiamento architettonico, il nuovo design garantisce che tu possa riutilizzare i modelli software esistenti con modifiche minime.

Per ottenere il miglior rendimento con l'architettura dual-chiplet, ti consigliamo i seguenti approcci:

• Utilizza il parallelismo dei tensori tra i chiplet: l'interfaccia D2D a larghezza di banda elevata è progettata per un parallelismo dei tensori efficiente. Ti consigliamo di suddividere i tensori tra i due dispositivi on-chip.

• Utilizza i collettivi gerarchici: per massimizzare l'efficienza della comunicazione, sfrutta la gerarchia di rete a due livelli: il collegamento D2D ultraveloce tra i chiplet on-chip e i collegamenti ICI veloci all'interno di una slice. Quando utilizzi il parallelismo automatico con SPMD (single program, multiple data), il compilatore XLA gestisce questa operazione generando automaticamente operazioni collettive gerarchiche. Quando partizioni manualmente il modello, devi anche progettare i pattern di comunicazione in base a questa gerarchia. Dai la priorità alla comunicazione tra i due dispositivi sullo stesso chip prima di comunicare con i dispositivi su altri chip.

• Sovrapponi la comunicazione al calcolo: per massimizzare l'utilizzo dell'hardware, trasferisci le operazioni di comunicazione collettiva, come all-reduce, su SparseCore. Queste operazioni, che non sono vincolate all'unità di moltiplicazione di matrici (MXU), possono essere eseguite su SparseCore contemporaneamente mentre TensorCore continua il calcolo. Questa tecnica può recuperare alcuni dei vantaggi in termini di rendimento che erano intrinseci alle operazioni unite nella precedente architettura MegaCore.

• Scarica su SparseCore per gli incorporamenti: nella progettazione dual-chiplet, le tabelle di incorporamento potrebbero essere partizionate nell'HBM di entrambi i chiplet. Per evitare un peggioramento del rendimento dovuto alla mancanza di condivisione della memoria, scarica le operazioni di raccolta degli incorporamenti su SparseCore. Questa strategia utilizza l'interconnessione D2D ad alta velocità per trasferire in modo efficiente i vettori di incorporamento tra i chiplet. Per ulteriori informazioni su SparseCore e sui modelli di incorporamento, consulta Un'analisi approfondita di SparseCore per i modelli di incorporamento di grandi dimensioni (LEM).

Per una guida completa all'ottimizzazione del rendimento su TPU7x, consulta Ottimizzazioni del rendimento di TPU7x (Ironwood).