TPU v4

Questo documento descrive l'architettura e le configurazioni supportate 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 seguente tabella mostra le specifiche principali per un pod TPU v4.

Specifiche principali	Valori pod v4
Picco di computing per chip	275 teraflops (bf16 o int8)
Capacità e larghezza di banda di HBM2	32 GiB, 1200 GBps
Potenza minima/media/massima misurata	90/170/192 W
Dimensione pod di TPU	4096 chip
Topologia Interconnect	Mesh 3D
Potenza di picco per pod	1.1 exaflop (bf16 o int8)
Riduci completamente la larghezza di banda per pod	1,1 PB/s
Larghezza di banda bisezione per pod	24 TB/s

Il seguente diagramma illustra un chip TPU v4.

Mesh 3D e toro 3D

Le TPU v4 hanno una connessione diretta ai chip vicini più vicini in 3 che generano un 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 un toro 3D sarà migliore di una configurazione mesh 3D. Per ulteriori informazioni, consulta Topologie di tori contorti.

Vantaggi delle prestazioni di TPU v4 rispetto a v3

Questa sezione descrive i vantaggi in termini di prestazioni di TPU v4

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 sua memoria sono chiamate nodo NUMA. NUMA nodi sono connessi ad altri NUMA nodi 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 memoria all'interno di un nodo NUMA.

Il software in esecuzione su una macchina multi-CPU può collocare i dati necessari a una CPU all'interno del suo nodo NUMA, aumentando la velocità effettiva della memoria. Per ulteriori informazioni su NUMA, Consulta Accesso alla memoria non uniforme su Wikipedia.

Puoi sfruttare i vantaggi della localizzazione NUMA associando lo script di addestramento al nodo NUMA 0.

Per attivare l'associazione di nodi NUMA:

1. Installa lo strumento a riga di comando numactl.

    $ sudo apt-get update
    $ sudo apt-get install numactl

2. Associa il codice dello script al nodo NUMA 0.

    $ numactl --cpunodebind=0 python3 your-training-script

Attiva la mappatura dei nodi NUMA se:

• Se il carico di lavoro dipende fortemente dai carichi di lavoro della CPU (ad esempio, classificazione, carichi di lavoro consigliati) a prescindere dal framework.
• Se utilizzi una versione runtime TPU senza un 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 in tutto il chip, consentendo un migliore coordinamento tra i due TensorCore on-chip.
• Miglioramento delle prestazioni HBM mediante l'utilizzo degli standard e delle velocità di memoria più recenti.
• Profilo di prestazioni DMA migliorato con il supporto integrato per la suddivisione ad alte prestazioni con granularità di 512 B.

TensorCores

• Il doppio del numero di MXU e una frequenza di clock superiore per un massimo di 275 TFLOPS.
• Larghezza di banda per la trasposizione e la permutazione del doppio.
• Modello di accesso alla memoria load-store per la memoria comune (Cmem).
• Maggiore larghezza di banda per il caricamento del peso MXU più veloce e supporto della modalità a 8 bit per consentire dimensioni del batch e una migliore latenza di inferenza.

Interconnessione tra chip

Sei link di interconnessione per chip per abilitare topologie di rete con dimensioni diametri della rete.

Altro

• Interfaccia PCIE 3ª gen. x16 con host (connessione diretta).
• Modello di sicurezza migliorato.
• Efficienza energetica migliorata.

Configurazioni

Un pod TPU v4 è composto da 4096 chip interconnessi con componenti link ad alta velocità. Il networking flessibile di TPU v4 ti consente di collegare i chip in un di una sezione di pod delle stesse dimensioni in vari modi. Quando crei una sezione di pod TPU, specifichi la versione di TPU e il numero di risorse TPU richieste. Quando una sezione di pod TPU v4, puoi specificarne tipo e dimensioni in uno dei due modi seguenti: AcceleratorType e AccleratorConfig.

Nota: AcceleratorType e AccleratorConfig sono mutuamente esclusivi; puoi utilizzarne solo uno durante la creazione di una TPU.

Uso: AcceleratorType

Usa AcceleratorType quando non specifichi una topologia. Per configurare le TPU v4 utilizzando AcceleratorType, utilizza il flag --accelerator-type quando crei la sezione del pod TPU. Imposta --accelerator-type su una stringa contenente la TPU e il numero di TensorCore che vuoi utilizzare. Ad esempio, per creare un'istanza La sezione di pod v4 con 32 TensorCore, useresti --accelerator-type=v4-32.

Il comando seguente crea una sezione di pod di TPU v4 con 512 TensorCore utilizzando il flag --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.17.0-pod-pjrt

Il numero dopo la versione di TPU (v4) specifica il numero di TensorCore. Una TPU v4 ha due TensorCore, quindi il numero di chip TPU sarà 512/2 = 256.

Uso: AcceleratorConfig

Utilizza AcceleratorConfig quando vuoi personalizzare la topologia fisica della sezione TPU. Di solito è necessario per l'ottimizzazione delle prestazioni con il pod più di 256 chip.

Per configurare le TPU v4 utilizzando AcceleratorConfig, utilizza i flag --version e --topology. Imposta --version sulla versione di TPU che vuoi utilizzare e --topology alla disposizione fisica dei chip TPU nella sezione di pod.

Si specifica una topologia TPU utilizzando una topologia a 3 tuple, AxBxC, dove A<=B<=C e A, B, C o tutti <= 4 o sono tutti multipli interi di 4. I valori A, B e C sono il chip conta in ognuna delle tre dimensioni. Ad esempio, per creare un pod v4 sezione con 16 chip, devi impostare --version=v4 e --topology=2x2x4.

Il comando seguente crea una sezione del pod TPU v4 con 128 chip TPU disposti in un array 4x4x8:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.17.0-pod-pjrt

Topologie in cui 2A=B=C o 2A=2B=C hanno anche varianti di topologia ottimizzato per la comunicazione all-to-all, ad esempio 4 x 4 x 8, 8 x 8 x 16 e 12 x 12 x 24. Questi sono noti come topologie tori intrecciati.

Le seguenti illustrazioni mostrano alcune topologie TPU v4 comuni.

Le sezioni di pod più grandi possono essere create da uno o più "cubi" 4x4x4 di chip.

Topologie tori intrecciate

Alcune forme delle sezioni del toro 3D della versione 4 offrono la possibilità di utilizzare la cosiddetta toro ruotato. Ad esempio, due cubi v4 possono essere disposti come una sezione 4x4x8 o 4x4x8_twisted. Le topologie contorte offrono una bisezione significativamente più alta e la larghezza di banda. Una maggiore larghezza di banda bisezione è utile per carichi di lavoro che utilizzano modelli di comunicazione. Le topologie contorte possono migliorare le prestazioni per la maggior parte di modelli, con un vantaggio maggiore per i grandi carichi di lavoro di embedding TPU.

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 contorta può variare a seconda del tipo di parallelismo (DP, MP e così via). La migliore pratica è addestrare il modello LLM con e senza una topologia a spirale per determinare quale offre le migliori prestazioni per il modello. Alcuni esperimenti Il modello FSDP MaxText ha abbiamo visto 1-2 miglioramenti degli MFU usando una topologia contorta.

Il vantaggio principale delle topologie intrecciate è che trasformano una topologia asimmetrica del toro (ad esempio 4 x 4 x 8) in una topologia simmetrica strettamente correlata. La topologia simmetrica offre molti vantaggi:

• Bilanciamento del carico migliorato
• Larghezza di banda bisezionale più elevata
• Percorsi dei pacchetti più brevi

Questi vantaggi si traducono in definitiva in un miglioramento del rendimento per molte aziende globali modelli di comunicazione.

Il software TPU supporta i tori contorti su sezioni in cui le dimensioni di ogni dimensione sono uguali o il doppio della dimensione più piccola. Ad esempio, 4x4x8, 4×8×8 o 12x12x24.

Ad esempio, considera questa topologia toro 4x2 con TPU etichettate con il loro Coordinate (X,Y) nella sezione:

Per maggiore chiarezza, i bordi in questo grafico della topologia sono mostrati come bordi non diretti. In pratica, ogni bordo è una connessione bidirezionale tra TPU. Ci riferiamo alle i bordi tra un lato di questa griglia e il lato opposto come bordi avvolgenti, come indicato nel diagramma.

Ruotando questa topologia, si finisce per ottenere un toro ritorto 4x2 completamente simmetrico topologia:

Tutto ciò che è cambiato tra questo diagramma e quello precedente è la Y a bordi tutto intorno. Invece di 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 generalizza in base a dimensioni diverse e numeri diversi di dimensioni. La rete risultante è simmetrica, a condizione che ogni dimensione sia uguale o doppia rispetto alla dimensione più piccola.

Per informazioni dettagliate su come, consulta l'articolo sull'utilizzo di AcceleratorConfig per specificare una configurazione tori intrecciata durante la creazione di una Cloud TPU.

La tabella seguente mostra le topologie contorte supportate e una di aumento della larghezza di banda di bisezione con loro rispetto alle topologie non intrecciate.

Topologia a spirale	Aumento teorico della bisezione larghezza di banda rispetto a un toro non intrecciato
4×4×8_twisted	Circa il 70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	Circa il 40%
8×16×16_twisted

Nota: la topologia a spirale non verrà utilizzata automaticamente a meno che il suffisso _twisted non sia specificato con il flag --topology. Ad esempio, per specificare una topologia contorta 4x4x8, devi specificare il flag topology come: 4x4x8_twisted.

Varianti di topologia TPU v4

Alcune topologie contenenti lo stesso numero di chip possono essere disposte in diversi in molti modi diversi. Ad esempio, una sezione di pod TPU con 512 chip (1024 TensorCore) può essere configurata utilizzando le seguenti topologie: 4x4x32, 4x8x16 o 8x8x8. Una sezione di pod 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 in parallelo con i dati. Altre topologie possono essere utili per carichi di lavoro con più tipi di parallelismo (ad es. parallelismo di modelli e dati o partizione spaziale di una simulazione). Questi carichi di lavoro offrono prestazioni migliori se la topologia corrisponde al parallelismo utilizzato. Ad esempio, posizionando il parallelismo del modello a 4 vie sulla dimensione X e sui dati a 256 vie il parallelismo delle dimensioni Y e Z corrisponde a una topologia 4x16x16.

I modelli con più dimensioni di parallelismo hanno il rendimento migliore quando le dimensioni del parallelismo sono mappate alle dimensioni della topologia TPU. Di solito si tratta di modelli linguistici di grandi dimensioni (LLM) paralleli di dati e modelli. Ad esempio, per una TPU v4 Sezione di pod con topologia 8x16x16, le dimensioni della topologia TPU sono 8, 16 e 16. È più performante per usare il parallelismo del modello a 8 o 16 vie (mappato a uno dei le dimensioni della topologia TPU fisica). Un parallelismo del modello a 4 vie sarebbe subottimale 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 v4 piccole

Cloud TPU supporta i seguenti slice di pod TPU v4 più piccoli 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

Le sezioni di pod TPU v4 sono disponibili con incrementi di 64 chip, con forme che sono múltiplori di 4 in tutte e tre le dimensioni. Le dimensioni devono essere in in ordine crescente. Nella tabella seguente sono riportati diversi esempi. Alcune di queste sono topologia "personalizzate" che 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
custom topology: devono essere utilizzati i flag --type e --topology	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8 x 8 x 16
custom topology: devono essere utilizzati i flag --type e --topology	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…