Introduzione a Cloud TPU

Le TPU (Tensor Processing Unit) sono circuiti integrati per applicazioni specifiche (ASIC) sviluppati da Google e utilizzati per accelerare i carichi di lavoro di machine learning. Per informazioni più dettagliate sull'hardware TPU, consulta Architettura TPU. Cloud TPU è un servizio web che rende disponibili le TPU come risorse di calcolo scalabili su Google Cloud.

Le TPU addestrano i 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 ad alta larghezza di banda (HBM) che ti consente di utilizzare modelli e dimensioni di batch più grandi. Le TPU possono essere collegate in gruppi chiamati slice che scalano i carichi di lavoro con poche o nessuna modifica al codice.

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

Quando utilizzare le TPU

Le Cloud TPU sono ottimizzate per carichi di lavoro specifici. In alcuni casi, potresti voler utilizzare GPU o CPU su istanze 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 di piccole dimensioni con dimensioni dei batch efficaci ridotte
  • 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 di destinazione

GPU

  • Modelli con un numero significativo di operazioni PyTorch/JAX personalizzate che devono essere eseguite almeno parzialmente su CPU
  • Modelli con operazioni TensorFlow non disponibili su Cloud TPU (vedi l'elenco delle operazioni TensorFlow disponibili)
  • Modelli di medie e grandi dimensioni con dimensioni di batch effettive più grandi

TPU

  • Modelli in cui predominano i calcoli matriciali
  • Modelli senza operazioni PyTorch/JAX personalizzate all'interno del loop di addestramento principale
  • Modelli che vengono addestrati per settimane o mesi
  • Modelli di grandi dimensioni con grandi dimensioni di batch effettive
  • Modelli con embedding di grandi dimensioni comuni nei workload di ranking e consigli avanzati

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

  • Programmi di algebra lineare che richiedono frequenti ramificazioni o contengono molte operazioni algebriche elementari
  • Carichi di lavoro che richiedono operazioni aritmetiche ad alta precisione
  • Carichi di lavoro di reti neurali che contengono operazioni personalizzate nel loop di addestramento principale

TPU in Google Cloud

Puoi utilizzare le TPU tramite le VM Cloud TPU, Google Kubernetes Engine e Vertex AI. La tabella seguente elenca le risorse per ogni Google Cloud servizio.

Google Cloud servizio Risorse
Cloud TPU Iniziare a utilizzare le VM Cloud TPU
Google Kubernetes Engine

Informazioni sulle TPU in GKE

Eseguire Ray su GKE con le TPU
Vertex AI

Addestramento su Vertex AI con TPU

Utilizzare le TPU per la previsione online su Vertex AI

Best practice per lo sviluppo del modello

Un programma il cui calcolo è dominato da operazioni non matriciali come addizione, trasformazione o concatenazione probabilmente non raggiungerà 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, inclusa la suddivisione di una moltiplicazione di matrici in blocchi più piccoli, per eseguire in modo efficiente i calcoli sull'unità di matrice (MXU). Il compilatore XLA utilizza la struttura dell'hardware MXU, un array sistolica 128x128, e il design del sottosistema di memoria delle TPU, che preferisce dimensioni che sono multipli di 8, per l'efficienza del tiling. Di conseguenza, alcuni layout sono più adatti alla suddivisione in riquadri, mentre altri richiedono l'applicazione di modifiche alla forma prima di poter essere suddivisi in riquadri. Le operazioni di ricreazione della forma sono spesso vincolate dalla memoria su Cloud TPU.

Forme

Il compilatore XLA compila un grafo ML just in time per il primo batch. Se i batch successivi hanno forme diverse, il modello non funziona. (La ricomposiz ione del grafico ogni volta che la forma cambia è troppo lenta.) Pertanto, qualsiasi modello con tensori con forme dinamiche non è adatto alle TPU.

Spaziatura interna

Un programma Cloud TPU ad alte prestazioni è un programma in cui il calcolo intensivo può essere suddiviso in blocchi di 128 x 128. Quando un calcolo della matrice non può occupare un'intera MXU, il compilatore completa i tensori con zeri. Esistono due svantaggi per la spaziatura interna:

  • I tensori riempiti con zeri sottoutilizzano il core TPU.
  • Il padding aumenta la quantità di spazio di archiviazione in memoria on-chip richiesta per un tensore e può portare a un errore di memoria insufficiente nei casi estremi.

Sebbene il padding venga eseguito automaticamente dal compilatore XLA, se necessario, è possibile determinare la quantità di padding eseguita tramite lo strumento op_profile. Puoi evitare il padding scegliendo dimensioni dei tensori ben adatte alle TPU.

Dimensioni

La scelta di dimensioni dei tensori adeguate è fondamentale per ottenere il massimo delle prestazioni dall'hardware TPU, in particolare dall'unità MXU. Il compilatore XLA tenta di utilizzare la dimensione del batch o una dimensione della funzionalità per utilizzare al meglio l'unità MXU. Pertanto, uno di questi deve essere un multiplo di 128. In caso contrario, il compilatore aggiungerà un a capo a uno di questi fino 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.

Inizia a utilizzare Cloud TPU

Richiesta di assistenza

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

  • Il tuo Google Cloud ID progetto
  • Il nome della TPU, se esistente
  • Altre informazioni che vuoi fornire

Passaggi successivi

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