Introduzione a Cloud TPU

Le TPU (Tensor Processing Unit) sono ASIC (Application-Specific Integrated Circuit) sviluppati appositamente 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 su matrici di grandi dimensioni, spesso presenti negli algoritmi di machine learning. Le TPU dispongono di memoria on-chip a larghezza di banda elevata (HBM) che consente di utilizzare modelli e dimensioni batch più grandi. Le TPU possono essere collegate in gruppi chiamati pod che fanno lo scale up dei carichi di lavoro con poche o nessuna modifica al codice.

Come funziona?

Per capire come funzionano le TPU, è utile comprendere in che modo altri acceleratori affrontano le sfide computazionali legate all'addestramento dei modelli 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 come questi:

Illustrazione di come funziona una CPU

Il principale vantaggio 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, 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 il risultato 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 è spesso definito come il 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 di solito contiene tra le 2500 e le 5000 ALU. L'elevato numero di processori significa che puoi eseguire migliaia di moltiplicazioni e aggiunte contemporaneamente.

Illustrazione di come funziona una GPU

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

Tuttavia, la GPU è ancora un processore per uso generico che deve supportare molte applicazioni e software diversi. Di conseguenza, 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 archiviare i risultati dei calcoli intermedi.

Come funziona una TPU

Google ha progettato le Cloud TPU come processori a matrice specializzato per i carichi di lavoro delle reti neurali. Le TPU non possono eseguire elaboratori di testi, controllare motori a razzo o eseguire transazioni bancarie, ma sono in grado di gestire a velocità elevate le operazioni matriciali utilizzate nelle reti neurali.

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

L'host TPU trasmette i flussi di dati in una coda in-feed. 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 uscita. L'host della TPU legge i risultati dalla coda di uscita e li archivia nella memoria dell'host.

Per eseguire le operazioni della matrice, la TPU carica i parametri dalla memoria HBM nella Matrix Multiplication Unit (MXU).

Illustrazione di come una TPU carica i parametri dalla memoria

Quindi, la TPU carica i dati dalla memoria HBM. Man mano che si esegue ogni moltiplicazione, il risultato viene passato all'accumulatore successivo. L'output è la somma di tutti i risultati di moltiplicazione tra dati e parametri. Non è necessario accedere alla memoria durante il processo di moltiplicazione matriciale.

Illustrazione di come una TPU carica i dati dalla memoria

Di conseguenza, le TPU possono raggiungere una velocità effettiva di calcolo elevata nei calcoli delle reti neurali.

Compilatore XLA

Il codice eseguito sulle TPU deve essere compilato dal compilatore accelerator linear algebra (XLA). XLA è un compilatore just-in-time che prende il grafico emesso da un'applicazione del 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 TPU. Il compilatore XLA fa parte dell'immagine VM TPU in esecuzione su una macchina host TPU.

Quando utilizzare le TPU

Le Cloud TPU sono ottimizzate per carichi di lavoro specifici. In alcune situazioni, potrebbe essere utile utilizzare GPU o CPU sulle istanze di Compute Engine per eseguire i 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 dimensioni ridotte ed effettive
  • Modelli che contengono molte operazioni TensorFlow personalizzate scritte in C++
  • Modelli che sono 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
  • Modelli con operazioni TensorFlow non disponibili su Cloud TPU (consulta l'elenco delle operazioni TensorFlow disponibili)
  • Modelli medio-grandi con dimensioni dei batch effettive più grandi

TPU

  • Modelli dominati dai calcoli matriciali
  • I modelli senza operazioni TensorFlow/PyTorch/JAX personalizzate all'interno del loop di addestramento principale
  • Modelli che si addestrano per settimane o mesi
  • Modelli di grandi dimensioni con grandi dimensioni dei batch effettive

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

  • Programmi di algebra lineare che richiedono frequenti ramificazioni o che contengono molte operazioni di algebra
  • Carichi di lavoro che accedono alla memoria in modo sparso
  • Carichi di lavoro che richiedono aritmetica di alta precisione
  • Carichi di lavoro di rete neurali che contengono operazioni personalizzate nel loop di addestramento principale

Best practice per lo sviluppo del modello

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

Layout

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

Forme

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

Spaziatura interna

Un programma Cloud TPU ad alte prestazioni può essere suddiviso in blocchi di 128 x 128. Quando un calcolo matriciale non può occupare un'intera MXU, il compilatore padroneggia i tensori con zeri. La spaziatura interna presenta due svantaggi:

  • I tensori imbottiti con zeri sottoutilizzano il core della TPU.
  • Il Padding aumenta la quantità di spazio di archiviazione on-chip richiesta per un tensore e può portare a un errore di esaurimento della memoria in caso estremo.

Mentre la spaziatura interna viene eseguita automaticamente dal compilatore XLA quando necessario, è possibile determinare la quantità di spaziatura interna eseguita dallo strumento op_profile. Puoi evitare il padding scegliendo le dimensioni del tensore più adatte alle TPU.

Dimensioni

La scelta di dimensioni del tensore adatte è un grande passo avanti nell'estrazione delle massime prestazioni dall'hardware TPU, in particolare dall'MXU. Il compilatore XLA tenta di utilizzare la dimensione del batch o una dimensione della funzionalità per utilizzare al massimo l'MXU. Pertanto, uno di questi deve essere un multiplo di 128. In caso contrario, il compilatore padrà una di queste a 128. Idealmente, la dimensione del batch e le dimensioni delle caratteristiche dovrebbero 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 attorno alle risorse Cloud TPU e di controllare il movimento dei dati attraverso i confini del perimetro. Per saperne di più su Controlli di servizio VPC, consulta Panoramica di Controlli di servizio VPC. Per saperne di più sulle limitazioni dell'utilizzo di Cloud TPU con i Controlli di servizio VPC, consulta prodotti supportati e limitazioni.

Edge TPU

I modelli di machine learning addestrati nel cloud hanno sempre più bisogno di eseguire attacchi "a livello perimetrale", ovvero su dispositivi che operano sul perimetro 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 la batteria), Google ha progettato il coprocessore Edge TPU per accelerare l'inferenza ML sui dispositivi a basso consumo. Una singola 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 mobile all'avanguardia come MobileNet V2 a quasi 400 frame al secondo in modo efficiente dal punto di vista energetico.

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

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

Introduzione a Cloud TPU

Richiesta di assistenza

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

  • L'ID del tuo progetto Google Cloud
  • Il nome della tua 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: