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 matriciale (MXU), un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche principali per unpod di TPUU v4.

Specifiche principali	Valori del pod v4
Picco computing per chip	275 teraflop (bf16 o int8)
Capacità e larghezza di banda HBM2	32 GiB, 1200 GB/s
Potenza minima/media/max misurata	90/170/192 E
Dimensione pod di TPU	4096 chip
Topologia di interconnessione	Mesh 3D
Picco computing per pod	1.1 exaflops (bf16 o int8)
Riduce 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 dimensioni, creando 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, 4 x 4 x 8, 4 x 8 x 8 o 12 x 12 x 24. In generale, le prestazioni di una configurazione di toro 3D sono migliori di una configurazione di mesh 3D. Per ulteriori informazioni, consulta la sezione Topologie contorte.

Vantaggi delle prestazioni di TPU v4 rispetto a v3

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

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 sua memoria sono chiamate nodo NUMA. NUMA nodi sono connessi ad altri nodi NUMA direttamente adiacenti l'uno all'altro. 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 richiesti da una CPU all'interno del nodo NUMA, aumentando la velocità effettiva della memoria. Per ulteriori informazioni su NUMA, consulta la pagina relativa a Non Uniform Memory Access su Wikipedia.

Puoi sfruttare i vantaggi delle località NUMA associando il tuo script di addestramento a NUMA Node 0.

Per abilitare 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 del tuo script a NUMA Node 0.

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

Abilita l'associazione di nodi NUMA se:

• Se il tuo carico di lavoro dipende in modo significativo dai carichi di lavoro della CPU (ad esempio classificazione delle immagini, carichi di lavoro per suggerimenti) indipendentemente dal framework.
• Se utilizzi una versione del 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 nell'intero chip, consentendo un migliore coordinamento tra i due TensorCore su chip.
• Prestazioni HBM migliorate grazie agli standard e alle velocità di memoria più recenti.
• Profilo di prestazioni DMA migliorato con supporto integrato per prestazioni elevate con granularità di 512 miliardi di dollari.

TensorCores

• 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 2x.
• Modello di accesso alla memoria dell'archivio di carico per la memoria comune (Cmem).
• Larghezza di banda per il caricamento del peso MXU più veloce e supporto della modalità a 8 bit per consentire dimensioni batch inferiori e una migliore latenza di inferenza.

Inter-chip Interconnect

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.
• Miglioramento dell'efficienza energetica.

Configurazioni

Un pod TPU v4 è composto da 4096 chip interconnessi con link ad alta velocità riconfigurabili. Il networking flessibile di TPU v4 consente di connettere i chip in una sezione di pod delle stesse dimensioni 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 tipo e dimensioni in uno dei due modi seguenti: AcceleratorType e AccleratorConfig.

Nota: AcceleratorType e AccleratorConfig si escludono a vicenda; 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, usa il flag --accelerator-type durante la creazione della sezione di 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 una sezione di pod v4 con 32 TensorCore, devi utilizzare --accelerator-type=v4-32.

Il seguente comando 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.16.1-pod-pjrt

Il numero dopo la versione TPU (v4) specifica il numero di TensorCore. In una TPU v4 sono presenti due TensorCore, quindi il numero di chip TPU sarebbe 512/2 = 256.

Uso: AcceleratorConfig

Usa AcceleratorConfig quando vuoi personalizzare la topologia fisica della sezione TPU. Questa operazione è generalmente necessaria 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 TPU che vuoi utilizzare e --topology sulla disposizione fisica dei chip TPU nella sezione di pod.

Puoi specificare una topologia TPU utilizzando AxBxC a 3 tuple, 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 il conteggio dei chip in ognuna delle tre dimensioni. Ad esempio, per creare una sezione di pod v4 con 16 chip, devi impostare --version=v4 e --topology=2x2x4.

Il seguente comando crea una sezione di pod di TPU v4 con 128 chip TPU disposti in un array 4 x 4 x 8:

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

Le topologie in cui 2A=B=C o 2A=2B=C hanno anche varianti di topologia ottimizzate per una comunicazione complessiva, ad esempio 4 × 4 × 8, 8 × 8 × 16 e 12 × 12 × 24. Queste sono note come topologie tori torsi.

Le illustrazioni seguenti 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 torsi

Alcune forme di sezione del toro 3D v4 hanno la possibilità di utilizzare quella che viene nota come topologia toro contorto. Ad esempio, due cubi v4 possono essere disposti come sezione 4x4x8 o 4x4x8_twisted. Le topologie intrecciate offrono una bisezione di larghezza di banda significativamente maggiore. Una maggiore larghezza di banda di bisezione è utile per carichi di lavoro che utilizzano schemi di comunicazione globali. Le topologie intrecciate possono migliorare le prestazioni per la maggior parte dei modelli, con grandi carichi di lavoro di incorporamento per TPU di maggiore vantaggio.

Per i carichi di lavoro che utilizzano il parallelismo dei dati come unica strategia di parallelismo, le topologie contorte potrebbero funzionare leggermente migliori. Per gli LLM, le prestazioni con una topologia intrecciata possono variare in base al tipo di parallelismo (DP, MP e così via). La best practice prevede l'addestramento dell'LLM con e senza una topologia contorta per determinare quale offre le migliori prestazioni per il tuo modello. Alcuni esperimenti sul modello FSDP MaxText hanno rilevato 1-2 miglioramenti MFU utilizzando una topologia contorta.

Il vantaggio principale delle topologie intrecciate è che trasforma una topologia torus asimmetrica (ad esempio, 4 × 4 × 8) in una topologia simmetrica strettamente correlata. La topologia simmetrica presenta numerosi vantaggi:

• Bilanciamento del carico migliorato
• Maggiore larghezza di banda di bisezione
• Route di pacchetti più brevi

Questi vantaggi si traducono in un miglioramento delle prestazioni per molti modelli di comunicazione globali.

Il software TPU supporta tori intrecciati sulle sezioni in cui la dimensione di ogni dimensione è uguale o doppia di quella della dimensione più piccola. Ad esempio, 4 x 4 x 8, 4 x 8 × 8 o 12 x 12 x 24.

Ad esempio, considera questa topologia toro 4 × 2 con TPU etichettate con le relative coordinate (X,Y) nella sezione:

I bordi in questo grafico della topologia sono mostrati come bordi non orientati per maggiore chiarezza. In pratica, ogni perimetro è una connessione bidirezionale tra TPU. I bordi tra un lato di questa griglia e l'altro vengono definiti bordi avvolgenti, come indicato nel diagramma.

Torcando questa topologia, si ottiene una topologia toroidale 4 × 2 completamente simmetrica:

Tutto è cambiato tra questo diagramma e quello precedente sono i bordi arrotondati a 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.

La stessa idea si applica a dimensioni di dimensioni diverse e a numeri diversi. La rete risultante è simmetrica, a condizione che ogni dimensione sia uguale o doppia della dimensione della dimensione più piccola.

Per informazioni dettagliate su come specificare una configurazione tori contorta durante la creazione di una Cloud TPU, consulta la pagina relativa all'utilizzo di AcceleratorConfig.

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 della larghezza di banda rispetto a un toro non distorto
4 × 4 × 8_torto	~70%
8x8x16_twisted
12 × 12 × 24_attorcigliato
4×8×8_twisted	Circa il 40%
8 × 16 × 16_torto

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

Varianti della 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 con parallelo dei dati. Altre topologie possono essere utili per carichi di lavoro con più tipi di parallelismo (ad esempio, parallelismo di modelli e dati o partizionamento spaziale di una simulazione). Questi carichi di lavoro funzionano meglio se la topologia è abbinata al parallelismo utilizzato. Ad esempio, il posizionamento 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 4 x 16 x 16.

I modelli con più dimensioni di parallelismo funzionano meglio con le dimensioni di parallelismo mappate a dimensioni di topologia TPU. Di solito si tratta di modelli linguistici di grandi dimensioni (LLM) paralleli a dati e modelli. Ad esempio, per una sezione di pod TPU v4 con topologia 8 x 16 x 16, 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 fisiche della topologia TPU). Un parallelismo del modello a 4 vie non sarebbe ottimale con questa topologia, poiché non è allineato con nessuna delle dimensioni della topologia TPU, ma sarebbe ottimale con una topologia 4 x 16 x 32 sullo stesso numero di chip.

Le configurazioni TPU v4 sono composte da due gruppi, quelli con topologie più piccole di 64 chip (topologie piccole) e quelli con topologie superiori a 64 chip (topologie grandi).

Topologie v4 piccole

Cloud TPU supporta le seguenti sezioni di pod TPU v4 più piccole di 64 chip, un cubo 4 x 4 x 4. Puoi creare queste piccole topologie v4 utilizzando il relativo nome basato su TensorCore (ad esempio v4-32) o la relativa topologia (ad esempio, 2x2x4):

Nome (basato sul 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 grandi

Le sezioni di pod TPU v4 sono disponibili in incrementi di 64 chip, con forme multiple di 4 su tutte e tre le dimensioni. Le dimensioni devono essere in ordine crescente. Nella tabella seguente sono riportati diversi esempi. Alcune di queste topologie sono topologie "personalizzate" che possono essere create solo utilizzando i flag --type e --topology perché esiste più di un modo per organizzare i chip.

Nome (basato sul conteggio di TensorCore)	Numero di chip	Topologia
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
N/A: è necessario utilizzare i flag --type e --topology	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
N/A: è necessario utilizzare i flag --type e --topology	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…