Introduzione a Cloud TPU

Le TPU (Tensor Processing Unit) sono ASIC (Application-Specific Integrated Circuit) appositamente sviluppati da Google per accelerare i carichi di lavoro di machine learning. Per informazioni più dettagliate sull'hardware TPU, consulta Architettura di sistema. Cloud TPU è un servizio web che rende le TPU disponibili come risorse di calcolo scalabili su Google Cloud.

Le TPU addestrano i tuoi modelli in modo più efficiente utilizzando hardware progettato per eseguire operazioni con matrici di grandi dimensioni, spesso presenti negli algoritmi di machine learning. Le TPU hanno memoria HBM (High-bandwidth Memory) su chip che ti consente di utilizzare modelli e dimensioni batch più grandi. Le TPU possono essere connesse in gruppi chiamati pod, in grado di fare lo scale up dei carichi di lavoro con modifiche minime o nulle al codice.

Come funziona?

Per capire come funzionano le TPU, è utile capire in che modo altri acceleratori affrontano le sfide di calcolo dell'addestramento dei modelli ML.

Come funziona una CPU

Una CPU è un processore per uso generico basato sull'architettura von Neumann. Ciò significa che una CPU funziona con software e memoria in questo modo:

Illustrazione del funzionamento di una CPU

Il maggiore vantaggio delle CPU è la loro flessibilità. Puoi caricare qualsiasi tipo di software su una CPU per tanti tipi diversi di applicazioni. Ad esempio, puoi utilizzare una CPU per l'elaborazione di testi su un PC, per il controllo dei motori a razzo, l'esecuzione di transazioni bancarie o la classificazione delle immagini con una rete neurale.

Una CPU carica i valori dalla memoria, esegue un calcolo sui valori e archivia i risultati in memoria per ogni calcolo. L'accesso alla memoria è lento rispetto alla velocità di calcolo e può limitare la velocità effettiva totale delle CPU. Questo viene spesso definito collo di bottiglia di von Neumann.

Come funziona una GPU

Per ottenere una velocità effettiva superiore, le GPU contengono migliaia di unità logiche aritmetiche (ALU) in un singolo processore. Una GPU moderna contiene solitamente tra 2500 e 5000 ALU. L'elevato numero di processori consente di eseguire migliaia di moltiplicazioni e aggiunte contemporaneamente.

Illustrazione del funzionamento di una GPU

Questa architettura GPU funziona bene su applicazioni con parallelismo elevato, come le operazioni con matrice in una rete neurale. Infatti, su un carico di lavoro di addestramento tipico per il deep learning, una GPU può offrire una velocità effettiva superiore di un ordine di grandezza rispetto a una CPU.

Tuttavia, la GPU è ancora un processore per uso generico che deve supportare molte applicazioni e software diversi. Le GPU presentano lo stesso problema delle CPU. Per ogni calcolo delle migliaia di ALU, una GPU deve accedere ai registri o alla memoria condivisa per leggere gli operandi e archiviare i risultati di calcolo intermedi.

Come funziona una TPU

Le Cloud TPU sono state progettate da Google come processori matriciali specializzati per i carichi di lavoro di rete neurale. Le TPU non possono eseguire elaboratori di testi, controllare motori a razzo o eseguire transazioni bancarie, ma sono in grado di gestire le enormi operazioni della matrice utilizzate nelle reti neurali ad alta velocità.

L'attività principale delle TPU è l'elaborazione matriciale, ovvero una combinazione di operazioni di moltiplicazione e accumulo. Le TPU contengono migliaia di moltiplicatori direttamente collegati tra loro per formare una grande matrice fisica. Questa è chiamata architettura di array sistolici. Cloud TPU v3 contiene due array sistolici di 128 x 128 ALU, su un singolo processore.

L'host TPU trasmette i flussi di dati in una coda di feed. La TPU carica i dati dalla coda di inserimento 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 dalla coda di outfeed e li archivia nella memoria dell'host.

Per eseguire le operazioni sulla matrice, la TPU carica i parametri dalla memoria HBM nell'unità di moltiplicazione matriciale (MXU).

Illustrazione di come una TPU carica i parametri dalla memoria

Quindi, la TPU carica i dati dalla memoria HBM. Quando viene eseguita ogni moltiplicazione, il risultato viene passato all'accumulatore di moltiplicazione successivo. L'output è la somma di tutti i risultati di moltiplicazione tra i dati e i parametri. Durante il processo di moltiplicazione della matrice non è richiesto l'accesso alla memoria.

Illustrazione di come una TPU carica i dati dalla memoria

Di conseguenza, le TPU possono raggiungere una velocità effettiva di calcolo elevata nei calcoli di rete neurale.

Compilatore XLA

Il codice eseguito sulle TPU deve essere compilato dal compilatore algebra lineare (XLA) dell'acceleratore. XLA è un compilatore just-in-time che prende il grafico emesso da un'applicazione di framework ML e compila i componenti di algebra lineare, perdita e gradiente del grafico nel codice macchina TPU. Il resto del programma viene eseguito sulla macchina host della TPU. Il compilatore XLA fa parte del runtime TPU che viene eseguito su una macchina host TPU.

Quando utilizzare le TPU

Le Cloud TPU sono ottimizzate per carichi di lavoro specifici. In alcune situazioni, potrebbe essere opportuno utilizzare GPU o CPU su istanze di Compute Engine per eseguire carichi di lavoro di machine learning. In generale, puoi decidere quale hardware è più adatto al tuo carico di lavoro in base alle seguenti linee guida:

CPU

Prototipazione rapida che richiede la massima flessibilità
Modelli semplici che non richiedono molto tempo per l'addestramento
Modelli piccoli con batch di piccole dimensioni ed efficaci
Modelli che contengono molte operazioni TensorFlow personalizzate scritte in C++
Modelli limitati dall'I/O disponibile o dalla larghezza di banda di rete del sistema host

GPU

I modelli con un numero significativo di operazioni TensorFlow/PyTorch/JAX personalizzate che devono essere eseguite almeno parzialmente sulle CPU
Modelli con operazioni TensorFlow non disponibili su Cloud TPU (consulta l'elenco delle operazioni TensorFlow disponibili)
Modelli medio-grandi con batch di dimensioni effettive più grandi

TPU

Modelli dominati da calcoli matriciali
Modelli senza operazioni TensorFlow/PyTorch/JAX personalizzate all'interno del loop di addestramento
Modelli con addestramento per settimane o mesi
Modelli di grandi dimensioni con batch di grandi dimensioni efficaci

Le Cloud TPU non sono adatte ai seguenti carichi di lavoro:

Programmi di algebra lineare che richiedono frequenti diramazioni o contengono molte operazioni algebriche per elemento
Carichi di lavoro che accedono alla memoria in modo sparso
Carichi di lavoro che richiedono operazioni aritmetiche ad alta precisione
Carichi di lavoro di rete neurale che contengono operazioni personalizzate nel loop di addestramento principale.

Best practice per lo sviluppo di modelli

Un programma il cui calcolo è dominato da operazioni non basate su matrice, come l'aggiunta, la rimodellazione o la concatena, probabilmente non otterrà un elevato utilizzo di MXU. Di seguito sono riportate alcune linee guida per aiutarti a scegliere e creare modelli adatti a Cloud TPU.

Layout

Il compilatore XLA esegue trasformazioni del codice, tra cui il tiling di una matrice moltiplicata in blocchi più piccoli, per eseguire in modo efficiente i calcoli sull'unità matrice (MXU). La struttura dell'hardware MXU, un array sistolico 128 x 128 e il design del sottosistema di memoria di TPU, che preferisce dimensioni multiple di 8, vengono utilizzati dal compilatore XLA per l'efficienza del tiling. Di conseguenza, alcuni layout sono più favoriti per la creazione di riquadri, mentre altri richiedono una rimodellazione prima di poter essere affiancati. Le operazioni di rimodellazione sono spesso associate alla memoria su Cloud TPU.

Forme

Il compilatore XLA compila un grafico ML appena in tempo per il primo batch. Se i batch successivi hanno forme diverse, il modello non funziona. (Ricompilare il grafico ogni volta che la forma cambia è troppo lento). Pertanto, qualsiasi modello che abbia tensori con forme dinamiche che cambiano in fase di runtime non è adatto alle TPU.

Spaziatura interna

In un programma Cloud TPU ad alte prestazioni, la densità di calcolo può essere facilmente suddivisa in blocchi 128 x 128. Quando un calcolo della matrice non può occupare un intero MXU, il compilatore riempie i tensori con zeri. La spaziatura interna presenta due svantaggi:

I tensori riempiti di zeri sottoutilizzano il core TPU.
La spaziatura interna aumenta la quantità di spazio di archiviazione su chip richiesta per un tensore e in casi estremi può causare un errore di esaurimento della memoria.

Mentre la spaziatura interna viene eseguita automaticamente dal compilatore XLA quando necessario, è possibile determinare la quantità di spaziatura interna eseguita utilizzando lo strumento op_profile. Puoi evitare la spaziatura interna scegliendo le dimensioni dei tensori più adatte alle TPU.

Dimensioni

La scelta delle dimensioni del tensore adatte è fondamentale per estrarre le massime prestazioni dall'hardware TPU, in particolare dall'MXU. Il compilatore XLA cerca di usare la dimensione batch o la dimensione della funzionalità per utilizzare al massimo l'MXU. Pertanto, uno di questi deve essere un multiplo di 128. In caso contrario, il compilatore ne aggiungerà uno a 128. Idealmente, le dimensioni del batch e le dimensioni delle funzionalità devono essere multipli di 8, il che consente di estrarre prestazioni elevate dal sottosistema di memoria.

Integrazione dei Controlli di servizio VPC

I Controlli di servizio VPC di Cloud TPU consentono di definire i perimetri di sicurezza intorno alle risorse Cloud TPU e controllare lo spostamento dei dati oltre i confini del perimetro. Per saperne di più sui Controlli di servizio VPC, consulta Panoramica dei Controlli di servizio VPC. Per informazioni sulle limitazioni dell'utilizzo di Cloud TPU con i Controlli di servizio VPC, consulta la pagina relativa ai prodotti e alle limitazioni supportati.

Edge TPU

I modelli di machine learning addestrati nel cloud hanno sempre più bisogno di eseguire l'inferenza "a livello perimetrale", ovvero su dispositivi che operano al limite dell'Internet of Things (IoT). Questi dispositivi includono sensori e altri smart device che raccolgono dati in tempo reale, prendono decisioni intelligenti e quindi agiscono o comunicano le loro informazioni ad altri dispositivi o al cloud.

Poiché tali dispositivi devono funzionare con una potenza limitata (inclusa l'alimentazione a batteria), Google ha progettato il coprocessore Edge TPU per accelerare l'inferenza ML sui dispositivi a bassa potenza. Un singolo Edge TPU può eseguire 4 trilioni di operazioni al secondo (4 TOPS), utilizzando solo 2 watt di potenza; in altre parole, ottieni 2 TOPS per watt. Ad esempio, Edge TPU può eseguire modelli di visione artificiale mobili all'avanguardia come MobileNet V2 a quasi 400 frame al secondo in modo efficiente dal punto di vista energetico.

Questo acceleratore ML a bassa potenza migliora Cloud TPU e Cloud IoT per fornire un'infrastruttura end-to-end (cloud-to-edge, hardware e software) che facilita le soluzioni basate su AI.

Edge TPU è disponibile per i tuoi dispositivi di prototipazione e produzione in diversi fattori di forma, tra cui un computer a scheda singola, un system-on-module, una scheda PCIe/M.2 e un modulo a montaggio superficiale. Per ulteriori informazioni su Edge TPU e su tutti i prodotti disponibili, visita coral.ai.

Introduzione a Cloud TPU

Richiesta di assistenza

Contatta l'assistenza di Cloud TPU. Se hai un progetto Google Cloud attivo, preparati a fornire le seguenti informazioni:

L'ID del tuo progetto Google Cloud
Il nome del nodo TPU, se esistente
Altre informazioni che vuoi fornire

Che cosa succede dopo?

Vuoi saperne di più su Cloud TPU? Le seguenti risorse potrebbero esserti utili: