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 matriciale (MXU), un vettoriale e un'unità scalare. La tabella seguente mostra le specifiche principali per un pod di 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 min/media/max misurata	90/170/192 O
Dimensione pod di TPU	4096 chip
Topologia Interconnect	Mesh 3D
Picco di computing per pod	1.1 exaflops (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 può essere configurato come toro 3D sulle 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 tori intrecciate.

Vantaggi delle prestazioni di TPU v4 rispetto a v3

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

Sistema di memoria

Non Uniform Memory Access (NUMA) è un'architettura di memoria del computer per le macchine con più CPU. Ogni CPU ha accesso diretto a un blocco di dati la memoria. 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 di NUMA località vincolando il tuo script di addestramento a NUMA nodo 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 a NUMA Node 0.

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

Attiva l'associazione 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 di runtime TPU senza un suffisso -pod (ad esempio, tpu-vm-tf-2.10.0-v4).

Altre differenze nel 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.
• Prestazioni HBM migliorate grazie agli standard e alle velocità di memoria più recenti.
• Profilo di rendimento DMA migliorato con supporto integrato per prestazioni elevate con una granularità di 512 miliardi.

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 del datastore per la memoria comune (Cmem).
• La larghezza di banda per il caricamento del peso MXU più veloce e il supporto della modalità a 8 bit per consentire una 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 gen3 x16 all'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 di TPU, specifica 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 si escludono a vicenda; puoi utilizzarne solo uno durante la creazione di una TPU.

Uso: AcceleratorType

Usa AcceleratorType quando non specifichi una topologia. Per configurare le TPU v4 usando AcceleratorType, usa il flag --accelerator-type durante la creazione Sezione di pod di 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 TensorCores, 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. In una TPU v4 ci sono due TensorCore, quindi il numero di chip TPU sarebbe 512/2 = 256.

Uso: AcceleratorConfig

Utilizza AcceleratorConfig per 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 --version e --topology flag. 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 di pod di TPU v4 con 128 chip TPU disposti in 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 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 grandi dimensioni, con carichi di lavoro di incorporamento TPU di grandi dimensioni trarne i maggiori vantaggi.

Per i carichi di lavoro che utilizzano il parallelismo dei dati come unica strategia di parallelismo, potrebbe funzionare leggermente meglio. Per gli LLM, le prestazioni sono basate su una può variare in base al tipo di parallelismo (DP, MP e così via). Migliore è addestrare l'LLM con e senza una topologia contorta per determinare che offre le migliori prestazioni per il tuo modello. Alcuni esperimenti Il modello FSDP MaxText ha abbiamo visto 1-2 miglioramenti degli MFU usando una topologia contorta.

Il vantaggio principale delle topologie contorte è che trasformano una torus (ad esempio, 4 × 4 × 8) in una topologia simmetrica strettamente correlata. La topologia simmetrica offre molti vantaggi:

• Bilanciamento del carico migliorato
• Maggiore larghezza di banda di bisezione
• Route di 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 intrecciati sulle sezioni in cui le dimensioni è uguale o doppia alla dimensione più piccola. Per ad esempio 4 x 4 x 8, 4 x 8 x 8 o 12 x 12 x 24.

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. Nel In pratica, ogni perimetro è 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 con bordi a tutto tondo. Invece di connetterti a un'altra TPU con la stessa 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, purché ogni dimensione sia uguale o doppia 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_torto	Circa il 70%
8x8x16_girata
12×12×24_intrecciato
4×8×8_torto	Circa il 40%
8×16×16_intrecciato

Nota: la topologia contorta non verrà utilizzata automaticamente a meno che il suffisso _twisted è specificato con il flag --topology. Ad esempio, per specificare un modello di , devi specificare il flag topology come: 4x4x8_twisted.

Varianti 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 di TPU con 512 chip (1024 TensorCore) può essere configurato utilizzando le seguenti topologie: 4x4x32, 4x8x16 o 8x8x8. Una sezione di pod di TPU con 2048 chip (4096 TensorCores) offre ancora più opzioni di topologia: 4 x 8 x 64, 4 x 16 x 32 e 8 x 16 x 16.

La topologia predefinita associata a un determinato conteggio di chip è quella più simile a un cubo. Questa forma è probabilmente la scelta migliore per l'addestramento con ML parallelo tra dati. Altro le topologie possono essere utili per carichi di lavoro con più tipi di parallelismo (ad esempio, il parallelismo di modello e dati o partizionamento 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 se dimensioni di parallelismo mappate alle dimensioni della topologia TPU. Di solito sono e modellare modelli linguistici di grandi dimensioni (LLM) paralleli. 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 non ottimale a questa topologia, poiché non è allineata ad alcuna dimensione della topologia TPU, ma sarebbe ottimale con una topologia 4x16x32 con lo stesso numero di chip.

Le configurazioni TPU v4 sono costituite da due gruppi, quelli con topologie più piccole 64 chip (piccole topologie) 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 4x4x4. Puoi creare queste piccole topologie v4 utilizzando Nome basato su TensorCore (ad esempio, v4-32) o relativa 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 grandi

Le sezioni di pod TPU v4 sono disponibili in incrementi di 64 chip, con forme multipli di 4 su tutte e tre le dimensioni. Le dimensioni devono essere in in ordine crescente. La seguente tabella mostra diversi esempi. Alcuni di che queste topologie sono "personalizzate" che possono essere create solo utilizzando --type e --topology perché esistono più modi per organizzare le patatine.

Nome (in base al 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
…	…	…