Configurazioni TPU

Configurazioni TPU v5p

Un pod TPU v5p è composto da 8960 chip interconnessi con link ad alta velocità riconfigurabili. Il networking flessibile di TPU v5p consente di connettere i chip in una sezione delle stesse dimensioni in più modi. Quando crei una sezione TPU utilizzando il comando gcloud compute tpus tpu-vm create, ne specifichi il tipo e la forma utilizzando i parametri AcceleratorType o AcceleratorConfig.

La seguente tabella mostra le forme a sezione singola più comuni supportate con v5p, oltre alla maggior parte delle forme di cubo complete (ma non tutte) maggiori di un cubo. La forma massima v5p è 16 x 16 x 24 (6144 chip, 96 cubi).

Forma della sezione Dimensioni VM # core # chip N. di macchine N. di cubi Supporta Twisted?
2x2x1 Host completo 8 4 1 N/A N/A
2x2x2 Host completo 16 8 2 N/A N/A
2x4x4 Host completo 64 32 8 N/A N/A
4x4x4 Host completo 128 64 16 1 N/A
4x4x8 Host completo 256 128 32 2 Sì
4x8x8 Host completo 512 256 64 4 Sì
8x8x8 Host completo 1024 512 128 8 N/A
8x8x16 Host completo 2048 1024 256 16 Sì
8x16x16 Host completo 4096 2048 512 32 Sì
16x16x16 Host completo 8192 4096 1024 64 N/A
16x16x24 Host completo 12288 6144 1536 96 N/A

L'addestramento su singola fetta è supportato per un massimo di 6144 chip. È estendibile a 18432 chip con multislice. Per informazioni dettagliate, consulta la Panoramica di Cloud TPU Multislice.

Utilizzo del parametro AcceleratorType

Quando assegni risorse TPU, utilizzi l'argomento --accelerator-type per specificare il numero di TensorCore in una sezione. --accelerator-type è una stringa formattata "v$VERSION_NUMBERp-$CORES_COUNT". Ad esempio, v5p-32 specifica una sezione TPU v5p con 32 TensorCore (16 chip).

Per eseguire il provisioning delle TPU per un job di addestramento v5p, utilizza uno dei seguenti tipi di acceleratore nella richiesta di creazione dell'API TPU o dell'interfaccia a riga di comando:

  • V5p-8
  • v5p-16
  • v5p-32
  • Versione 5p-64
  • v5p-128 (un cubo/rack pieno)
  • v5p-256 (2 cubi)
  • v5p-512
  • v5p-1024 ... v5p-12288

Utilizzo del parametro AcceleratorConfig

Per le versioni Cloud TPU v5p e successive, AcceleratorConfig viene utilizzato più o meno come con Cloud TPU v4. La differenza è che, invece di specificare il tipo di TPU come --type=v4, lo specifichi come la versione di TPU che stai utilizzando (ad esempio, --type=v5p per la release v5p).

Resilienza ICI di Cloud TPU

La resilienza ICI aiuta a migliorare la tolleranza di errore dei collegamenti ottici e degli switch di circuito ottici (OCS) che collegano le TPU tra i cubi. (le connessioni ICI all'interno di un cubo usano maglie in rame non interessate). La resilienza ICI consente il routing delle connessioni ICI intorno a OCS e guasti ICI ottici. Di conseguenza, migliora la disponibilità di pianificazione delle sezioni TPU, con il compromesso del degrado temporaneo delle prestazioni ICI.

Analogamente a Cloud TPU v4, la resilienza ICI è abilitata per impostazione predefinita per le sezioni v5p di almeno un cubo:

  • v5p-128 quando si specifica il tipo di accleratore
  • 4 x 4 x 4 quando si specifica la configurazione dell'acceleratore

Proprietà di VM, host e sezioni

Proprietà Valore in una TPU
N. di chip v5p 4
N. di vCPU 208 (solo la metà è utilizzabile se si utilizza l'associazione NUMA per evitare una penalizzazione delle prestazioni di più NUMA)
RAM (GB) 448 (solo la metà è utilizzabile se si utilizza l'associazione NUMA per evitare una penalizzazione delle prestazioni di più NUMA)
N. di NUMA nodi 2
Velocità effettiva NIC (Gbps) 200

Relazione tra il numero di TensorCore, chip, host/VM e cubi in un pod:

Core Chip Host/VM Cubi
Organizzatore 8 4 1
Cubo (noto anche come rack) 128 64 16 1
Sezione più grande supportata 12288 6144 1536 26
Pod completo v5p 17920 8960 2240 140

Configurazioni TPU v5e

Cloud TPU v5e è un prodotto combinato di addestramento e inferenza (pubblicazione). Per differenziare tra un ambiente di addestramento e un ambiente di inferenza, utilizza i flag AcceleratorType o AcceleratorConfig con l'API TPU o il flag --machine-type durante la creazione di un pool di nodi GKE.

I job di addestramento sono ottimizzati per velocità effettiva e disponibilità, mentre i job di servizio sono ottimizzati per la latenza. Quindi, un job di addestramento sulle TPU di cui è stato eseguito il provisioning per l'addestramento potrebbe avere una disponibilità inferiore e, analogamente, un job di addestramento eseguito sulle TPU di cui è stato eseguito il provisioning per l'addestramento potrebbe avere una latenza maggiore.

Usa AcceleratorType per specificare il numero di TensorCore che vuoi utilizzare. Puoi specificare il AcceleratorType durante la creazione di una TPU utilizzando gcloud CLI o la console Google Cloud. Il valore specificato per AcceleratorType è una stringa con il formato: v$VERSION_NUMBER-$CHIP_COUNT.

Puoi anche usare AcceleratorConfig per specificare il numero di TensorCore che vuoi utilizzare. Tuttavia, poiché non esistono varianti della topologia 2D personalizzata per TPU v5e, non esiste alcuna differenza tra l'uso di AcceleratorConfig e AcceleratorType.

Per configurare una TPU v5e utilizzando AcceleratorConfig, utilizza i flag --version e --topology. Imposta --version sulla versione TPU che vuoi utilizzare e --topology sulla disposizione fisica dei chip TPU nella sezione. Il valore specificato per AcceleratorConfig è una stringa con formato AxB, dove A e B corrispondono al numero di chip in ciascuna direzione.

Le seguenti forme di sezione 2D sono supportate per v5e:

Topologia Numero di chip TPU Numero di host
1x1 1 8/1
2x2 4 1/2
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

Ogni VM TPU in una sezione TPU v5e contiene 1, 4 o 8 chip. Nelle sezioni a 4 chip o più piccole, tutti i chip TPU condividono lo stesso nodo NUMA (Non Uniform Memory Access).

Per le VM TPU v5e a 8 chip, la comunicazione CPU-TPU sarà più efficiente all'interno di NUMA partizioni. Ad esempio, nella figura seguente, la comunicazione con CPU0-Chip0 sarà più veloce rispetto a quella con CPU0-Chip4.

Comunicazione nodo NUMA

Tipi di Cloud TPU v5e per la gestione

La pubblicazione su un singolo host è supportata per un massimo di 8 chip v5e. Sono supportate le seguenti configurazioni: sezioni 1x1, 2x2 e 2x4. Ogni sezione ha rispettivamente 1, 4 e 8 chip.

Configurazioni TPU v5e che supportano la pubblicazione: 1x1, 2x2 e 2x4.

Per eseguire il provisioning delle TPU per un job di gestione, utilizza uno dei seguenti tipi di acceleratore nella richiesta di creazione di TPU dell'interfaccia a riga di comando o di TPU dell'API:

Tipo di acceleratore (API TPU) Tipo di macchina (API GKE)
v5litepod-1 ct5lp-hightpu-1t
v5litepod-4 ct5lp-hightpu-4t
v5litepod-8 ct5lp-hightpu-8t

La pubblicazione su più di 8 chip v5e, chiamata anche gestione multi-host, è supportata utilizzando Sax. Per ulteriori informazioni, consulta Pubblicazione di modelli linguistici di grandi dimensioni (Large Language Model Serving).

Tipi di Cloud TPU v5e per l'addestramento

L'addestramento è supportato per un massimo di 256 chip.

Per eseguire il provisioning delle TPU per un job di addestramento v5e, utilizza uno dei seguenti tipi di acceleratore nella richiesta di creazione di TPU dell'interfaccia a riga di comando o di TPU dell'API:

Tipo di acceleratore (API TPU) Tipo di macchina (API GKE) Topologia
v5litepod-16 ct5lp-hightpu-4t 4x4
v5litepod-32 ct5lp-hightpu-4t 4x8
v5litepod-64 ct5lp-hightpu-4t 8x8
v5litepod-128 ct5lp-hightpu-4t 8x16
v5litepod-256 ct5lp-hightpu-4t 16x16

Confronto tra i tipi di VM TPU v5e:

Tipo di VM n2d-48-24-v5lite-tpu n2d-192-112-v5lite-tpu n2d-384-224-v5lite-tpu
N. di chip v5e 1 4 8
N. di vCPU 24 112 224
RAM (GB) 48 192 384
N. di NUMA nodi 1 1 2
Si applica a V5litepod-1 v5litepod-4 V5litepod-8
Interruzione Alto Medie Bassa

Per fare spazio ai carichi di lavoro che richiedono più chip, gli scheduler potrebbero prerilasciare le VM con meno chip. È probabile quindi che le VM a 8 chip prerilasciano le VM a 1 e 4 chip.

Configurazioni TPU v4

Un pod TPU v4 è composto da 4096 chip interconnessi con link ad alta velocità riconfigurabili. Il networking flessibile di TPU v4 consente di connettere i chip in una sezione di pod delle stesse dimensioni in più modi. Quando crei una sezione di pod di TPU, specifichi la versione di TPU e il numero di risorse TPU richieste. Quando crei una sezione di pod TPU v4, puoi specificarne tipo e dimensioni in uno dei due modi seguenti: AcceleratorType e AcceleratorConfig.

Uso: AcceleratorType

Utilizza AcceleratorType quando non specifichi una topologia. Per configurare le TPU v4 utilizzando AcceleratorType, usa il flag --accelerator-type durante la creazione della sezione di pod di TPU. Imposta --accelerator-type su una stringa contenente la versione TPU e il numero di TensorCore che vuoi utilizzare. Ad esempio, per creare una sezione di pod v4 con 32 TensorCore, devi utilizzare --accelerator-type=v4-32.

Il seguente comando crea una sezione di pod di TPU v4 con 512 TensorCore utilizzando il flag --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-512
    --version=tpu-vm-tf-2.16.1-pod-pjrt

Il numero dopo la versione TPU (v4) specifica il numero di TensorCore. In una TPU v4 sono presenti due TensorCore, quindi il numero di chip TPU sarebbe 512/2 = 256.

Uso: AcceleratorConfig

Usa AcceleratorConfig quando vuoi personalizzare la topologia fisica della sezione TPU. Questa operazione è generalmente necessaria per l'ottimizzazione delle prestazioni con sezioni di pod superiori a 256 chip.

Per configurare le TPU v4 utilizzando AcceleratorConfig, utilizza i flag --version e --topology. Imposta --version sulla versione TPU che vuoi utilizzare e --topology sulla disposizione fisica dei chip TPU nella sezione di pod.

Puoi specificare una topologia TPU utilizzando AxBxC a 3 tuple, dove A<=B<=C e A, B, C sono tutti <= 4 o sono tutti multipli interi di 4. I valori A, B e C sono il conteggio dei chip in ognuna delle tre dimensioni. Ad esempio, per creare una sezione di pod v4 con 16 chip, devi impostare --version=v4 e --topology=2x2x4.

Il seguente comando crea una sezione di pod di TPU v4 con 128 chip TPU disposti in un array 4 x 4 x 8:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=4x4x8
    --version=tpu-vm-tf-2.16.1-pod-pjrt

Le topologie in cui 2A=B=C o 2A=2B=C hanno anche varianti di topologia ottimizzate per una comunicazione complessiva, ad esempio 4 × 4 × 8, 8 × 8 × 16 e 12 × 12 × 24. Queste sono note come topologie tori torsi.

Le illustrazioni seguenti mostrano alcune topologie TPU v4 comuni.

immagine

Le sezioni di pod più grandi possono essere create da uno o più "cubi" 4x4x4 di chip.

Topologie Tori intrecciate

Alcune forme di sezione del toro 3D v4 hanno la possibilità di utilizzare quella che viene nota come topologia toro contorto. Ad esempio, due cubi v4 possono essere disposti come sezione 4x4x8 o 4x4x8_twisted. Le topologie intrecciate offrono una bisezione di larghezza di banda significativamente maggiore. Una maggiore larghezza di banda di bisezione è utile per carichi di lavoro che utilizzano schemi di comunicazione globali. Le topologie intrecciate possono migliorare le prestazioni per la maggior parte dei modelli, tra i maggiori vantaggi per carichi di lavoro di incorporamento TPU di grandi dimensioni.

Per i carichi di lavoro che utilizzano il parallelismo dei dati come unica strategia di parallelismo, le topologie contorte potrebbero funzionare leggermente migliori. Per gli LLM, le prestazioni con una topologia intrecciata possono variare in base al tipo di parallelismo (DP, MP e così via). La best practice prevede l'addestramento dell'LLM con e senza una topologia contorta per determinare quale offre le migliori prestazioni per il tuo modello. Alcuni esperimenti sul modello FSDP MaxText hanno rilevato 1-2 miglioramenti MFU utilizzando una topologia contorta.

Il vantaggio principale delle topologie intrecciate è che trasforma una topologia torus asimmetrica (ad esempio, 4 × 4 × 8) in una topologia simmetrica strettamente correlata. La topologia simmetrica presenta numerosi vantaggi:

  • Bilanciamento del carico migliorato
  • Maggiore larghezza di banda di bisezione
  • Route di pacchetti più brevi

Questi vantaggi si traducono in un miglioramento delle prestazioni per molti modelli di comunicazione globali.

Il software TPU supporta tori intrecciati sulle sezioni in cui la dimensione di ogni dimensione è uguale o doppia di quella della dimensione più piccola. Ad esempio, 4 x 4 x 8, 4 x 8 × 8 o 12 x 12 x 24.

Ad esempio, considera questa topologia toro 4 × 2 con TPU etichettate con le relative coordinate (X,Y) nella sezione:

I bordi in questo grafico della topologia sono mostrati come bordi non orientati per maggiore chiarezza. In pratica, ogni perimetro è una connessione bidirezionale tra TPU. I bordi tra un lato di questa griglia e l'altro vengono definiti bordi avvolgenti, come indicato nel diagramma.

disegno

Torcando questa topologia, si ottiene una topologia toroidale 4 × 2 completamente simmetrica:

disegno

Tutto è cambiato tra questo diagramma e quello precedente sono i bordi arrotondati a Y. Invece di connettersi a un'altra TPU con la stessa coordinata X, sono stati spostati per connettersi alla TPU con la coordinata X+2 mod 4.

La stessa idea si applica a dimensioni di dimensioni diverse e a numeri diversi. La rete risultante è simmetrica, a condizione che ogni dimensione sia uguale o doppia della dimensione della dimensione più piccola.

La seguente tabella mostra le topologie intrecciate supportate e un aumento teorico della larghezza di banda di bisezione rispetto alle topologie non distorte.

Topologia contorta Aumento teorico della bisezione
della larghezza di banda rispetto a un toro non distorto
4 × 4 × 8_torto ~70%
8x8x16_twisted
12 × 12 × 24_attorcigliato
4×8×8_twisted Circa il 40%
8 × 16 × 16_torto

Varianti della topologia TPU v4

Alcune topologie contenenti lo stesso numero di chip possono essere disposte in modi diversi. Ad esempio, una sezione di pod di TPU con 512 chip (1024 TensorCore) può essere configurata utilizzando le seguenti topologie: 4x4x32, 4x8x16 o 8x8x8. Una sezione di pod di TPU con 2048 chip (4096 TensorCore) offre ancora più opzioni di topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16. Una sezione di pod di TPU con 2048 chip (4096 TensorCores) offre ancora più opzioni di topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16.

La topologia predefinita associata a un determinato numero di chip è quella più simile a un cubo (vedi Topologia v4). Questa forma è probabilmente la scelta migliore per l'addestramento ML con parallelo dei dati. Altre topologie possono essere utili per carichi di lavoro con più tipi di parallelismo (ad esempio, parallelismo di modelli e dati o partizionamento spaziale di una simulazione). Questi carichi di lavoro funzionano meglio se la topologia è abbinata al parallelismo utilizzato. Ad esempio, il posizionamento del parallelismo del modello a 4 vie nella dimensione X e del parallelismo dei dati a 256 vie nelle dimensioni Y e Z corrisponde a una topologia 4 x 16 x 16.

I modelli con più dimensioni di parallelismo funzionano meglio con le dimensioni di parallelismo mappate a dimensioni di topologia TPU. In genere si tratta di modelli linguistici di grandi dimensioni (LLM) paralleli dati e modelli. Ad esempio, per una sezione di pod TPU v4 con topologia 8 x 16 x 16, le dimensioni della topologia TPU sono 8, 16 e 16. È più efficiente utilizzare il parallelismo dei modelli a 8 o 16 vie (mappato a una delle dimensioni fisiche della topologia TPU). Un parallelismo del modello a 4 vie non sarebbe ottimale con questa topologia, poiché non è allineato con nessuna delle dimensioni della topologia TPU, ma sarebbe ottimale con una topologia 4 x 16 x 32 sullo stesso numero di chip.

Le configurazioni TPU v4 sono composte da due gruppi, quelli con topologie più piccole di 64 chip (topologie piccole) e quelli con topologie superiori a 64 chip (topologie grandi).

Topologie v4 piccole

Cloud TPU supporta le seguenti sezioni di pod TPU v4 più piccole di 64 chip, un cubo 4 x 4 x 4. Puoi creare queste piccole topologie v4 utilizzando il relativo nome basato su TensorCore (ad esempio v4-32) o la relativa topologia (ad esempio, 2x2x4):

Nome (basato sul conteggio di TensorCore) Numero di chip Topologia
v4-8 4 2x2x1
v4-16 8 2x2x2
v4-32 16 2x2x4
v4-64 32 2x4x4

Topologie v4 grandi

Le sezioni di pod TPU v4 sono disponibili in incrementi di 64 chip, con forme multiple di 4 su tutte e tre le dimensioni. Anche le dimensioni devono essere in ordine crescente. Nella tabella seguente sono riportati diversi esempi. Alcune di queste topologie sono topologie "personalizzate" che possono essere create solo utilizzando i flag --type e --topology perché esiste più di un modo per organizzare i chip.

Il seguente comando crea una sezione di pod di TPU v4 con 512 chip TPU disposti in un array 8 x 8 x 8:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=8x8x8
    --version=tpu-vm-tf-2.16.1-pod-pjrt

Puoi creare una sezione di pod di TPU v4 con lo stesso numero di TensorCore utilizzando --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-1024
    --version=tpu-vm-tf-2.16.1-pod-pjrt
Nome (basato sul conteggio di TensorCore) Numero di chip Topologia
v4-128 64 4x4x4
v4-256 128 4x4x8
v4-512 256 4x8x8
N/A: è necessario utilizzare i flag --type e --topology 256 4x4x16
v4-1024 512 8x8x8
v4-1536 768 8x8x12
v4-2048 1024 8x8x16
N/A: è necessario utilizzare i flag --type e --topology 1024 4x16x16
v4-4096 2048 8x16x16
… … …

Configurazioni TPU v3

Un pod TPU v3 è composto da 1024 chip interconnessi con link ad alta velocità. Per creare una sezione di pod o dispositivo TPU v3, utilizza il flag --accelerator-type per il comando gcloud compute tpus tpu-vm. Puoi specificare il tipo di acceleratore specificando la versione TPU e il numero di core TPU. Per una singola TPU v3, utilizza --accelerator-type=v3-8. Per una sezione di pod v3 con 128 TensorCore, utilizza --accelerator-type=v3-128.

Il seguente comando mostra come creare una sezione di pod di TPU v3 con 128 TensorCore:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v3-128
    --version=tpu-vm-tf-2.16.1-pjrt

Nella tabella seguente sono elencati i tipi di TPU v3 supportati:

Versione TPU Fine del supporto
v3-8 (Data di fine non ancora impostata)
v3-32 (Data di fine non ancora impostata)
v3-128 (Data di fine non ancora impostata)
v3-256 (Data di fine non ancora impostata)
v3-512 (Data di fine non ancora impostata)
v3-1024 (Data di fine non ancora impostata)
v3-2048 (Data di fine non ancora impostata)

Per saperne di più sulla gestione delle TPU, vedi Gestire le TPU. Per ulteriori informazioni sulle diverse versioni di Cloud TPU, consulta Architettura di sistema.

Configurazioni TPU v2

Un pod TPU v2 è composto da 512 chip interconnessi con link ad alta velocità riconfigurabili. Per creare una sezione di pod TPU v2, utilizza il flag --accelerator-type per il comando gcloud compute tpus tpu-vm. Puoi specificare il tipo di acceleratore specificando la versione di TPU e il numero di core TPU. Per una singola TPU v2, utilizza --accelerator-type=v2-8. Per una sezione di pod v2 con 128 TensorCore, utilizza --accelerator-type=v2-128.

Il seguente comando mostra come creare una sezione di pod di TPU v2 con 128 TensorCore:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v2-128
    --version=tpu-vm-tf-2.16.1-pjrt

Per saperne di più sulla gestione delle TPU, vedi Gestire le TPU. Per ulteriori informazioni sulle diverse versioni di Cloud TPU, consulta Architettura di sistema.

La seguente tabella elenca i tipi di TPU v2 supportati

Versione TPU Fine del supporto
v2-8 (Data di fine non ancora impostata)
v2-32 (Data di fine non ancora impostata)
v2-128 (Data di fine non ancora impostata)
v2-256 (Data di fine non ancora impostata)
v2-512 (Data di fine non ancora impostata)

Compatibilità del tipo di TPU

Puoi cambiare il tipo di TPU scegliendo un altro tipo di TPU con lo stesso numero di TensorCore o chip (ad esempio, v3-128 e v4-128) ed eseguire lo script di addestramento senza modifiche al codice. Tuttavia, se passi a un tipo di TPU con un numero maggiore o minore di TensorCore, dovrai eseguire un'ottimizzazione e un'ottimizzazione significativi. Per maggiori informazioni, consulta Addestramento sui pod TPU.

Versioni software delle VM TPU

Questa sezione descrive le versioni del software TPU da utilizzare per una TPU con l'architettura VM TPU. Per l'architettura dei nodi TPU, consulta la pagina relativa alle versioni software dei nodi TPU.

Sono disponibili versioni software TPU per i framework TensorFlow, PyTorch e JAX.

TensorFlow

Utilizza la versione software TPU corrispondente alla versione di TensorFlow con cui è stato scritto il modello.

A partire da TensorFlow 2.15.0, devi anche specificare il runtime dell'esecutore di flussi (SE) o il runtime PJRT. Ad esempio, se utilizzi TensorFlow 2.16.1 con il runtime PJRT, usa la versione software della TPU tpu-vm-tf-2.16.1-pjrt. Le versioni precedenti a TensorFlow 2.15.0 supportano solo l'esecutore dei flussi. Per ulteriori informazioni su PJRT, consulta la pagina relativa all'assistenza di TensorFlow PJRT.

Le attuali versioni software delle VM TensorFlow TPU supportate sono:

  • tpu-vm-tf-2.16.1-pjrt
  • tpu-vm-tf-2.16.1-se
  • tpu-vm-tf-2.15.0-pjrt
  • tpu-vm-tf-2.15.0-se
  • tpu-vm-tf-2.14.1
  • tpu-vm-tf-2.14.0
  • tpu-vm-tf-2.13.1
  • tpu-vm-tf-2.13.0
  • tpu-vm-tf-2.12.1
  • tpu-vm-tf-2.12.0
  • tpu-vm-tf-2.11.1
  • tpu-vm-tf-2.11.0
  • tpu-vm-tf-2.10.1
  • tpu-vm-tf-2.10.0
  • tpu-vm-tf-2.9.3
  • tpu-vm-tf-2.9.1
  • tpu-vm-tf-2.8.4
  • tpu-vm-tf-2.8.3
  • tpu-vm-tf-2.8.0
  • tpu-vm-tf-2.7.4
  • tpu-vm-tf-2.7.3

Per ulteriori informazioni sulle versioni della patch TensorFlow, consulta Versioni della patch di TensorFlow supportate.

Supporto di TensorFlow PJRT

A partire da TensorFlow 2.15.0, puoi utilizzare l'interfaccia PJRT per TensorFlow su TPU. PJRT offre la deframmentazione automatica della memoria dei dispositivi e semplifica l'integrazione dell'hardware con i framework. Per ulteriori informazioni su PJRT, consulta il documento PJRT: Semplificare l'integrazione di Hardware e Framework per ML sul blog open source di Google.

Non tutte le funzionalità di TPU v2, v3 e v4 sono state migrate al runtime PJRT. La seguente tabella descrive quali funzionalità sono supportate su PJRT o esecutore a vapore.

Acceleratore Selezione delle Supportato su PJRT Supportato sull'esecutore degli stream
TPU v2-v4 Dense computing (nessuna API di incorporamento TPU) Sì Sì
TPU v2-v4 API Dense Compute + API Embedding TPU No Sì
TPU v2-v4 tf.summary/tf.print con posizionamento soft device No Sì
TPU v5e Dense computing (nessuna API di incorporamento TPU) Sì No
TPU v5e API TPU Embedding N/A: TPU v5e non supporta l'API TPU Embedding N/A
TPU v5p Dense computing (nessuna API di incorporamento TPU) Sì No
TPU v5p API TPU Embedding Sì No

TPU v4 con TensorFlow versioni 2.10.0 e precedenti

Se stai addestrando un modello su TPU v4 con TensorFlow, le versioni 2.10.0 e precedenti di TensorFlow utilizzano versioni specifiche di v4 mostrate nella tabella seguente. Se la versione di TensorFlow che stai utilizzando non è mostrata nella tabella, segui le indicazioni riportate nella sezione TensorFlow.

Versione TensorFlow Versione software TPU
2.10.0 tpu-vm-tf-2.10.0-v4, tpu-vm-tf-2.10.0-pod-v4
2.9.3 tpu-vm-tf-2.9.3-v4, tpu-vm-tf-2.9.3-pod-v4
2.9.2 tpu-vm-tf-2.9.2-v4, tpu-vm-tf-2.9.2-pod-v4
2.9.1 tpu-vm-tf-2.9.1-v4, tpu-vm-tf-2.9.1-pod-v4

Versioni Libtpu

Le VM TPU vengono create con TensorFlow e con la libreria Libtpu corrispondente preinstallata. Se stai creando una tua immagine VM, specifica le seguenti versioni software TensorFlow TPU e le versioni libtpu corrispondenti:

Versione TensorFlow Versione libtpu.so
2.16.1 1.10.1
2.15.0 1.9.0
2.14.1 1.8.1
2.14.0 1.8.0
2.13.1 1.7.1
2.13.0 1.7.0
2.12.1 1.6.1
2.12.0 1.6.0
2.11.1 1.5.1
2.11.0 1.5.0
2.10.1 1.4.1
2.10.0 1.4.0
2.9.3 1.3.2
2.9.1 1.3.0
2.8.3 1.2.3
2,8.* 1.2.0
2.7.3 1.1.2

PyTorch

Utilizza la versione del software TPU corrispondente alla versione di PyTorch con cui è stato scritto il modello. Ad esempio, se utilizzi PyTorch 1.13 e TPU v2 o v3, utilizza la versione software TPU tpu-vm-pt-1.13. Se usi TPU v4, usa la versione software TPU tpu-vm-v4-pt-1.13. Per i pod TPU viene utilizzata la stessa versione software TPU (ad esempio v2-32, v3-128, v4-32). Le versioni attuali del software TPU supportate sono:

TPU v2/v3:

  • tpu-vm-pt-2.0 (pytorch-2.0)
  • tpu-vm-pt-1.13 (pytorch-1.13)
  • tpu-vm-pt-1.12 (pytorch-1.12)
  • tpu-vm-pt-1.11 (pytorch-1.11)
  • tpu-vm-pt-1.10 (pytorch-1.10)
  • v2-alpha (pytorch-1.8.1)

TPU v4:

  • tpu-vm-v4-pt-2.0 (pytorch-2.0)
  • tpu-vm-v4-pt-1.13 (pytorch-1.13)

TPU v5 (v5e e v5p):

  • v2-alpha-tpuv5 (pytorch-2.0)

Quando crei una VM TPU, la versione più recente di PyTorch è preinstallata sulla VM TPU. Quando installi PyTorch, viene installata automaticamente la versione corretta di libtpu.so.

Per modificare la versione corrente del software PyTorch, consulta Modificare la versione di PyTorch.

JAX

Devi installare manualmente JAX sulla tua VM TPU. Non esiste una versione del software TPU (runtime) specifica di JAX per TPU v2 e v3. Per le versioni TPU successive, utilizza le seguenti versioni software:

  • TPU v4: tpu-vm-v4-base
  • TPU v5e: v2-alpha-tpuv5
  • TPU v5p: v2-alpha-tpuv5

La versione corretta di libtpu.so viene installata automaticamente quando installi JAX.

Versioni software dei nodi TPU

Questa sezione descrive le versioni del software TPU da utilizzare per una TPU con l'architettura Nodo TPU. Per l'architettura delle VM TPU, consulta la pagina relativa alle versioni del software delle VM TPU.

Sono disponibili versioni software TPU per i framework TensorFlow, PyTorch e JAX.

TensorFlow

Utilizza la versione software TPU corrispondente alla versione di TensorFlow con cui è stato scritto il modello. Ad esempio, se utilizzi TensorFlow 2.12.0, usa la versione software TPU 2.12.0. Le versioni software TPU specifiche per TensorFlow sono:

  • 2.12.1
  • 2.12.0
  • 2.11.1
  • 2.11.0
  • 2.10.1
  • 2.10.0
  • 2.9.3
  • 2.9.1
  • 2.8.4
  • 2.8.2
  • 2.7.3

Per ulteriori informazioni sulle versioni della patch TensorFlow, consulta Versioni della patch di TensorFlow supportate.

Quando crei un nodo TPU, la versione più recente di TensorFlow è preinstallata sul nodo TPU.

PyTorch

Utilizza la versione del software TPU corrispondente alla versione di PyTorch con cui è stato scritto il tuo modello. Ad esempio, se usi PyTorch 1.9, usa la versione software pytorch-1.9.

Le versioni software TPU specifiche di PyTorch sono:

  • Pytorch-2.0
  • Pytorch-1.13
  • Pytorch-1.12
  • Pytorch-1.11
  • pytorch-1.10
  • Pytorch-1.9
  • Pytorch-1.8
  • Pytorch-1.7

  • Pytorch-1.6

  • pytorch-notte

Quando crei un nodo TPU, la versione più recente di PyTorch è preinstallata sul nodo TPU.

JAX

Devi installare manualmente JAX sulla tua VM TPU, altrimenti non esiste una versione del software TPU specifica per JAX preinstallata. Puoi utilizzare una qualsiasi delle versioni software elencate per TensorFlow.

Passaggi successivi