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 seguente tabella mostra le specifiche principali per un pod di TPU v4.

Specifiche principali	Valori del pod v4
Potenza di calcolo di picco per chip	275 teraflops (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 dell'interconnessione	Maglia 3D
Potenza di picco 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 mostra un chip TPU v4.

Mesh 3D e toro 3D

Le TPU v4 hanno una connessione diretta ai chip vicini più vicini in 3 dimensioni, con 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 una configurazione di un toro 3D saranno migliori rispetto a quelle di una configurazione di un mesh 3D. Per ulteriori informazioni, consulta Topologie di tori contorti.

Vantaggi delle prestazioni della TPU v4 rispetto alla 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 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 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 multi-CPU può posizionare i dati necessari per una CPU all'interno del suo nodo NUMA, aumentando il throughput della memoria. Per ulteriori informazioni 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 il binding dei 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 tuo carico di lavoro dipende molto dai carichi della CPU (ad esempio, classificazione delle immagini, carichi di consigli) indipendentemente 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 GB nell'intero chip, che consente una migliore coordinazione tra i due TensorCore on-chip.
• Miglioramento delle prestazioni HBM mediante l'utilizzo degli standard e delle velocità di memoria più recenti.
• Profilo delle prestazioni DMA migliorato con il supporto integrato per gli intervalli ad alte prestazioni con granularità di 512 B.

TensorCores

• Il doppio del numero di MXU e una frequenza di clock più elevata che offre 275 TFLOPS di potenza massima.
• 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ù elevata e supporto della modalità a 8 bit per consentire dimensioni dei batch inferiori e latenza di inferenza migliorata.

Interconnessione tra chip

Sei link di interconnessione per chip per consentire topologie di rete con diametri inferiori.

Altro

• Interfaccia PCIE 3ª gen. x16 con 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. La rete flessibile di TPU v4 ti consente di collegare i chip in una sezione del pod di dimensioni uguali in più modi. Quando crei una sezione di pod di TPU, specifichi la versione di TPU e il numero di risorse TPU richieste. Quando crei una sezione di pod TPU v4, puoi specificarne il tipo e le dimensioni in due modi:AcceleratorType e AccleratorConfig.

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

Uso: AcceleratorType

Utilizza AcceleratorType quando non specifichi una topologia. Per configurare le TPU v4 utilizzando AcceleratorType, utilizza il flag --accelerator-type quando crei la sezione del pod di TPU. Imposta --accelerator-type su una stringa contenente la versione TPU e il numero di TensorCore che vuoi utilizzare. Ad esempio, per creare un slice del pod v4 con 32 TensorCore, devi utilizzare --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.18.0-pod-pjrt

Il numero dopo la versione 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. In genere è necessario per l'ottimizzazione delle prestazioni con sezioni di pod superiori a 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 sulla disposizione fisica dei chip TPU nella sezione del pod.

Specifica una topologia TPU utilizzando una tupla di 3 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. Ad esempio, per creare una sezione del pod v4 con 16 chip, devi impostare --version=v4 e --topology=2x2x4.

Il comando seguente crea una sezione pod di 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.18.0-pod-pjrt

Le topologie in cui 2A=B=C o 2A=2B=C hanno anche varianti di topologia ottimizzate per la comunicazione all'infinito, ad esempio 4×4×8, 8×8×16 e 12×12×24. Queste sono note come topologie torus attorti.

Le seguenti illustrazioni mostrano alcune topologie TPU v4 comuni.

È possibile creare slice di pod più grandi da uno o più "cubi" di chip 4x4x4.

Topologie di tori contorti

Per alcuni tipi di sezioni di toro 3D v4 è possibile utilizzare la topologia del toro contorto. Ad esempio, due cubi v4 possono essere disposti come sezione 4x4x8 o 4x4x8_twisted. Le topologie con cavo intrecciato offrono una larghezza di banda di bisezione molto più elevata. L'aumento della larghezza di banda della bisezione è utile per i carichi di lavoro che utilizzano modelli di comunicazione globali. Le topologie contorte possono migliorare le prestazioni per la maggior parte di modelli, con grandi carichi di lavoro di embedding TPU che 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 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 sul modello FSDP MaxText hanno registrato miglioramenti di 1-2 MFU utilizzando una topologia a spirale.

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 presenta molti vantaggi:

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

Questi vantaggi si traducono in un miglioramento del rendimento per molti modelli di comunicazione a livello mondiale.

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 a toro 4 x 2 con TPU etichettate con le relative coordinate (X,Y) nello slice:

Per chiarezza, gli archi in questo grafo di topologia sono mostrati come archi non diretti. In pratica, ogni bordo è una connessione bidirezionale tra TPU. I bordi tra un lato di questa griglia e il lato opposto sono definiti bordi con a capo, come indicato nel diagramma.

Se ruotiamo questa topologia, otteniamo una topologia di toro contorto 4x2 completamente simmetrica:

L'unica differenza tra questo diagramma e quello precedente è rappresentata dai bordi con rientro Y. 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.

Lo stesso concetto si applica a dimensioni di dimensioni diverse e a un numero diverso di dimensioni. La rete risultante è simmetrica, a condizione che ogni dimensione sia uguale o doppia rispetto alla dimensione più piccola.

Consulta Utilizzare AcceleratorConfig per informazioni dettagliate su come specificare una configurazione di tori attorcigliati durante la creazione di una Cloud TPU.

La tabella seguente mostra le topologie intrecciate supportate e un aumento teorico della larghezza di banda della bisezione rispetto alle topologie non intrecciate.

Topologia a spirale	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

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 modi diversi. Ad esempio, una sezione di pod di TPU con 512 chip (1024 TensorCore) può essere configurata utilizzando le seguenti topologie: 4x4x32, 4x8x16 o 8x8x8. Una sezione di pod di 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 i workload con più tipi di parallelismo (ad es. parallelismo di modelli e dati o partizione spaziale di una simulazione). Questi workload hanno il rendimento migliore se la topologia corrisponde al parallelismo utilizzato. Ad esempio, l'applicazione del parallelismo del modello 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 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 sezione di pod 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 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).

Piccole topologie v4

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 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 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 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 numero 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	8x8x16
custom topology: devi utilizzare i flag --type e --topology	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…