Architettura TPU
Le TPU (Tensor Processing Unit) sono circuiti integrati per applicazioni specifiche (ASIC) progettati da Google per accelerare i carichi di lavoro di machine learning. Cloud TPU è un Google Cloud servizio che rende disponibili le TPU come risorsa scalabile.
Le TPU sono progettate per eseguire rapidamente operazioni sulle 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 capire in che modo altri acceleratori risolvono le sfide computazionali dell'addestramento dei modelli di ML.
Come funziona una CPU
Una CPU è un processore generico basato sull'architettura di Von Neumann. Ciò significa che una CPU lavora con software e memoria in questo modo:
Il vantaggio principale delle CPU è la loro 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, per controllare i motori dei razzi, eseguire transazioni bancarie o classificare le immagini con una rete neurale.
Una CPU carica i valori dalla memoria, esegue un calcolo sui valori e memorizza nuovamente 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 collo di bottiglia di von Neumann.
Come funziona una GPU
Per ottenere un throughput più elevato, le GPU contengono migliaia di unità aritmetiche logiche (ALU) in un singolo processore. Una GPU moderna di solito contiene tra 2500 e 5000 ALU. Il grande numero di processori ti consente di eseguire contemporaneamente migliaia di moltiplicazioni e addizioni.
Questa architettura GPU funziona bene in applicazioni con parallelismo massiccio, come le operazioni matriciali in una rete neurale. Infatti, in un tipico carico di lavoro 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 molti diversi software e applicazioni. 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 intermedi del calcolo.
Come funziona una TPU
Google ha progettato le Cloud TPU come processori a matrice specializzati per i carichi di lavoro delle reti neurali. Le TPU non possono eseguire elaboratori di testi, controllare motori di razzi o eseguire transazioni bancarie, ma possono gestire a velocità elevate operazioni di matrice di grandi dimensioni utilizzate nelle reti neurali.
Il compito principale delle TPU è l'elaborazione delle matrici, ovvero 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 è chiamata architettura array systolic. Cloud TPU v3, contengono due array systolici di ALU 128 x 128 su un singolo processore.
L'host TPU esegue lo streaming dei dati in una coda infeed. La TPU carica i dati dalla coda infeed e li archivia nella memoria HBM. Al termine del calcolo, la TPU carica i risultati nella coda di outfeed. L'host TPU legge quindi i risultati della 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).
Successivamente, la TPU carica i dati dalla memoria HBM. Quando viene eseguita ogni moltiplicazione, il risultato viene passato all'accumulatore di moltiplicazioni successivo. L'output è la somma di tutti i risultati della moltiplicazione tra i dati e i parametri. Non è richiesto alcun accesso alla memoria durante il processo di moltiplicazione delle matrici.
Di conseguenza, le TPU possono raggiungere una velocità di calcolo elevata per i calcoli delle reti neurali.
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 TensorCore è 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 l'addestramento di reti neurali profonde.
Un'unità MXU è composta da accumulatori di moltiplicazione 256 x 256 (TPU v6e) o 128 x 128 (versioni TPU precedenti alla v6e) in un array sistolica. Le MXU forniscono la maggior parte della potenza di calcolo in un TensorCore. Ogni MXU è in grado di eseguire 16.000 operazioni di moltiplicazione e accumulo per ciclo. Tutte le moltiplicazioni accettano input bfloat16, ma tutte le accumulazioni vengono eseguite in formato numerico FP32.
L'unità di vettore 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 TPU dipende dalla versione della TPU.
Sezione
Una sezione è una raccolta di chip tutti all'interno dello stesso pod TPU connessi da interconnessioni interchip ad alta velocità (ICI). 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 singolo slice
Multislice è un gruppo di sezioni che estende la connettività delle TPU oltre le connessioni inter-chip (ICI) e sfrutta la rete del data center (DCN) per trasmettere i dati oltre una sezione. I dati all'interno di ogni segmento vengono comunque trasmessi dall'ICI. Grazie a questa connettività ibrida, Multislice consente il parallelismo tra le slice e ti permette di utilizzare un numero maggiore di core TPU per un singolo job rispetto a quanto può supportare una singola slice.
Le TPU possono essere utilizzate per eseguire un job su un singolo slice o su più slice. Per ulteriori dettagli, consulta la Introduzione a Multislice.
Cube TPU
Una topologia 4x4x4 di chip TPU interconnessi. Questo vale solo per le topiarie 3D (a partire da TPU v4).
SparseCore
SparseCore sono processori di dataflow che accelerano i modelli basati sugli incorporamenti trovati nei modelli di consigli. La versione v5p include quattro SparseCore per chip e la versione v6e include due SparseCore per chip.
Resilienza dell'ICI di Cloud TPU
La resilienza dell'ICI contribuisce a migliorare la tolleranza agli errori dei link ottici e degli interruttori di circuiti ottici (OCS) che collegano le TPU tra i cubetti. (le connessioni ICI all'interno di un cubo utilizzano cavi in rame non interessati). La resilienza dell'ICI consente di instradare le connessioni ICI attorno ai guasti dell'OCS e dell'ICI ottico. Di conseguenza, migliora la disponibilità di pianificazione delle sezioni TPU, con il compromesso del degrado temporaneo delle prestazioni dell'ICI.
Per Cloud TPU v4 e v5p, la resilienza ICI è abilitata per impostazione predefinita per le sezioni di almeno un cubo, ad esempio:
- v5p-128 quando si specifica il tipo di acceleratore
- 4x4x4 quando specifichi la configurazione dell'acceleratore
Versioni TPU
L'architettura esatta di un chip TPU dipende dalla versione della TPU utilizzata. Ogni versione di TPU supporta anche dimensioni e configurazioni diverse delle sezioni. Per maggiori informazioni sull'architettura di sistema e sulle configurazioni supportate, consulta le seguenti pagine:
Architettura cloud TPU
Google Cloud rende disponibili le TPU come risorse di calcolo tramite VM TPU. Puoi utilizzare direttamente le VM TPU per i tuoi carichi di lavoro o tramite Google Kubernetes Engine o Vertex AI. Le sezioni seguenti descrivono i componenti chiave dell'architettura cloud TPU.
Architettura VM TPU
L'architettura VM TPU ti consente di connetterti direttamente alla VM collegata fisicamente al dispositivo TPU tramite SSH. Una VM TPU, nota anche come worker, è una macchina virtuale che esegue Linux e ha accesso alle TPU sottostanti. Hai accesso amministrativo 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 host secondario
Un host TPU è una VM in esecuzione su un computer fisico collegato all'hardware TPU. I carichi di lavoro TPU possono utilizzare uno o più host.
Un carico di lavoro a singolo host è limitato a una VM TPU. Un carico di lavoro multi-host distribuisce l'addestramento su più VM TPU. Un carico di lavoro del sottohost non utilizza tutti i chip di una VM TPU.
Architettura del 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 TPU.
Passaggio dall'architettura del nodo TPU a quella della VM TPU
Se hai TPU che utilizzano l'architettura dei nodi TPU, segui questi passaggi per identificarle, eliminarle e eseguirne il nuovo provisioning 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 eseguire nuovamente il provisioning della TPU.
Elimina e esegui nuovamente il provisioning della TPU.
Consulta la sezione Gestione delle TPU per le istruzioni sull'eliminazione e sul riprovisioning delle TPU.