Architettura TPU
Le Tensor Processing Unit (TPU) sono circuiti integrati specifici per le applicazioni (ASIC) progettati da Google per accelerare i carichi di lavoro di machine learning. Cloud TPU è un Google Cloud servizio che rende le TPU disponibili come risorsa scalabile.
Le TPU sono progettate per eseguire rapidamente operazioni con matrici, il che le rende ideali per i carichi di lavoro di machine learning. Puoi eseguire carichi di lavoro di machine learning su TPU utilizzando framework come PyTorch e JAX.
Come funzionano le TPU?
Per capire come funzionano le TPU, è utile comprendere come altri acceleratori affrontano le sfide computazionali dell'addestramento dei modelli di ML.
Come funziona una CPU
Una CPU è un processore per uso generico basato sull'architettura di von Neumann. Ciò significa che una CPU funziona con software e memoria in questo modo:
Il vantaggio principale delle CPU è la flessibilità. Puoi caricare qualsiasi tipo di software su una CPU per molti tipi diversi di applicazioni. Ad esempio, puoi utilizzare una CPU per l'elaborazione di testi su un PC, controllare i motori di un razzo, eseguire transazioni bancarie o classificare immagini con una rete neurale.
Una CPU carica i valori dalla memoria, esegue un calcolo sui valori e memorizza il risultato nella memoria per ogni calcolo. L'accesso alla memoria è lento rispetto alla velocità di calcolo e può limitare il throughput totale delle CPU. Questo è spesso indicato come il collo di bottiglia di von Neumann.
Come funziona una GPU
Per ottenere una velocità effettiva maggiore, le GPU contengono migliaia di unità aritmetico-logiche (ALU) in un singolo processore. Una GPU moderna di solito contiene tra 2500 e 5000 ALU. Il gran numero di processori ti consente di eseguire migliaia di moltiplicazioni e addizioni contemporaneamente.
Questa architettura GPU funziona bene su applicazioni con parallelismo massiccio, come le operazioni con matrici in una rete neurale. Infatti, in un tipico workload di addestramento per il deep learning, una GPU può fornire un throughput di un ordine di grandezza superiore rispetto a una CPU.
Tuttavia, la GPU è ancora un processore generico che deve supportare molte applicazioni e software diversi. Pertanto, le GPU hanno lo stesso problema delle CPU. Per ogni calcolo nelle migliaia di ALU, una GPU deve accedere ai registri o alla memoria condivisa per leggere gli operandi e memorizzare i risultati del calcolo intermedio.
Come funziona una TPU
Google ha progettato le Cloud TPU come processore di matrici specializzato per i carichi di lavoro delle reti neurali. Le TPU non possono eseguire programmi di videoscrittura, controllare motori di razzi o eseguire transazioni bancarie, ma possono gestire operazioni con matrici di grandi dimensioni utilizzate nelle reti neurali a velocità elevate.
L'attività principale delle TPU è l'elaborazione delle matrici, che è una combinazione di operazioni di moltiplicazione e accumulo. Le TPU contengono migliaia di accumulatori di moltiplicazione collegati direttamente tra loro per formare una grande matrice fisica. Questa architettura è chiamata array sistolico. Cloud TPU v3, contengono due array sistolici di 128 x 128 ALU, su un singolo processore.
L'host TPU trasmette i dati in una coda infeed. La TPU carica i dati dalla coda di input e li archivia nella memoria HBM. Al termine del calcolo, la TPU carica i risultati nella coda di output. L'host TPU legge quindi i risultati dalla coda di outfeed e li memorizza nella memoria dell'host.
Per eseguire le operazioni sulle matrici, la TPU carica i parametri dalla memoria HBM nell'unità di moltiplicazione a matrice (MXU).
Quindi, la TPU carica i dati dalla memoria HBM. Man mano che ogni moltiplicazione viene eseguita, il risultato viene passato al successivo moltiplicatore-accumulatore. L'output è la somma di tutti i risultati della moltiplicazione tra i dati e i parametri. Non è necessario l'accesso alla memoria durante il processo di moltiplicazione delle matrici.
Di conseguenza, le TPU possono raggiungere un throughput di calcolo elevato sui calcoli della rete neurale.
Architettura del sistema TPU
Le sezioni seguenti descrivono i concetti chiave di un sistema TPU. Per ulteriori informazioni sui termini comuni del machine learning, consulta il glossario del machine learning.
Se non hai mai utilizzato Cloud TPU, consulta la home page della documentazione di TPU.
Chip TPU
Un chip TPU contiene uno o più TensorCore. Il numero di TensorCore dipende dalla versione del chip TPU. Ogni Tensor Core è costituito da una o più unità di moltiplicazione a matrice (MXU), un'unità vettoriale e un'unità scalare. Per ulteriori informazioni sui Tensor Core, consulta Un supercomputer specifico per il dominio per l'addestramento di reti neurali profonde.
Una MXU è composta da accumulatori di moltiplicazione 256 x 256 (TPU v6e) o 128 x 128 (versioni di TPU precedenti alla v6e) in un array sistolico. Le MXU forniscono la maggior parte della potenza di calcolo in un Tensor Core. Ogni MXU è in grado di eseguire 16.000 operazioni di moltiplicazione e accumulo per ciclo. Tutte le moltiplicazioni utilizzano input bfloat16, ma tutti gli accumuli vengono eseguiti nel formato numerico FP32.
L'unità vettoriale viene utilizzata per calcoli generali come attivazioni e softmax. L'unità scalare viene utilizzata per il flusso di controllo, il calcolo degli indirizzi di memoria e altre operazioni di manutenzione.
pod di TPU
Un pod di TPU è un insieme contiguo di TPU raggruppate in una rete specializzata. Il numero di chip TPU in un pod di TPU dipende dalla versione di TPU.
Sezione
Una sezione è un insieme di chip tutti situati all'interno dello stesso pod di TPU collegati da interconnessioni interchip (ICI) ad alta velocità. Le sezioni sono descritte in termini di chip o TensorCore, a seconda della versione della TPU.
Forma del chip e topologia del chip si riferiscono anche alle forme delle sezioni.
Multislice e single slice
Multislice è un gruppo di sezioni che estende la connettività TPU oltre le connessioni inter-chip interconnect (ICI) e sfrutta la rete di data center (DCN) per trasmettere dati oltre una sezione. I dati all'interno di ogni segmento vengono comunque trasmessi da ICI. Utilizzando questa connettività ibrida, Multislice consente il parallelismo tra le sezioni e ti permette di utilizzare un numero maggiore di core TPU per un singolo job rispetto a quelli che può ospitare una singola sezione.
Le TPU possono essere utilizzate per eseguire un job su una singola sezione o su più sezioni. Per ulteriori dettagli, consulta la presentazione di Multislice.
Cubo TPU
Una topologia 4x4x4 di chip TPU interconnessi. Questo vale solo per le topologie 3D (a partire dalla TPU v4).
SparseCore
SparseCore sono processori di dataflow che accelerano i modelli basati sugli incorporamenti trovati nei modelli di suggerimenti. v5p include quattro SparseCore per chip e v6e include due SparseCore per chip.
Resilienza dell'ICI di Cloud TPU
La resilienza ICI contribuisce a migliorare la tolleranza agli errori dei collegamenti ottici e degli interruttori di circuiti ottici (OCS) che collegano le TPU tra i cubi. (le connessioni ICI all'interno di un cubo utilizzano collegamenti in rame che non sono interessati). La resilienza ICI consente di instradare le connessioni ICI intorno ai guasti OCS e ICI ottici. Di conseguenza, migliora la disponibilità di pianificazione delle sezioni TPU, con il compromesso di un peggioramento temporaneo delle prestazioni dell'ICI.
Per Cloud TPU v4 e v5p, la resilienza ICI è abilitata per impostazione predefinita per le sezioni di un cubo o più grandi, ad esempio:
- v5p-128 quando specifichi il tipo di acceleratore
- 4x4x4 quando specifichi la configurazione dell'acceleratore
Versioni di TPU
L'architettura esatta di un chip TPU dipende dalla versione della TPU che utilizzi. Ogni versione di TPU supporta anche dimensioni e configurazioni diverse delle sezioni. Per maggiori informazioni sull'architettura del sistema e sulle configurazioni supportate, consulta le seguenti pagine:
Architettura cloud TPU
Google Cloud rende disponibili le TPU come risorse di calcolo tramite le VM TPU. Puoi utilizzare direttamente le VM TPU per i tuoi carichi di lavoro oppure utilizzarle tramite Google Kubernetes Engine o Vertex AI. Le sezioni seguenti descrivono i componenti chiave dell'architettura cloud TPU.
Architettura della VM TPU
L'architettura VM TPU ti consente di connetterti direttamente alla VM collegata fisicamente al dispositivo TPU utilizzando SSH. Una VM TPU, nota anche come worker, è una macchina virtuale che esegue Linux e ha accesso alle TPU sottostanti. Hai accesso root alla VM, quindi puoi eseguire codice arbitrario. Puoi accedere ai log di debug e ai messaggi di errore del compilatore e del runtime.
Host singolo, multi-host e sub-host
Un host TPU è una VM in esecuzione su un computer fisico connesso all'hardware TPU. I workload TPU possono utilizzare uno o più host.
Un carico di lavoro su un singolo host è limitato a una VM TPU. Un workload multi-host distribuisce l'addestramento su più VM TPU. Un carico di lavoro sub-host non utilizza tutti i chip di una VM TPU.
Architettura nodo TPU (deprecata)
L'architettura del nodo TPU è costituita da una VM utente che comunica con l'host TPU tramite gRPC. Quando utilizzi questa architettura, non puoi accedere direttamente all'host TPU, il che rende difficile il debug dell'addestramento e degli errori della TPU.
Passaggio dall'architettura del nodo TPU alla VM TPU
Se hai TPU che utilizzano l'architettura dei nodi TPU, segui questi passaggi per identificarle, eliminarle e riprovisionarle come VM TPU.
Vai alla pagina TPU:
Individua la TPU e la relativa architettura nella sezione Architettura. Se l'architettura è"VM TPU", non devi fare nulla. Se l'architettura è "Nodo TPU", devi eliminare e riprovisionare la TPU.
Elimina e riprovisiona la TPU.
Consulta la sezione Gestione delle TPU per le istruzioni su come eliminare e riprovisionare le TPU.