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 TensorCore ha quattro unità di moltiplicazione a matrice (MXU), un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche principali di un pod di TPU v4.
| Specifica | Valori |
|---|---|
| Picco di computing 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 del pod di TPU | 4096 chip |
| Topologia di interconnessione | Mesh 3D |
| Picco di computing per pod | 1,1 exaflop (bf16 o int8) |
| Larghezza di banda 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 delle prestazioni di TPU v4, consulta TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings.
Mesh 3D e toroide 3D
Le TPU v4 hanno una connessione diretta ai chip adiacenti più vicini in 3 dimensioni, il che si traduce in una mesh 3D di connessioni di rete. Le sezioni di uno o più cubi 4x4x4 completi (dove ogni dimensione A, B e C nella topologia AxBxC è un multiplo di 4) includono link di ritorno a capo che collegano i lati opposti della sezione in un toroide 3D, mentre le sezioni più piccole formano un mesh 3D senza link di ritorno a capo. Sulle sezioni in cui la topologia, AxBxC, è
2A=B=C o 2A=2B=C e ogni dimensione è un multiplo di 4 (ad esempio,
4x4x8, 4x8x8 o 12x12x24), le connessioni possono essere configurate anche come un
toroide 3D intrecciato. In generale, le prestazioni di una configurazione torus 3D sono migliori di quelle di una configurazione mesh 3D. Per saperne di più sulle configurazioni di tori intrecciati, consulta Topologie di tori intrecciati.
Vantaggi in termini di prestazioni di TPU v4 rispetto a v3
Questa sezione mostra un modo efficiente dal punto di vista della 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
Non Uniform Memory Access (NUMA) è un'architettura di memoria per computer 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 connessi ad altri nodi NUMA direttamente adiacenti. Una CPU di un nodo NUMA può accedere alla memoria di 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 suo nodo NUMA, aumentando il throughput della memoria. Per saperne di più su NUMA, consulta Non Uniform Memory Access su Wikipedia.
Puoi sfruttare i vantaggi della località NUMA associando lo script di addestramento al nodo NUMA 0.
Per attivare l'associazione dei nodi NUMA:
Installa lo strumento a riga di comando numactl. numactl consente di eseguire processi con una specifica policy di pianificazione o posizionamento della memoria NUMA.
$ 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
Attiva l'associazione dei nodi NUMA se:
- Se il tuo workload dipende fortemente dai workload della CPU (ad esempio, classificazione delle immagini, workload di suggerimenti), indipendentemente dal framework.
- Se utilizzi una versione di runtime 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 di 32 GiB sull'intero chip, che consente 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 lo striding ad alte prestazioni con granularità di 512 B.
TensorCores
- Il doppio del numero di MXU e una frequenza di clock più elevata che offre un massimo di 275 TFLOPS.
- Larghezza di banda di trasposizione e permutazione raddoppiata.
- Modello di accesso alla memoria di caricamento-archiviazione per la memoria comune (Cmem).
- Larghezza di banda di caricamento del peso MXU più veloce e supporto della modalità a 8 bit per consentire dimensioni dei batch inferiori 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 x16 PCIE gen3 per l'host (connessione diretta).
- Modello di sicurezza migliorato.
- Maggiore efficienza energetica.
Configurazioni
Un pod TPU v4 è composto da 4096 chip interconnessi con link ad alta velocità riconfigurabili. Il networking flessibile di TPU v4 ti consente di connettere i chip in una sezione della stessa dimensione in diversi modi. Quando crei una sezione TPU, specifichi la versione di TPU e il numero di risorse TPU richieste.
Puoi specificare una topologia TPU utilizzando una tupla di tre elementi, AxBxC, dove A<=B<=C e A, B, C sono tutti <= 4 o sono tutti multipli interi di 4. I valori A, B e C sono i conteggi dei 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 toro ritorto.
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 Gestisci le TPU. Per saperne di più sull'architettura di sistema di Cloud TPU, consulta Architettura di sistema.
Topologie di toroidi intrecciati
Per alcune forme di slice 3D, puoi utilizzare una topologia toroide ritorto. Queste topologie offrono una larghezza di banda di bisezione significativamente più elevata. Ad esempio, una topologia twisted 4x4x8 fornisce un aumento teorico del 70% della larghezza di banda di bisezione rispetto a una sezione 4x4x8 non twisted. Questa maggiore larghezza di banda aiuta i workload che utilizzano pattern di comunicazione globali. Le topologie twisted possono migliorare le prestazioni della maggior parte dei modelli, con i workload di TPU embedding di grandi dimensioni che offrono il massimo vantaggio. Il software TPU supporta le topologie twisted su sezioni in cui ogni dimensione è uguale o doppia rispetto alla dimensione più piccola. Ad esempio, 4x4x8, 4×8×8 o 12x12x24. Le topologie intrecciate sono supportate su TPU v4 e TPU v5p tramite l'API Cloud TPU.
Per i workload che utilizzano il parallelismo dei dati come unica strategia di parallelismo, le topologie twisted potrebbero avere un rendimento leggermente migliore. Con i modelli linguistici di grandi dimensioni (LLM), le prestazioni della topologia twisted variano a seconda del tipo di parallelismo utilizzato (ad esempio, parallelismo dei dati o parallelismo dei modelli). Per trovare le prestazioni migliori per il tuo modello, addestra l'LLM sia con che senza una topologia twisted. Alcuni esperimenti sul modello FSDP MaxText hanno mostrato miglioramenti di 1-2 punti percentuali nell'utilizzo dei FLOP del modello (MFU) quando si utilizza una topologia twisted.
Il vantaggio principale delle topologie twisted è che trasformano una topologia torus asimmetrica (ad esempio 4×4×8) in una simmetrica. Una topologia simmetrica offre:
- Bilanciamento del carico migliorato
- Larghezza di banda bisezionale più elevata
- Percorsi dei pacchetti più brevi
Questi vantaggi si traducono in un miglioramento delle prestazioni per molti pattern di comunicazione globali.
Ad esempio, considera questa topologia torus 4×2 con TPU etichettate con coordinate X e Y nella sezione:
Per maggiore chiarezza, il grafico mostra le connessioni come archi non orientati. In pratica, ogni arco è una connessione bidirezionale tra le TPU. I bordi tra un lato di questa griglia e il lato opposto sono bordi di ritorno a capo.
Se questa topologia viene intrecciata, si ottiene una topologia a toro intrecciato 4x2 simmetrica:
La differenza tra questo diagramma con torsione e quello senza torsione risiede nei bordi di avvolgimento Y. Anziché connettersi a un'altra TPU con la stessa coordinata X, questi bordi di avvolgimento si spostano per connettersi alla TPU alla coordinata X+2 mod 4.
Questo principio si applica a diverse dimensioni e numeri di dimensioni. La rete risultante è simmetrica se ogni dimensione è uguale o doppia rispetto alla dimensione più piccola.
La tabella seguente mostra alcune delle topologie twisted supportate e l'aumento teorico della larghezza di banda di bisezione che forniscono rispetto alle topologie non twisted.
| Topologia | Aumento teorico della larghezza di banda di bisezione rispetto a un toro non intrecciato |
|---|---|
| 4×4×8 | ~70% |
| 8x8x16 | |
| 12×12×24 | |
| 4×8×8 | ~40% |
| 8×16×16 |
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 con parallelismo dei dati. Altre topologie possono essere utili per workload con più tipi di parallelismo (ad esempio, parallelismo di modelli e dati o partizionamento spaziale di una simulazione). Questi workload funzionano al meglio se la topologia corrisponde al parallelismo utilizzato. Ad esempio, l'inserimento del parallelismo dei modelli a 4 vie nella dimensione X e del parallelismo dei dati a 256 vie nelle dimensioni Y e Z corrisponde a una topologia 4x16x16.
I modelli con più dimensioni di parallelismo offrono un rendimento migliore se le dimensioni di parallelismo sono mappate alle dimensioni della topologia TPU. In genere si tratta di modelli linguistici di grandi dimensioni (LLM) paralleli di 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 dei modelli a 8 o 16 vie (mappato a una delle dimensioni della topologia TPU fisica). Un parallelismo dei modelli a 4 vie sarebbe non ottimale con questa topologia, poiché non è allineato con 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: quelli con topologie inferiori a 64 chip (topologie piccole) e quelli con topologie superiori a 64 chip (topologie grandi).
Topologie v4 di piccole dimensioni
Cloud TPU supporta le seguenti sezioni 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 numero 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
Le sezioni TPU v4 sono disponibili in incrementi di 64 chip, con forme che sono multipli di 4 su tutte e tre le dimensioni. Le dimensioni devono essere in ordine crescente. Nella tabella seguente sono riportati alcuni esempi. Alcune di queste topologie sono "personalizzate" e possono essere create solo utilizzando i flag --type e --topology perché esiste più di un modo per disporre i chip.
| Nome (in base al numero di TensorCore) | Numero di chip | Topologia |
| v4-128 | 64 | 4x4x4 |
| v4-256 | 128 | 4x4x8 |
| v4-512 | 256 | 4x8x8 |
topologia personalizzata: devi utilizzare i flag --type e --topology |
256 | 4x4x16 |
| v4-1024 | 512 | 8x8x8 |
| v4-1536 | 768 | 8x8x12 |
| v4-2048 | 1024 | 8x8x16 |
topologia personalizzata: devi utilizzare i flag --type e --topology |
1024 | 4x16x16 |
| v4-4096 | 2048 | 8x16x16 |
| … | … | … |