TPU v4
Questo documento descrive l'architettura e le configurazioni supportate di Cloud TPU v4.
Architettura di sistema
Ogni chip TPU v4 contiene due TensorCore. Ogni Tensor Core ha quattro unità di moltiplicazione a matrice (MXU), un'unità vettoriale e un'unità scalare. La seguente tabella mostra le specifiche chiave per unpod di TPUU v4.
| Specifiche principali | Valori del pod v4 |
|---|---|
| Picco di calcolo per chip | 275 teraflop (bf16 o int8) |
| Capacità e larghezza di banda HBM2 | 32 GiB, 1200 GBps |
| Potenza minima/media/massima misurata | 90/170/192 W |
| Dimensioni pod di TPU | 4096 chip |
| Topologia di interconnessione | Mesh 3D |
| Picco di computing per pod | 1,1 exaflop (bf16 o int8) |
| Larghezza di banda di All-reduce per pod | 1,1 PB/s |
| Larghezza di banda bisezionale per pod | 24 TB/s |
Il seguente diagramma illustra un chip TPU v4.
Per ulteriori informazioni sui dettagli dell'architettura e sulle caratteristiche di prestazioni della TPU v4, consulta TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings.
Mesh 3D e toro 3D
Le TPU v4 hanno una connessione diretta ai chip vicini più vicini in 3 dimensioni, il che si traduce in una mesh 3D di connessioni di rete. Le connessioni possono essere configurate come un toro 3D su sezioni in cui la topologia, AxBxC, è 2A=B=C o 2A=2B=C, dove ogni dimensione è un multiplo di 4. Ad esempio, 4x4x8, 4x8x8 o 12x12x24. In generale, le prestazioni di una configurazione a toro 3D saranno migliori rispetto a una configurazione a mesh 3D. Per saperne di più, consulta Topologie di tori attorcigliati.
Vantaggi in termini di prestazioni della TPU v4 rispetto alla v3
Questa sezione mostra un modo efficiente in termini di memoria per eseguire uno script di addestramento di esempio su TPU v4, nonché i miglioramenti delle prestazioni per TPU v4 rispetto a TPU v3.
Sistema di memoria
L'accesso alla memoria non uniforme (NUMA) è un'architettura di memoria del computer per le 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. I nodi NUMA sono collegati ad altri nodi NUMA direttamente adiacenti. Una CPU di 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 con più CPU può inserire i dati necessari a una CPU all'interno del relativo nodo NUMA, aumentando la velocità effettiva della memoria. Per saperne di più su NUMA, consulta la pagina Non Uniform Memory Access su Wikipedia.
Puoi sfruttare i vantaggi della località NUMA associando lo script di addestramento al nodo NUMA 0.
Per abilitare il binding dei nodi NUMA:
Installa lo strumento a riga di comando numactl. numactl ti consente di eseguire processi con una specifica policy di pianificazione NUMA o di posizionamento della memoria.
$ sudo apt-get update $ sudo apt-get install numactl
Associa il codice dello script al nodo NUMA 0. Sostituisci your-training-script con il percorso dello script di addestramento.
$ numactl --cpunodebind=0 python3 your-training-script
Abilita l'associazione di nodi NUMA se:
- Se il tuo workload dipende molto dai workload della CPU (ad esempio, classificazione delle immagini, workload di raccomandazione) indipendentemente dal framework.
- Se utilizzi una versione runtime della TPU senza il suffisso -pod (ad esempio,
tpu-vm-tf-2.10.0-v4).
Altre differenze del sistema di memoria:
- I chip TPU v4 hanno uno spazio di memoria HBM unificato da 32 GiB sull'intero chip, il che consente una migliore coordinazione 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 lo striding ad alte prestazioni con granularità di 512 byte.
TensorCores
- Il doppio delle MXU e una frequenza di clock più elevata che offre 275 TFLOPS massimi.
- Larghezza di banda di trasposizione e permutazione 2x.
- Modello di accesso alla memoria load-store per la memoria comune (Cmem).
- Larghezza di banda di caricamento dei pesi MXU più veloce e supporto della modalità a 8 bit per consentire batch più piccoli e una latenza di inferenza migliorata.
Interconnessione tra chip
Sei link di interconnessione per chip per abilitare topologie di rete con diametri di rete più piccoli.
Altro
- Interfaccia PCIE gen3 x16 all'host (connessione diretta).
- Modello di sicurezza migliorato.
- Maggiore efficienza energetica.
Configurazioni
Un pod TPU v4 è composto da 4096 chip interconnessi con link riconfigurabili ad alta velocità. Il networking flessibile della TPU v4 ti consente di connettere i chip in una sezione delle stesse dimensioni in più modi. Quando crei una sezione TPU, specifichi la versione di TPU e il numero di risorse TPU richieste.
Specifichi una topologia TPU utilizzando una tupla di tre elementi, AxBxC, dove A<=B<=C e A, B, C sono tutti <= 4 o tutti multipli interi di 4. I valori A, B e C sono il numero di chip in ciascuna delle tre dimensioni.
Le topologie in cui 2A=B=C o 2A=2B=C hanno anche varianti di topologia ottimizzate per la comunicazione all-to-all, ad esempio 4×4×8, 8×8×16 e 12×12×24. Queste sono note come topologie toroidi attorcigliati.
Le seguenti illustrazioni mostrano alcune topologie TPU v4 comuni.
Le sezioni più grandi possono essere create da uno o più "cubi" di chip 4x4x4.
Per saperne di più sulla gestione delle TPU, consulta Gestire le TPU. Per saperne di più sull'architettura di sistema di Cloud TPU, consulta la sezione Architettura di sistema.
Topologie di tori attorcigliati
Alcune forme di sezioni di tori 3D v4 hanno la possibilità di utilizzare quella che è nota come topologia di toro attorcigliato. Ad esempio, due cubi v4 possono essere disposti come una sezione 4x4x8 o 4x4x8_twisted. Le topologie contorte offrono una larghezza di banda di bisezione significativamente superiore. Ad esempio, una sezione con topologia 4x4x8_twisted fornisce un aumento teorico del 70% della larghezza di banda di bisezione rispetto a una sezione 4x4x8 non twisted. L'aumento della larghezza di banda di bisezione è utile per i carichi di lavoro che utilizzano pattern di comunicazione globali. Le topologie contorte possono migliorare le prestazioni della maggior parte dei modelli, con i carichi di lavoro di incorporamento TPU di grandi dimensioni che ne traggono il massimo vantaggio.
Per i carichi di lavoro che utilizzano il parallelismo dei dati come unica strategia di parallelismo, le topologie contorte potrebbero avere un rendimento leggermente migliore. Per gli LLM, il rendimento con una topologia twisted può variare a seconda del tipo di parallelismo (DP, MP e così via). La best practice consiste nell'addestrare l'LLM con e senza una topologia contorta per determinare quale offre le prestazioni migliori per il modello. Alcuni esperimenti sul modello FSDP MaxText hanno registrato miglioramenti di 1-2 MFU utilizzando una topologia twisted.
Il vantaggio principale delle topologie twisted è che trasformano una topologia torus asimmetrica (ad esempio, 4×4×8) in una topologia simmetrica strettamente correlata. La topologia simmetrica presenta molti vantaggi:
- Bilanciamento del carico migliorato
- Larghezza di banda bisezionale più elevata
- Percorsi dei pacchetti più brevi
Questi vantaggi si traducono in un miglioramento del rendimento per molti modelli di comunicazione globali.
Il software TPU supporta i tori attorcigliati sulle sezioni in cui la dimensione di ogni dimensione è uguale o doppia rispetto alla dimensione più piccola. Ad esempio, 4x4x8, 4×8×8 o 12x12x24.
Ad esempio, considera questa topologia a toro 4×2 con TPU etichettate con le relative coordinate (X,Y) nella sezione:
I bordi di questo grafico della topologia vengono mostrati come bordi non diretti per chiarezza. In pratica, ogni bordo è una connessione bidirezionale tra le TPU. I bordi tra un lato di questa griglia e il lato opposto sono chiamati bordi di avvolgimento, come indicato nel diagramma.
Se torciamo questa topologia, otteniamo una topologia di toro ritorto 4×2 completamente simmetrica:
L'unica differenza tra questo diagramma e quello precedente sono i bordi di avvolgimento Y. Anziché connettersi a un'altra TPU con la stessa coordinata X, sono state spostate per connettersi alla TPU con coordinata X+2 mod 4.
Lo stesso concetto si generalizza a dimensioni diverse e a numeri diversi di dimensioni. La rete risultante è simmetrica, a condizione che ogni dimensione sia uguale o doppia rispetto alla dimensione più piccola.
La seguente tabella mostra le topologie twisted supportate e un aumento teorico della larghezza di banda di bisezione rispetto alle topologie untwisted.
| Topologia contorta | Aumento teorico della larghezza di banda della bisezione rispetto a un toro non attorcigliato |
|---|---|
| 4×4×8_twisted | ~70% |
| 8x8x16_twisted | |
| 12×12×24_twisted | |
| 4×8×8_twisted | ~40% |
| 8×16×16_twisted |
Varianti della topologia TPU v4
Alcune topologie contenenti lo stesso numero di chip possono essere disposte in modi diversi. Ad esempio, una sezione TPU con 512 chip (1024 TensorCore) può essere configurata utilizzando le seguenti topologie: 4x4x32, 4x8x16 o 8x8x8. Una sezione TPU con 2048 chip (4096 TensorCore) offre ancora più opzioni di topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16.
La topologia predefinita associata a un determinato numero di chip è quella più simile a un cubo. Questa forma è probabilmente la scelta migliore per l'addestramento ML data-parallel. Altre topologie possono essere utili per i workload con più tipi di parallelismo (ad esempio, parallelismo del modello e dei dati o partizionamento spaziale di una simulazione). Questi workload hanno un rendimento ottimale se la topologia corrisponde al parallelismo utilizzato. Ad esempio, il parallelismo del modello a 4 vie sulla dimensione X e il parallelismo dei dati a 256 vie sulle dimensioni Y e Z corrisponde a una topologia 4x16x16.
I modelli con più dimensioni di parallelismo funzionano meglio se le dimensioni di parallelismo sono mappate alle dimensioni della topologia TPU. Si tratta in genere di modelli linguistici di grandi dimensioni (LLM) paralleli per dati e modelli. Ad esempio, per una sezione TPU v4 con topologia 8x16x16, le dimensioni della topologia TPU sono 8, 16 e 16. È più efficiente utilizzare il parallelismo del modello a 8 o 16 vie (mappato a una delle dimensioni della topologia TPU fisica). Un parallelismo del modello a 4 vie non sarebbe ottimale con questa topologia, poiché non è allineato a nessuna delle dimensioni della topologia TPU, ma sarebbe ottimale con una topologia 4x16x32 sullo stesso numero di chip.
Le configurazioni TPU v4 sono costituite da due gruppi: quelle con topologie inferiori a 64 chip (topologie piccole) e quelle con topologie superiori a 64 chip (topologie grandi).
Topologie Small v4
Cloud TPU supporta le seguenti sezioni di TPU v4 più piccole di 64 chip, un cubo 4x4x4. Puoi creare queste piccole topologie v4 utilizzando il nome basato su TensorCore (ad esempio v4-32) o la topologia (ad esempio 2x2x4):
| Nome (in base al conteggio di TensorCore) | Numero di chip | Topologia |
| v4-8 | 4 | 2x2x1 |
| v4-16 | 8 | 2x2x2 |
| v4-32 | 16 | 2x2x4 |
| v4-64 | 32 | 2x4x4 |
Topologie v4 di grandi dimensioni
Gli slice di TPU v4 sono disponibili con incrementi di 64 chip, con forme che sono
multipli di 4 in tutte e tre le dimensioni. Le dimensioni devono essere in
ordine crescente. Nella tabella seguente sono riportati diversi esempi. Alcune di queste topologie sono "personalizzate" e possono essere create solo utilizzando i flag --type e --topology perché esistono più modi per disporre i chip.
| Nome (in base al conteggio di TensorCore) | Numero di chip | Topologia |
| v4-128 | 64 | 4x4x4 |
| v4-256 | 128 | 4x4x8 |
| v4-512 | 256 | 4x8x8 |
topologia personalizzata: devono essere utilizzati i flag --type e --topology |
256 | 4x4x16 |
| v4-1024 | 512 | 8x8x8 |
| v4-1536 | 768 | 8x8x12 |
| v4-2048 | 1024 | 8x8x16 |
topologia personalizzata: devono essere utilizzati i flag --type e --topology |
1024 | 4x16x16 |
| v4-4096 | 2048 | 8x16x16 |
| … | … | … |