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 che venga caricato un flusso costante di dati. 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 batch globale di 128, ogni core riceve una dimensione batch di 16 (128 / 8).

    Per un utilizzo ottimale della memoria, utilizza la dimensione 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.

  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 è sufficientemente 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 la dimensione 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 la dimensione 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 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 in grado di produrre file binari per TPU, CPU, GPU e altre piattaforme. Sebbene XLA faccia parte del codebase 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 della TPU.
  2. Il padding aumenta la quantità di spazio di archiviazione della memoria on-chip richiesta per un tensore e può causare un errore di memoria insufficiente.

Sebbene il padding venga eseguito automaticamente dal compilatore XLA quando necessario, puoi determinare la quantità di padding eseguita utilizzando lo strumento di visualizzazione della memoria. Puoi evitare il padding scegliendo dimensioni dei tensori 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 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 memorizzare 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 scadenti 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)`

Suggerimenti sulle prestazioni per l'architettura dual-chiplet Ironwood

Il modello di programmazione Ironwood consente di accedere a due dispositivi TPU anziché all'architettura a singolo core logico (nota anche come MegaCore) utilizzata nelle generazioni precedenti (TPU v4 e v5p). Questa modifica migliora l'efficacia in termini di costi e l'efficienza della produzione del chip. Sebbene ciò rappresenti 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 dividere i tensori tra i due dispositivi on-chip.

  • Utilizza 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 sezione. Quando utilizzi il parallelismo automatico con SPMD (single program, multiple data), il compilatore XLA gestisce questa operazione per te 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, scarica le operazioni di comunicazione collettiva, come all-reduce, su SparseCores. Queste operazioni, che non sono vincolate all'unità di moltiplicazione matriciale (MXU), possono essere eseguite contemporaneamente sugli SparseCore mentre i TensorCore continuano il calcolo. Questa tecnica può recuperare alcuni dei vantaggi in termini di prestazioni inerenti alle operazioni combinate nella precedente architettura MegaCore.

  • Offload a SparseCore per gli incorporamenti:nella progettazione a doppio chiplet, le tabelle di incorporamento potrebbero essere partizionate nella HBM di entrambi i chiplet. Per evitare il peggioramento delle prestazioni dovuto a questa mancanza di condivisione della memoria, scarica le operazioni di raccolta degli incorporamenti in SparseCore. Questa strategia utilizza l'interconnessione D2D ad alta velocità per trasferire in modo efficiente i vettori di incorporamento tra i chiplet. Per saperne di più su SparseCore e sui modelli di incorporamento, consulta Un'analisi approfondita di SparseCore per i modelli di incorporamento di grandi dimensioni (LEM).

Per saperne di più sull'architettura Ironwood in TPU7x, consulta TPU7x (Ironwood).