Architettura di sistema
Le TPU (Tensor Processing Unit) sono circuiti integrati specifici per le applicazioni (ASIC) progettati da Google per accelerare i carichi di lavoro di machine learning. Cloud TPU è un servizio Google Cloud che rende le TPU disponibili come risorsa scalabile.
Le TPU sono progettate per eseguire rapidamente operazioni di matrice, il che le rende ideali per i carichi di lavoro di machine learning. Puoi eseguire carichi di lavoro di machine learning sulle TPU utilizzando framework come TensorFlow, Pytorch e JAX.
Termini per Cloud TPU
Le seguenti sezioni illustrano i termini e i concetti correlati utilizzati in questo documento.
Chip TPU
Un chip TPU contiene uno o più TensorCore. Ogni TensorCore è costituito da una o più unità MXU (matrici-moltiplica), un'unità vettoriale e un'unità scalare.
Un MXU è composto da 128 x 128 moltiplicatori/accumulatori in una matrice sistolica. Gli MXU forniscono la maggior parte della potenza di calcolo di un ambiente TensorCore. Ogni MXU è in grado di eseguire 16.000 operazioni ad accumulo multiplo per ciclo. Tutti i moltiplicatori prendono gli input bfloat16, ma tutte le accumulazioni vengono eseguite nel formato numerico FP32.
L'unità vettoriale viene utilizzata per calcoli generali quali attivazioni e softmax. L'unità scalare viene utilizzata per il flusso di controllo, il calcolo degli indirizzi di memoria e altre operazioni di manutenzione.
dispositivo TPU
Un dispositivo TPU contiene quattro chip TPU interconnessi tramite interconnessioni ad alta velocità. Il numero di TensorCore in un chip TPU e il numero di MXU in un chip TensorCore dipende dalla versione del dispositivo TPU in uso.
Computer host
Ogni dispositivo TPU è connesso a una macchina host basata su CPU (VM TPU). La macchina host è responsabile del caricamento dei dati da Cloud Storage, della pre-elaborazione dei dati, dell'invio dei dati alla TPU e dell'esecuzione di operazioni che non possono essere eseguite su una TPU.
Singoli dispositivi e pod
Cloud TPU consente di lavorare con un singolo dispositivo TPU o con più dispositivi TPU interconnessi, chiamati pod. Puoi lavorare con un intero pod di TPU (il numero di dispositivi TPU in un pod varia a seconda della versione TPU) o con un sottoinsieme di un pod di TPU chiamato sezione.
Dispositivi TPU singoli
La configurazione TPU per un singolo dispositivo è una scheda autonoma senza connessioni di rete ad altre schede TPU. Le TPU a scheda singola non fanno parte della configurazione di un pod di TPU e non possono comunicare con altri dispositivi TPU.
Sezioni TPU
In una sezione TPU, i chip TPU sono connessi tramite un'interconnessione ad alta velocità. Ogni chip TPU comunica direttamente con gli altri chip nella sezione TPU. Il software TPU gestisce automaticamente la distribuzione dei dati a ogni TensorCore in una sezione. Le sezioni TPU v4 sono disponibili con TensorCore 16, 32, 64, 128, 256, 512, 1024, 2048 o 4096. Le sezioni TPU v3 sono disponibili con TensorCore 32, 128, 512, 1024 o 2048. Le sezioni TPU v2 sono disponibili con 32, 128, 256 o 512 TensorCore.
Versioni TPU
Il layout esatto di una TPU dipende dalla versione della TPU utilizzata. I dettagli dell'architettura e le caratteristiche prestazionali di TPU v2 e v3 sono disponibili in Un supercomputer specifico per il dominio per l'addestramento di reti neurali profonde.
TPU v4
Ogni chip TPU v4 contiene due TensorCore. Ogni TensorCore ha quattro MXU, un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche principali per un pod di TPU v4.
| Specifiche principali | Valori pod v4 |
|---|---|
| Picco di computing per chip | 275 teraflops (bf16 o int8) |
| Capacità e larghezza di banda HBM2 | 32 GiB, 1200 GB/s |
| Potenza min/media/max misurata | 90/170/192 E |
| Dimensioni pod di TPU | 4096 chip |
| Topologia di interconnessione | Mesh 3D |
| Picco di computing per pod | 1.1 exaflops (bf16 o int8) |
| Riduce completamente la larghezza di banda per pod | 1,1 PB/s |
| Larghezza di banda di bisezione per pod | 24 TB/s |
Il seguente diagramma illustra un chip TPU v4.
Mesh 3D e torus 3D
Le TPU v4 hanno una connessione diretta ai chip vicini più vicini in 3 dimensioni, ottenendo un mesh 3D di connessioni di rete. Quando la sezione è uguale o più grande di un cubo singolo, le connessioni possono essere configurate come un torus 3D. Se tutte le dimensioni delle sezioni sono un torus 3D completo (collegamenti a tutto tondo connessi), le prestazioni saranno migliori rispetto a una configurazione mesh 3D.
Topologia v4
Alcune forme del toro 3D v4 hanno la possibilità di utilizzare quella che è nota come topologia toro intrecciata, aumentando la larghezza di banda della sezione trasversale della sezione. Una topologia contorta è una topologia alternativa per le TPU v4.
Alcune topologie v4 hanno la possibilità di una configurazione intrecciata, ad esempio due cubi v4 possono essere disposti come una sezione 4x4x8 o 4x4x8_twisted. Torto tori che offre una larghezza di banda di bisezione significativamente più elevata. Questa larghezza di banda maggiore può essere utile per i carichi di lavoro che contengono pattern di comunicazione globali. È stato osservato che le topologie contorte migliorano le prestazioni per la maggior parte dei modelli, con i grandi vantaggi derivanti dai modelli di incorporamento di grandi dimensioni.
La variante tori contorto è consigliata per l'addestramento di modelli LLM (Large Language Model) (che coinvolge i collettivi di sottogruppi). Il tori contorto sarà probabilmente più vantaggioso per l'incorporamento dei carichi di lavoro. Per i carichi di lavoro con solo il parallelismo dei dati, il tori tortuoso potrebbe avere prestazioni leggermente migliori. Per gli LLM con più forme di parallelismo (DP, MP e così via), i tori intrecciati avranno probabilmente prestazioni inferiori rispetto a un torus normale.
Il vantaggio principale della "torsione" di un toro è che trasforma una topologia toro asimmetrica (ad esempio 4 × 4 × 8) in una topologia strettamente correlata e simmetrica. La topologia simmetrica offre molti vantaggi: migliore bilanciamento del carico, maggiore larghezza di banda di bisezione e route di pacchetti più brevi. Questi vantaggi in definitiva si traducono in un miglioramento delle prestazioni per molti modelli di comunicazione globali.
Il software TPU supporta i tori intrecciati su sezioni in cui la dimensione di ogni dimensione è uguale o doppia della dimensione più piccola, ad esempio 4 x 4 x 8, 4 x 8 x 8 o 12 x 12 x 24.
Ad esempio, considera questa topologia torus 4 × 2 con TPU etichettate con le coordinate (X,Y) nella sezione:
I bordi in questo grafico di topologia sono mostrati come bordi non diretti per chiarezza. In pratica, ogni bordo è una connessione bidirezionale tra le TPU. I pacchetti possono scorrere in entrambe le direzioni attraverso questi bordi. Ci riferiamo ai bordi tra un lato di questa griglia e il lato opposto come bordi avvolgenti, come indicato nel diagramma.
Con la torsione di questa topologia, si ottiene una topologia a toro contorto 4x2 completamente simmetrica:
Sono cambiati solo i bordi del diagramma a Y. Anziché connettersi a un'altra TPU con la stessa coordinata X, sono stati spostati per connettersi alla TPU con la coordinata X+2 mod 4.
La stessa idea si riferisce a dimensioni e numeri diversi. La rete risultante è simmetrica, purché ogni dimensione sia uguale o doppia delle dimensioni della dimensione più piccola.
Consulta la pagina relativa all'utilizzo di AcceleratorConfig per informazioni dettagliate su come specificare una configurazione tori ritorta durante la creazione di una Cloud TPU.
La seguente tabella mostra le topologie intrecciate supportate e un aumento teorico della larghezza di banda di bisezione rispetto alle topologie non distorte.
| Topologia contorta | Aumento teorico della bisezione rispetto a un toru non ritorto |
|---|---|
| 4×4×8_Intrecciata | Circa il 70% |
| 8x8x16_torto | |
| 12×12×24_Intrecciato | |
| 4×8×8_Intrecciata | Circa il 40% |
| 8×16×16_Intrecciato |
TPU v3
Ogni chip TPU v3 contiene due TensorCore. Ogni TensorCore ha due MXU, un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche principali e i relativi valori per un pod di TPU v3.
| Specifiche principali | Valori pod v3 |
|---|---|
| Picco di computing per chip | 123 teraFLOPS (bf16) |
| Capacità e larghezza di banda HBM2 | 32 GiB, 900 GB/s |
| Potenza min/media/max misurata | 123/220/262 E |
| Dimensioni pod di TPU | 1024 chip |
| Topologia di interconnessione | Toru 2D |
| Picco di computing per pod | 126 petaflops (bf16) |
| Riduce completamente la larghezza di banda per pod | 340 TB/s |
| Larghezza di banda di bisezione per pod | 6,4 TB/s |
Il seguente diagramma illustra un chip TPU v3.
Vantaggi delle prestazioni di TPU v4 rispetto a v3
Sistema di memoria:
NUMA (Non Uniform Memory Access) è un'architettura di memoria per computer per macchine con più CPU. Ogni CPU ha accesso diretto a un blocco di memoria ad alta velocità. Una CPU e la relativa memoria sono chiamate nodo NUMA. NUMA nodi sono connessi a nodi NUMA direttamente adiacenti l'uno all'altro. Una CPU da un nodo NUMA può accedere alla memoria in un altro nodo NUMA, ma questo accesso è più lento rispetto all'accesso alla memoria all'interno di un nodo NUMA.
Il software in esecuzione su una macchina multi-CPU può inserire i dati necessari a una CPU all'interno del nodo NUMA, aumentando la velocità effettiva di memoria. Per ulteriori informazioni su NUMA, vedi Non Uniform Memory Access su Wikipedia.
Il runtime TPU v4 non è compatibile con NUMA in modo nativo. Tuttavia, puoi sfruttare i vantaggi della località NUMA associando lo script di addestramento a NUMA Node 0.
Per abilitare l'associazione di nodi NUMA:
Installa lo strumento a riga di comando numactl.
$ $ sudo apt-get update $ $ sudo apt-get install numactl
Usa
numactl --cpunodebind=0quando avvii lo script di addestramento. Questo associa il codice dello script al nodo NUMA 0.$ $ numactl --cpunodebind=0 python3 your-training-script
Abilita l'associazione di nodi NUMA se:
- Se il carico di lavoro dipende fortemente dai carichi di lavoro della CPU (ad esempio classificazione di immagini, carichi di lavoro per suggerimenti) indipendentemente dal framework.
- Se utilizzi una versione di runtime TPU senza un suffisso -pod (ad esempio tpu-vm-tf-2.10.0-v4).
Altre differenze di sistema di memoria:
- I chip TPU v4 hanno uno spazio di memoria HBM unificato di 32 GiB nell'intero chip, consentendo un migliore coordinamento tra i due TensorCore on-chip.
- Prestazioni HBM migliorate grazie agli standard e alle velocità di memoria più recenti.
- Profilo delle prestazioni DMA migliorato con supporto integrato per un avanzamento ad alte prestazioni con granularità a 512 B.
TensorCore:
- Il doppio del numero di MXU e una frequenza di clock superiore per un massimo di 275 TFLOPS.
- Larghezza di banda di trasposizione e permutazione raddoppiata.
- Modello di accesso alla memoria dell'archivio del carico per la memoria comune (Cmem).
- La larghezza di banda per il caricamento del peso MXU più veloce e il supporto della modalità a 8 bit per consentire dimensioni dei batch inferiori e una migliore latenza di inferenza.
Inter-chip Interconnect:
Sei link di interconnessione per chip per consentire topologie di rete con diametri di rete più piccoli.
Altro:
- Interfaccia PCIE gen3 x16 per l'host (connessione diretta).
- Modello di sicurezza migliorato.
- Efficienza energetica migliorata.
Vantaggi delle prestazioni di TPU v3 rispetto a v2
L'aumento dei FLOPS per TensorCore e la capacità di memoria nelle configurazioni TPU v3 possono migliorare le prestazioni dei tuoi modelli nei seguenti modi:
Le configurazioni TPU v3 offrono vantaggi significativi in termini di prestazioni per TensorCore per i modelli legati al calcolo. I modelli associati alla memoria nelle configurazioni TPU v2 potrebbero non raggiungere lo stesso miglioramento delle prestazioni se sono associati anche alla memoria nelle configurazioni TPU v3.
Nei casi in cui i dati non rientrano in memoria nelle configurazioni TPU v2, TPU v3 può offrire prestazioni migliori e una riduzione del ricalcolo dei valori intermedi (rimaterializzazione).
Le configurazioni di TPU v3 possono eseguire nuovi modelli con dimensioni batch non adatte alle configurazioni TPU v2. Ad esempio, TPU v3 potrebbe consentire reti ResNet più profonde e immagini più grandi con RetinaNet.
I modelli quasi legati all'input ("infeed") su TPU v2 perché i passaggi di addestramento sono in attesa di input potrebbero essere associati all'input anche con Cloud TPU v3. La guida sul rendimento della pipeline può aiutarti a risolvere i problemi di inserimento.
Architetture delle VM di Cloud TPU
La modalità di interazione con l'host TPU (e la scheda TPU) dipende dall'architettura della VM TPU che stai utilizzando: nodi TPU o VM TPU.
Nodi TPU
L'architettura del nodo TPU è composta da una VM utente che comunica con l'host TPU su gRPC. Quando utilizzi questa architettura, non puoi accedere direttamente all'host TPU, il che rende difficile il debug degli errori di addestramento e TPU.
VM TPU
L'architettura della VM TPU consente di connettersi tramite SSH alla VM fisicamente connessa al dispositivo TPU. Hai accesso root alla VM, quindi puoi eseguire codice arbitrario. Puoi accedere ai log di debug del compilatore e di runtime e ai messaggi di errore.
