Architettura di sistema
Le Tensor Processing Unit (TPU) sono acceleratori di machine learning progettati da Google. Cloud TPU rende le TPU disponibili come risorsa Google Cloud scalabile. Puoi eseguire carichi di lavoro di machine learning su Cloud TPU utilizzando framework di machine learning come TensorFlow, Pytorch e JAX.
Ogni TensorCore su un chip TPU è costituito da una o più unità di moltiplicazione Matrix, da un'unità vettoriale e da un'unità scalare.
Un MXU è composto da 128 x 128 moltiplicatori/accumulatori in un array sistolico. Le MXU forniscono la maggior parte della potenza di calcolo di un TensorCore. Ogni MXU è in grado di eseguire 16.000 operazioni di moltiplicazione ad accumulo in ogni ciclo. Tutti i moltiplicatori prendono gli input bFlo16, ma tutti gli accumuli vengono eseguiti nel formato numerico FP32.
L'unità vettore viene utilizzata per il calcolo generale, come attivazioni e softmax. L'unità scalare viene utilizzata per il flusso di controllo, il calcolo dell'indirizzo di memoria e altre operazioni di manutenzione.
Versioni TPU
Il layout esatto di una TPU dipende dalla versione TPU che utilizzi. I dettagli architettonici e le caratteristiche prestazionali di TPU v2 e v3 sono disponibili in Un supercomputer specifico per il dominio per l'addestramento di reti neurali profonde.
TPU v4
Ogni chip TPU v4 contiene due TensorCore. Ogni TensorCore ha quattro MXU, un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche della chiave e i relativi valori per un pod TPU v4.
| Specifiche principali | Valori dei pod v4 |
|---|---|
| Picco di calcolo per chip | 275 teraflop (bf16 o int8) |
| Capacità e larghezza di banda HBM2 | 32 GiB, 1200 Gbps |
| Potenza misurata min/media/max | 90/170/192 W |
| Dimensione pod TPU | 4096 chip |
| Topologia Interconnect | Toro 3D |
| Picco di calcolo per pod | 1.1 exaflop (bf16 o int8) |
| Riduci la larghezza di banda per pod | 1,1 PB/s |
| Larghezza di banda della bisezione per pod | 24 TB/s |
Il seguente diagramma illustra un chip TPU v4.
TPU v3
Ogni chip TPU v3 contiene due TensorCore. Ogni TensorCore ha due MXU, un'unità vettoriale e un'unità scalare. La tabella seguente mostra le specifiche della chiave e i relativi valori per un pod TPU v3.
| Specifiche principali | Valori dei pod v3 |
|---|---|
| Picco di calcolo per chip | 123 teraflop (bf16) |
| Capacità e larghezza di banda HBM2 | 32 GiB, 900 Gbps |
| Potenza misurata min/media/max | 123/220/262 W |
| Dimensione pod TPU | 1024 chip |
| Topologia Interconnect | Toro 2D |
| Picco di calcolo per pod | 126 petaflop (bf16) |
| Riduci la larghezza di banda per pod | 340 TB/s |
| Larghezza di banda della bisezione per pod | 6,4 TB/s |
Il seguente diagramma illustra un chip TPU v3.
TPU v2
La configurazione TPU v2 più piccola contiene quattro chip TPU e 16 GiB di HBM. Ogni chip TPU contiene due TensorCore. Ogni TensorCore ha un'unità MXU, un'unità vettoriale e un'unità di scalabilità. Il seguente diagramma illustra un chip TPU v2.
Cloud TPU fornisce le seguenti configurazioni TPU:
- Un singolo dispositivo TPU
- Un pod di TPU: un gruppo di dispositivi TPU connessi da interconnessioni ad alta velocità
- Una sezione di TPU: una suddivisione di un pod di TPU
Vantaggi delle prestazioni di TPU v4 rispetto a v3
Sistema di memoria:
- I chip TPU v4 hanno uno spazio di memoria HBM a 32 GiB unificato per l'intero chip, consentendo un migliore coordinamento tra i due TensorCore sul chip.
- Prestazioni HBM migliorate utilizzando i più recenti standard di memoria e velocità.
- Profilo di prestazioni DMA migliorato con supporto integrato per sequenze di prestazioni elevate con granularità di 512 miliardi.
TensorCore:
- Raddoppia il numero di MXU e la frequenza di orologio più elevata, offrendo un massimo di 275 TFLOPS.
- 2x larghezza di banda di trasposizione e permutazione.
- Modello di accesso alla memoria di archiviazione per Common Memory (Cmem).
- Maggiore velocità di larghezza di banda di caricamento MXU e supporto per la 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'hosting (connessione diretta).
- Modello di sicurezza migliorato.
- Miglioramento dell'efficienza energetica.
Vantaggi delle prestazioni di TPU v3 rispetto a v2
L'aumento dei FLOPS per capacità di memoria e core nelle configurazioni TPU v3 può migliorare le prestazioni dei tuoi modelli nei seguenti modi:
Le configurazioni di TPU v3 offrono vantaggi significativi in termini di prestazioni per core per i modelli legati al calcolo. I modelli legati alla memoria nelle configurazioni TPU v2 potrebbero non ottenere questo stesso miglioramento delle prestazioni se sono legati anche alla memoria nella configurazione della TPU v3.
Nei casi in cui i dati non rientrano nella memoria nelle configurazioni TPU v2, TPU v3 può fornire prestazioni migliori e una riduzione del calcolo dei valori intermedi (riformizzazione).
Le configurazioni di TPU v3 possono eseguire nuovi modelli con dimensioni del batch non adatte alle configurazioni di TPU v2. Ad esempio, TPU v3 potrebbe consentire ResNet più dettagliati e immagini più grandi con RetinaNet.
Anche i modelli quasi in entrata ("infeed") su TPU v2 perché i passaggi di addestramento sono in attesa di input potrebbero essere vincolati all'input con Cloud TPU v3. La guida alle prestazioni della pipeline può aiutarti a risolvere i problemi di feed.
Configurazioni TPU v4
Un pod TPU v4 è composto da 4096 chip interconnessi con link ad alta velocità riconfigurabili. Il networking flessibile di TPU v4 ti consente di collegare i chip in una porzione della stessa dimensione in diversi modi. Le TPU v4 vengono configurate utilizzando AcceleratorConfig.
AcceleratorConfig ti consente di specificare le dimensioni della sezione in termini di chip tramite il campo topology.
La topologia TPU v4 viene specificata utilizzando una tupla a 3 tuffi che descrive il numero, la forma e le interconnessioni tra i chip TPU. Le seguenti illustrazioni mostrano alcune topologie comuni di TPU v4.
Le fette più grandi possono essere create da uno o più "cubetti" 4x4x4.
Le sezioni TPU di un determinato numero di chip possono essere configurate in diversi modi. Ad esempio, una sezione di TPU con AcceleratorType v4-1024 ora può essere configurata come: 4x4x32, 4x8x16 o 8x8x8.
Configurazioni di TPU v2 e v3
Le TPU sono disponibili nelle seguenti configurazioni:
- Una singola scheda TPU
- Un pod TPU
- Una porzione di pod TPU
Scheda TPU singola
Una configurazione TPU a scheda singola è una scheda autonoma con quattro chip TPU (otto TensorCore) senza connessioni di rete ad altre schede TPU. Le TPU a scheda singola non fanno parte della configurazione di un pod di TPU e non occupano una porzione del pod di TPU.
Pod e fette di TPU
In un pod TPU o porzione di pod TPU, i chip TPU vengono connessi tramite un'interconnessione ad alta velocità. Ogni chip TPU comunica direttamente con gli altri chip sul dispositivo TPU. Il software TPU gestisce automaticamente la distribuzione dei dati a ogni core TPU in un pod o una sezione. Le sezioni dei pod sono disponibili con 32, 128, 512, 1024 o 2048 core.
Architetture VM Cloud TPU
Le TPU sono progettate per eseguire rapidamente operazioni con le matrici. Ogni scheda TPU è connessa a una macchina host basata su CPU per eseguire operazioni non eseguibili sulla TPU. Le macchine host sono responsabili del caricamento dei dati da Cloud Storage, della pre-elaborazione e dell'invio dei dati alla TPU.
In un pod di TPU è presente un host TPU per ogni scheda.
La tua interazione con l'host TPU (e la scheda TPU) dipende dall'architettura della VM TPU che stai utilizzando: nodi TPU o VM TPU.
Nodi TPU
L'architettura del nodo TPU è composta da una VM utente che comunica con l'host TPU su gRPC. Quando utilizzi questa architettura, non puoi accedere direttamente all'host TPU, rendendo difficile il debug degli errori di addestramento e TPU.
VM TPU
L'architettura delle VM TPU consente di connetterti tramite SSH alla VM fisicamente connessa al dispositivo TPU. Hai accesso root alla VM e puoi eseguire codice arbitrario. Puoi accedere ai log di debug e ai messaggi di errore del compilatore e di runtime.
Framework supportati
I framework come JAX, PyTorch e TensorFlow accedono alle TPU utilizzando una libreria condivisa denominata libtpu presente in ogni VM TPU. libtpu include il compilatore XLA, il software di runtime TPU e il driver TPU.
TensorFlow
VM TPU
Con le VM TPU, invece di eseguire il codice Python su una VM utente, questo può essere eseguito direttamente sull'host TPU.
Per ulteriori informazioni su TensorFlow e sulla VM Cloud TPU, consulta la Guida dell'utente alle VM Cloud TPU.
Nodi TPU
L'architettura di sistema dei nodi Cloud TPU è stata creata originariamente per TensorFlow. Gli host TPU non sono accessibili all'utente ed eseguono una copia headless di un server TensorFlow. Non eseguono Python o codice utente non rappresentato come un grafico TensorFlow. Il codice utente viene eseguito in una VM utente separata che comunica con gli host TPU sulla rete gRPC.
PyTorch
VM TPU
Con le VM TPU, il codice PyTorch viene eseguito direttamente sugli host TPU.
Per ulteriori informazioni su PyTorch e Cloud TPU, consulta la Guida dell'utente di PyTorch/XLA.
Nodi TPU
PyTorch viene eseguita sull'architettura del nodo Cloud TPU utilizzando una libreria denominata XRT. XRT invia grafici XLA e istruzioni di runtime su gRPC da eseguire sui server TensorFlow. È necessaria una VM utente per ogni host TPU.
JAX
VM TPU
Con le VM TPU, il codice JAX viene eseguito direttamente sugli host TPU.
Per ulteriori informazioni sull'esecuzione di JAX su Cloud TPU, consulta la guida rapida di JAX.
Nodi TPU
JAX sui nodi Cloud TPU viene eseguito in modo simile a PyTorch, in quanto per ogni VM host è necessaria una VM utente separata.
