Tipi e topologie di TPU

Configurazioni TPU

Quando crei una TPU, devi specificare una configurazione TPU. Puoi specificare una configurazione TPU in termini di TensorCore o chip di TPU. La specifica di una configurazione TPU in termini di chip è necessaria solo negli scenari avanzati in cui addestra un modello di grandi dimensioni e ne ottimizzi le prestazioni.

Configurazione basata su TensorCore

Puoi utilizzare una configurazione basata su TensorCore con le TPU v2-v4. Descrivi le configurazioni basate su TensorCore con una versione TPU e il numero di TensorCore che vorresti utilizzare. Ad esempio, v4-128 specifica una configurazione con TPU v4 e 128 TensorCore. Specifica una configurazione basata su TensorCore con il flag --accelerator-type del comando gcloud compute tpus tpu-vm create:

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

Configurazione TPU basata su chip

Puoi utilizzare una configurazione basata su chip TPU con TPU v4 e versioni successive in TPU. Descrivi le configurazioni basate su chip con una versione TPU e una topologia (una disposizione fisica dei chip TPU). Specifica una versione TPU con una stringa come "v4". Specifica una topologia con una tupla a tre tuffi che descrive come i chip sono disposti in tre dimensioni. Puoi specificare una configurazione basata su chip con i flag --version e --topology del comando gcloud compute tpus tpu-vm create. Ad esempio, il seguente comando specifica TPU v4 come versione TPU e 2x2x4 (16 chip, ognuno con 2 TensorCore) come topologia: 

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=2x2x4
    --version=tpu-vm-tf-2.12.0

Ulteriori informazioni sulle diverse versioni e topologie di Cloud TPU sono disponibili in Architettura di sistema.

La rappresentazione di una topologia v4 è AxBxC, dove A<=B<=C e A, B, C sono tutti <= 4 o tutti i multipli interi di 4. I valori A, B e C sono i conteggi dei chip in ciascuna delle tre dimensioni. Topologie in cui 2A=B=C o 2A=2B=C hanno anche varianti di topologia ottimizzate per la comunicazione all-to-all, ad esempio 4×4×8, 8×8×16 e 12×12×24.

Per ulteriori informazioni sulle topologie v4, consulta la sezione Configurazioni TPU v4.

Varianti di topologia

La topologia predefinita associata a un determinato numero di chip è quella più simile a un cubo (vedi Forme di topologia). Questa forma è probabilmente la scelta migliore per l'addestramento ML con dati paralleli. Altre topologie possono essere utili per i 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 hanno prestazioni migliori se la topologia viene abbinata al parallelismo utilizzato. Ad esempio, posizionando il parallelismo dei modelli a 4 vie sulla dimensione X e il parallelismo dei dati a 256 vie sulle dimensioni Y e Z corrisponde a una topologia di 4x16x16.

I modelli con più dimensioni di parallelismo hanno un rendimento migliore perché le loro dimensioni di parallelismo vengono mappate a quelle della topologia di TPU. Di solito si tratta di modelli LLM (Data Language Model) paralleli di dati + modelli. Ad esempio, per un pod pod v4 con topologia 8x16x16, le dimensioni per topologia TPU sono 8, 16 e 16. È più efficiente utilizzare le parallele del modello a 8 o 16 vie (mappate a una delle dimensioni topologiche di TPU fisica). Un parallelismo del modello a 4 vie sarebbe non ottimale con questa topologia, poiché non è allineato con nessuna delle dimensioni della topologia TPU, ma sarebbe ottimale con una topologia 4x16x32 sullo stesso numero di chip.

Topologie piccolo v4

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

Nome (in base al conteggio di TensorCore) Numero di chip Topologia
versione 4-8 4 2x2x1
Versione 4-16 8 2x2x2
Versione 4-32 16 2x2x4
Versione 4-64 32 2x4x4

Topologie v4 di grandi dimensioni

Le sezioni TPU v4 sono disponibili in incrementi di 64 chip, con forme che sono multipli di 4 in tutte e tre le dimensioni. Anche le dimensioni devono essere in ordine crescente. Di seguito sono riportati diversi esempi. Alcune di queste topologie sono topologie "personalizzate" che possono essere avviate solo utilizzando l'API topology perché hanno lo stesso numero di chip di una topologia con nome più comunemente utilizzata.

Nome (in base al conteggio di TensorCore) Numero di chip Topologia
versione 4-128 64 4x4x4
Versione 4-256 128 4x4x8
Versione 4-512 256 4x8x8
Vedi API Topology 256 4x4x16
versione 4-1024 512 8x8x8
versione 4-1536 768 8x8x12
versione 4-2048 1024 8x8x16
Vedi API Topology 1024 4x16x16
versione 4-4096 2048 8x16x16
… … …

API Topology

Per creare sezioni di pod di Cloud TPU con topologia personalizzata, è possibile utilizzare l'API TPU gcloud come segue:

     $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=us-central2-b \
        --subnetwork=tpusubnet \
        --type=v4 \
        --topology=4x4x16 \
        --version=runtime-version

Forme di topologia

L'immagine seguente mostra tre possibili forme di topologia e le relative impostazioni dei flag topology associate.

immagine

Compatibilità tipo TPU

Puoi cambiare il tipo di TPU in 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 TPU con un numero maggiore o minore di chip o TensorCore, dovrai eseguire un'ottimizzazione e un'ottimizzazione significative. Per ulteriori informazioni, consulta la pagina relativa all'addestramento sui pod TPU.

Passaggi successivi