Architettura di sistema

Le TPU (Tensor Processing Unit) sono acceleratori ML progettati da Google. Cloud TPU rende le TPU disponibili come risorsa cloud GCP scalabile. Per eseguire carichi di lavoro di machine learning su Cloud TPU, utilizza framework di machine learning come TensorFlow, Pytorch e JAX.

L'unità moltiplicatore di matrice (MXU) è composta da moltiplicatori/accumulatori 128 x 128 in un array sistolico. Gli MXU forniscono la maggior parte della potenza di calcolo di un TensorCore (core TPU). Ogni MXU è in grado di eseguire 16.000 operazioni per moltiplicare-accumulare in ogni ciclo utilizzando il formato numerico bfloat16.

La VPU viene utilizzata per i calcoli generici, come attivazioni, softmax e così via. L'unità scalare viene utilizzata per il flusso di controllo, il calcolo dell'indirizzo di memoria e altre operazioni di manutenzione.

Versioni TPU

Il layout esatto di una TPU dipende dalla versione TPU in uso. I dettagli architettonici e le caratteristiche delle prestazioni di TPU v2 e v3 sono disponibili in Un supercomputer specifico del dominio per l'addestramento di reti neurali profonde.

TPU v4

La configurazione più piccola della TPU TPU v4 contiene quattro chip TPU e 32 GiB di memoria HBM e 128 MiB di memoria comune condivisa. Ogni chip TPU contiene due TensorCore. Ogni TensorCore ha quattro MXU, un'unità vettoriali e una scala. Il seguente diagramma illustra un chip TPU v4.

immagine

TPU v3

La configurazione TPU più piccola contiene quattro chip TPU e 32 GiB di HBM. Ogni chip TPU contiene due TensorCore. Ogni TensorCore ha due MXU, un'unità vettoriale e una scalabile. Il seguente diagramma illustra un chip TPU v3.

immagine

TPU v2

La configurazione più piccola di TPU v2 contiene quattro chip TPU e 16 GiB di HBM. Ogni chip TPU contiene due TensorCore. Ogni TensorCore ha un'unità MXU, un'unità vettoriali e una scalabile. Il seguente diagramma illustra un chip TPU v2.

immagine

Cloud TPU fornisce le seguenti configurazioni TPU:

  • Un singolo dispositivo TPU
  • Un pod di TPU - un gruppo di dispositivi TPU connessi da interconnessioni ad alta velocità
  • Una sezione TPU: una suddivisione di un pod TPU

Vantaggi delle prestazioni di TPU v4 rispetto a v3

Sistema di memoria:

  • I chip TPU v4 hanno uno spazio di memoria HBM unificato di 32 GiB in tutto il chip, consentendo un migliore coordinamento tra i due TensorCore sul chip.
  • Prestazioni HBM migliorate utilizzando i più recenti standard di memoria e velocità.
  • È stato migliorato il profilo delle prestazioni DMA con il supporto nativo per una velocità di elaborazione ad alte prestazioni con granularità 512B.

TensorCore:

  • Il doppio del numero di MXU e una frequenza di clock superiore con una frequenza massima di 275 TFLOPS.
  • Larghezza di banda doppia di trasposizione e permutazione.
  • Carica il modello di accesso alla memoria per la memoria comune (Cmem).
  • Larghezza di banda per il caricamento del peso MXU più rapida e supporto per la modalità a 8 bit per consentire dimensioni del batch inferiori e migliore latenza di inferenza.

Interconnessione chip:

6 link di interconnessione per chip per abilitare topologie di rete con diametri di rete più piccoli.

Altro:

  • Interfaccia x16 PCIE gen3 per l'hosting (connessione diretta).
  • Modello di sicurezza migliorato.
  • Maggiore efficienza energetica.

Vantaggi delle prestazioni di TPU v3 rispetto a v2

L'aumento dei FLOPS per capacità di core e memoria nelle configurazioni di TPU v3 può migliorare le prestazioni dei tuoi modelli nei seguenti modi:

  • Le configurazioni di TPU v3 offrono vantaggi significativi in termini di prestazioni per core per i modelli con vincoli di calcolo. I modelli basati sulla memoria sulle configurazioni TPU v2 potrebbero non ottenere lo stesso miglioramento delle prestazioni se sono limitati anche sulla memoria nelle configurazioni TPU v3.

  • Nei casi in cui i dati non rientrano nella memoria sulle configurazioni TPU v2, TPU v3 può fornire prestazioni migliori e una riduzione del calcolo dei valori intermedi (riprogettazione).

  • Le configurazioni di TPU v3 possono eseguire nuovi modelli con dimensioni batch che non erano adatte alle configurazioni di TPU v2. Ad esempio, TPU v3 potrebbe consentire ResNet più approfonditi e immagini più grandi con RetinaNet.

I modelli che sono quasi associati a un input ("infeed") su TPU v2 perché i passaggi di addestramento in attesa di input potrebbero essere associati anche a Cloud TPU v3. La guida alle prestazioni della pipeline può aiutarti a risolvere i problemi relativi agli annunci in-feed.

Configurazioni di TPU v4

Un pod TPU v4 è composto da 4096 chip interconnessi con link riconfigurabili ad alta velocità. Il networking flessibile di TPU v4 ti consente di collegare i chip di una sezione della stessa dimensione in diversi modi. Per configurare le TPU v4 utilizzando AcceleratorConfig. AcceleratorConfig consente di specificare le dimensioni della sezione in termini di chip nel campo della topologia.

La topologia TPU v4 viene specificata utilizzando una tupla tre, che descrive il numero, la forma e le interconnessioni tra i chip TPU. Le seguenti illustrazioni mostrano alcune topologie TPU v4 comuni.

immagine

Le fette più grandi possono essere costruite con uno o più cubetti di 4 x 4 x 4 pollici.

Le sezioni di TPU di un determinato numero di chip possono essere configurate in diversi modi. Ad esempio, una porzione di TPU con v4-1024 può ora essere configurata come: 4x4x32, 4x8x16, 8x8x8. Una TPU v4-4096 offre ancora più flessibilità: 4x4 x 128, 4 x 86, 4 x 16 x 32, 8 x 16 x 16, 8 x 16 x 16_twisted. Per ulteriori informazioni, consulta la pagina Tipi e topologie TPU.

Configurazioni di TPU v2 e v3

Le TPU sono disponibili nelle seguenti configurazioni:

  1. Una singola lavagna TPU
  2. Un pod TPU
  3. Una porzione di pod di TPU

Scheda TPU singola

Una configurazione TPU a scheda singola è una scheda autonoma con 4 chip TPU (8 TensorCore) senza connessioni di rete ad altre schede TPU. Le TPU singolo non fanno parte della configurazione di un pod di TPU e non occupano una parte del pod.

Pod e sezioni di TPU

In un'istanza di pod pod o TPU, i chip TPU sono connessi tramite un'interconnessione ad alta velocità e ogni chip TPU comunica direttamente con gli altri chip sul dispositivo TPU. Il software TPU gestisce automaticamente la distribuzione dei dati a ciascun core TPU in un pod o una sezione. Le sezioni dei pod sono disponibili con core da 32, 128, 512, 1024 o 2048.

Architetture VM di Cloud TPU

Le TPU possono eseguire solo operazioni di matrice, quindi ogni scheda TPU è connessa a una macchina host basata su CPU per eseguire operazioni che non possono essere eseguite sulla TPU. Le macchine host sono responsabili del caricamento dei dati da Cloud Storage, della pre-elaborazione e dell'invio dei dati alla TPU.

immagine

In un pod TPU c'è un host TPU per ogni scheda TPU.

immagine

Il modo in cui interagisci con l'host della TPU (e la scheda TPU) dipende dall'architettura della VM TPU che utilizzi: nodi nodo o VM TPU.

Nodi TPU

L'architettura dei nodi TPU è costituita da una VM utente che comunica con l'host TPU su gRPC. Se utilizzi questa architettura, non puoi accedere direttamente all'host della TPU. Questo può rendere difficile il debug degli errori di addestramento e TPU.

immagine

VM TPU

L'architettura delle VM TPU elimina la necessità di utilizzare la VM dell'utente e puoi utilizzare SSH per connetterti direttamente alla VM collegata fisicamente al dispositivo TPU. Hai accesso root alla VM, quindi puoi eseguire codice arbitrario. Puoi accedere ai log di debug e di debug del compilatore e di runtime.

immagine

Framework supportati

Framework come JAX, PyTorch e TensorFlow accedono alle TPU tramite una libreria condivisa denominata libtpu presente su ogni VM TPU. Questa libreria include il compilatore XLA utilizzato per compilare i programmi TPU, il software TPU utilizzato per eseguire i programmi compilati e il driver TPU utilizzato dal runtime per l'accesso di basso livello alla TPU.

immagine

TensorFlow

VM TPU

Con le VM TPU, anziché eseguire il codice Python sulla VM utente, puoi eseguirlo direttamente sull'host TPU.

immagine

Per saperne di più su TensorFlow e Cloud TPU, vedi Esecuzione di modelli TensorFlow su Cloud TPU.

Nodi TPU

L'architettura di sistema del nodo Cloud TPU è stata creata in origine per TensorFlow. Gli host TPU non sono accessibili all'utente ed eseguono una copia headless del server TensorFlow. Non eseguono Python o qualsiasi codice utente non rappresentato come un grafico TensorFlow. Il codice utente viene eseguito in una VM utente separata che comunica con gli host TPU sulla rete gRPC.

immagine

PyTorch

VM TPU

Con le VM TPU, il codice PyTorch viene eseguito direttamente sugli host TPU.

immagini

Per ulteriori informazioni su PyTorch e Cloud TPU, vedi Esecuzione di modelli PyTorch su Cloud TPU

Nodi TPU

PyTorch viene eseguito sull'architettura dei nodi Cloud TPU utilizzando una libreria denominata XRT, che consente di inviare grafici XLA e istruzioni di runtime sulle connessioni TensorFlow TensorFlow e di eseguirle sui server TensorFlow. Per ogni host TPU è richiesta una VM dell'utente.

immagine

JAX

VM TPU

Con le VM TPU non è necessario utilizzare le VM degli utenti, perché puoi eseguire il codice direttamente sugli host delle TPU.

immagine

Per ulteriori informazioni sull'esecuzione di JAX su Cloud TPU, consulta la guida rapida di JAX.

Nodi TPU

JAX sui nodi Cloud TPU viene eseguito in modo simile a PyTorch, in quanto è richiesta una VM utente separata per ogni VM host.

immagine

Passaggi successivi