Panoramica di più sezioni di Cloud TPU

Cloud TPU Multislice è una tecnologia full stack che ottimizza le prestazioni che consente a un job di addestramento di utilizzare più sezioni TPU all'interno di un singolo pod in più pod con parallelismo dei dati semplice. Con i chip TPU v4, significa che i job di addestramento possono utilizzare più di 4096 chip in una singola esecuzione. Per addestramento per i job che richiedono meno di 4096 chip, una singola sezione può offrire delle prestazioni. Tuttavia, più sezioni più piccole sono più facilmente disponibili consentendo un tempo di avvio più rapido quando l'uso di più sezioni sezioni.

Quando viene eseguito il deployment in configurazioni multisezione, i chip TPU in ogni sezione comunicano tramite ICI (Inter-Chip-Interconnect). di chip TPU in diverse le sezioni comunicano trasferendo i dati alle CPU (host) che a loro volta trasmettono i dati attraverso la rete del data center (DCN).

Gli sviluppatori non devono scrivere codice per implementare la comunicazione DCN tra sezioni. Il compilatore XLA genera questo codice per te e si sovrappone alla comunicazione con calcolo per ottenere le massime prestazioni.

Concetti

Tipo di acceleratore

La forma di ogni sezione TPU che comprende una sezione multipla. Ciascuna la sezione in una richiesta multisezione è dello stesso tipo di acceleratore. Un acceleratore consiste in un tipo di TPU (v4 o v5e) seguito dal numero TensorCore. Ad esempio, v4-128 specifica una TPU v4 con 128 TensorCore.

Riparazione automatica

Quando una sezione rileva un evento di manutenzione, un prerilascio o un errore hardware, Cloud TPU creerà una nuova sezione. Nel raro caso in cui risorse insufficienti per creare una nuova sezione, la creazione non verrà completata finché l'hardware non sarà disponibile. Dopo aver creato la nuova sezione, tutte le altre le sezioni nell'ambiente Multislice verranno riavviate, quindi l'addestramento con uno script di avvio configurato correttamente, lo script di addestramento può riavviare automaticamente senza l'intervento dell'utente, oltre al caricamento e al ripristino dall'ultimo checkpoint.

Set di dati

I dati utilizzati da un modello per l'addestramento o l'inferenza.

Networking dei data center (DCN)

Una rete con latenza più elevata e velocità effettiva inferiore (se confrontata con ICI) connette le sezioni TPU in una configurazione con più sezioni.

Programmazione per le gang

Quando viene eseguito contemporaneamente il provisioning di tutte le sezioni TPU, garantisce il provisioning di tutte le sezioni o di nessuna è stato eseguito correttamente.

Organizzatore

Un host è un computer fisico che esegue le VM. Un host può eseguire al massimo quattro VM contemporaneamente. Ogni VM ha una TPU dedicata.

Inferenza

Carica un modello di machine learning preaddestrato su un host e fai previsioni in base dati.

Interchip Interconnect (ICI)

Collegamenti interni ad alta velocità e bassa latenza che collegano le TPU all'interno di un pod di TPU.

Più sezioni

Due o più sezioni di chip TPU che possono comunicare tramite DCN.

Nodo

Nel contesto di più sezioni, il nodo si riferisce a una singola sezione TPU. Ciascuna alla sezione TPU in una sezione multipla viene assegnato un ID nodo.

Pod

Una raccolta di chip TPU collegati da interfacce di rete ICI dedicate. R Il pod consente di distribuire il carico di elaborazione su più TPU.

Risorsa in coda (QR)

Una rappresentazione delle risorse TPU, utilizzata per accodare e gestire una richiesta per un ambiente TPU a sezione singola o multisezione.

Script di avvio

Uno script di avvio di Compute Engine standard che viene eseguito a ogni avvio o riavvio di una VM. Per Multisezione, viene specificato nella richiesta di creazione QR. Per ulteriori informazioni per informazioni sugli script di avvio di Cloud TPU, vedi Gestire le risorse TPU.

Sezione TPU

Una sottosezione logica di un pod di TPU composta da chip TPU. Tutti i chip in un le sezioni comunicano tra loro tramite la rete ICI.

VM TPU

Una macchina virtuale che esegue Linux e che ha accesso alle TPU sottostanti. Per TPU v4, ogni VM TPU ha accesso diretto a quattro chip. A volte chiamiamo TPU Invia una VM a un worker.

Tensor

Una struttura di dati utilizzata per rappresentare dati multidimensionali in una macchina modello di machine learning.

Tensor Processing Unit (TPU)

Chip di accelerazione ML sviluppato internamente da Google. Sono progettate per offrono computing veloce ed efficiente per le principali attività di machine learning come moltiplicazione matriciale.

Tipi di capacità di Cloud TPU

Le TPU possono essere create da diversi tipi di capacità (vedi Opzioni di utilizzo in Come funzionano i prezzi delle TPU) :

• Prenotazione: ha come target la quota prenotata. Per utilizzare la quota richiesta, devi avere un contratto di prenotazione con Google. Utilizza il flag --reserved durante la creazione le tue risorse.
• Spot: ha come target la quota prerilasciabile utilizzando le VM spot. Il tuo le risorse possono essere prerilasciate per fare spazio a richieste un job prioritario. Utilizza il flag --spot durante la creazione delle risorse.
• On demand: ha come target la quota on demand, che non richiede una prenotazione senza essere prerilasciata. La richiesta TPU verrà accodata a un modulo per la quota offerta da Cloud TPU, la disponibilità delle risorse non garantito. Selezionato per impostazione predefinita, nessun flag necessario.

Nota: Cloud TPU limita la dimensione della sezione per tutti i tipi di acceleratori. Per fetta limiti di dimensioni per le richieste on demand, consulta Dimensioni predefinite delle sezioni per i codici QR on demand.

Inizia

Se non hai mai utilizzato TPU, inizia installando Google Cloud CLI, e configurare l'ambiente Cloud TPU. Per utilizzare Più sezioni, le risorse TPU devono essere gestite come risorse in coda.

Se sei un utente TPU v4 esistente e hai una prenotazione, potresti dover eseguire la migrazione la prenotazione a un nuovo sistema di prenotazione. Per ulteriori informazioni, contatta il rappresentante dell'account Google Cloud.

Esempio introduttivo

Questo tutorial utilizza il codice del repository GitHub di MaxText. MaxText è un servizio ad alte prestazioni, arbitrariamente scalabile, open source e ben collaudato un LLM di base scritto in Python e Jax. MaxText è stato progettato per essere addestrato in modo efficiente con Cloud TPU.

Il codice in shardings.py è progettato per aiutarti a iniziare a sperimentare diversi parallelismo le opzioni di CPU e memoria disponibili. Ad esempio, parallelismo dei dati, parallelismo dei dati con sharding completo (FSDP), e il parallelismo tensore. Il codice scala da una sezione singola a una sezione multipla ambienti cloud-native.

Parallelismo ICI

ICI è l'interconnessione ad alta velocità che collega le TPU in un sezione. Lo sharding ICI corrisponde allo sharding all'interno di una sezione. shardings.py fornisce tre parametri di parallelismo di ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

I valori specificati per questi parametri determinano il numero di shard per ogni metodo di parallelismo.

Questi input devono essere vincolati in modo che ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism è uguale a al numero di chip nella sezione.

La tabella seguente mostra esempi di input utente per il parallelismo di ICI per i quattro disponibili nella versione 4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP a 4 vie	1	4	1
Parallelismo tensore a quattro vie	1	1	4
FSDP a 2 vie + parallelismo tensore a due vie	1	2	2

Tieni presente che nella maggior parte dei casi ici_data_parallelism deve essere lasciato 1 perché il valore La rete ICI è abbastanza veloce da preferire quasi sempre l'FSDP al parallelismo dei dati.

Questo esempio presuppone che tu abbia familiarità con l'esecuzione di codice su una singola sezione TPU come in Esegui un calcolo su una VM Cloud TPU utilizzando JAX. Questo esempio mostra come eseguire shardings.py su una singola sezione.

1. Configura l'ambiente:

   $ gcloud auth login
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone
   

2. Crea chiavi SSH per gcloud. Ti consigliamo di lasciare una password vuota (premi inseriscilo due volte dopo aver eseguito questo comando). Se ti viene richiesto google_compute_engine file esiste già. Sostituisci la versione esistente.

   $ ssh-keygen -f ~/.ssh/google_compute_engine
   

3. Esegui il provisioning delle TPU con il seguente comando:

   $ gcloud alpha compute tpus queued-resources \
   create your-qr-id \
   --accelerator-type your-accelerator-type \
   --runtime-version tpu-ubuntu2204-base \
   --node-id qr-id \
   [--reserved |--spot]
   

   Descrizioni flag di comando

   your-qr-id

   Una stringa definita dall'utente che identifica la richiesta QR.

   accelerator-type

   Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi Versioni TPU.

   runtime-version

   La [versione software di Cloud TPU](/tpu/docs/supported-tpu-configurations#tpu_software_versions).

   node-id

   L'ID delle risorse TPU che verranno create in risposta alla richiesta Richiesta QR.

   reserved

   Utilizza la quota prenotata durante la creazione delle sezioni.

   best-effort

   Utilizza la quota best effort durante la creazione delle sezioni [predefinita].

   Google Cloud CLI non supporta tutte le opzioni di creazione QR, come i tag. Per maggiori informazioni, consulta la sezione Creare codici QR.

4. Attendi che il QR sia nello stato ACTIVE, il che significa che i nodi worker sono nello stato READY. Una volta avviato il provisioning del codice QR, potrebbero essere necessari da uno a cinque minuti per il completamento, a seconda delle dimensioni del QR. Puoi controllare lo stato di una richiesta QR con il comando seguente:

   $ gcloud compute tpus queued-resources \
     list --filter=your-qr-id
   

5. Una sezione v4-8 ha una singola VM TPU. Connettiti alla VM TPU tramite SSH:

   $ gcloud compute tpus tpu-vm ssh your-qr-id
   

6. Clona MaxText (che include shardings.py) sulla VM TPU.

7. All'interno della directory del repository MaxText, esegui lo script di configurazione per installare JAX e altre dipendenze dalla sezione TPU. Lo script di configurazione richiede minuti per l'esecuzione.

   $ bash setup.sh
   

8. Esegui questo comando per eseguire shardings.py sulla sezione TPU.

   $ python3 pedagogical_examples/shardings.py \
     --ici_fsdp_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048
   

   Puoi vedere i risultati nei log. Le TPU dovrebbero raggiungere circa 260 TFLOP al secondo o un impressionante utilizzo di FLOP superiore al 90%! In questo caso, abbiamo è stato selezionato approssimativamente il batch massimo che rientra nell'impostazione Larghezza di banda (HBM).

9. Scopri pure altre strategie dello sharding invece di ICI. Ad esempio, puoi provare la seguente combinazione:

   $ python3 pedagogical_examples/shardings.py \
     --ici_tensor_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048
   

10. Al termine, elimina il QR e la sezione TPU. Devi eseguire questa pulizia passaggi dall'ambiente in cui hai configurato la sezione (prima esegui exit per uscire dalla sessione SSH). Il completamento dell'eliminazione richiederà da due a cinque minuti. e possono essere eseguiti in background con il flag facoltativo --async.

    Importante: i codici QR devono essere eliminati dopo l'uso. In caso contrario continueranno per consumare la quota QR che viene monitorata separatamente da quella di TPU.

    $ gcloud compute tpus queued-resources
      delete your-qr-id --force (--async)
    

Partizionamento orizzontale multisezione utilizzando il parallelismo della DCN

Lo script shardings.py accetta tre parametri che specificano il parallelismo della rete DCN: corrispondente al numero di shard di ogni tipo di parallelismo dei dati:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

I valori di questi parametri devono essere vincolati in modo da dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism è uguale a il numero di sezioni.

Come esempio per due sezioni, utilizza --dcn_data_parallelism = 2.

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	N. di fette
Parallelismo dei dati bidirezionale	2	1	1	2

dcn_tensor_parallelism deve essere sempre impostato su 1 perché la rete DCN è scadente adatto a questo sharding. Per i tipici carichi di lavoro LLM su chip v4, Anche dcn_fsdp_parallelism deve essere impostato su 1 e pertanto dcn_data_parallelism deve essere impostato sul numero di sezioni, ma questo è dall'applicazione.

Man mano che si aumenta il numero di sezioni (supponendo di mantenere le dimensioni della sezione e il batch per costante di sezione), aumenti la quantità di parallelismo dei dati.

Esecuzione di shardings.py in un ambiente con più sezioni

Puoi eseguire shardings.py in un ambiente con più sezioni utilizzando multihost_runner.py oppure eseguendo shardings.py su ogni VM TPU. In questo caso utilizziamo multihost_runner.py. I seguenti passaggi sono molto simili a quelli Per iniziare: esperimenti rapidi su più sezioni dal repository MaxText, ma in questo caso eseguiamo shardings.py anziché LLM più complesso in train.py.

Lo strumento multihost_runner.py è ottimizzato per esperimenti rapidi, ripetutamente riutilizzano le stesse TPU. Poiché lo script multihost_runner.py dipende da per le connessioni SSH di lunga durata, lo sconsigliamo per i job a lunga esecuzione. Se vuoi eseguire un job più lungo (ad esempio, ore o giorni), ti consigliamo di usa multihost_job.py.

In questo tutorial, utilizziamo il termine runner per indicare la macchina su cui esegui lo script multihost_runner.py. Usiamo il termine lavoratori per indicare VM TPU che compongono le tue sezioni. Puoi eseguire multihost_runner.py su una o su qualsiasi VM di Compute Engine nello stesso progetto delle tue sezioni. Esecuzione L'istruzione multihost_runner.py su un worker non è supportata.

multihost_runner.py si connette automaticamente ai worker TPU tramite SSH.

In questo esempio, eseguiamo shardings.py su due sezioni v4-16, per un totale di quattro e 16 chip TPU. Puoi modificare l'esempio in modo che venga eseguito su più TPU.

Configura l'ambiente

1. Clona MaxText sul tuo runner in una macchina virtuale.

2. Vai alla directory del repository.

3. Crea chiavi SSH per gcloud, consigliamo di lasciare una password vuota (premi inseriscilo due volte dopo aver eseguito questo comando). Se ti viene richiesto il file google_compute_engine esiste già, scegli di non mantenere una versione esistente.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. Aggiungi una variabile di ambiente per impostare il conteggio delle sezioni TPU su 2.

     $ export SLICE_COUNT=2
     

5. Crea un ambiente con più sezioni utilizzando queued-resources create.

   Nota: assicurati di disporre della quota corrispondente prima di selezionare --reserved oppure --best_effort o la quota on demand predefinita. Per informazioni sulla quota consulta i Criteri per le quote.

   Il seguente comando mostra come creare una TPU Multislice v4. Per utilizzare v5e, specifica un valore accelerator-type v5e (ad esempio v5litepod-16) e la Versione 5e runtime-version (v2-alpha-tpuv5-lite).

     $ gcloud alpha compute tpus queued-resources 
   
     create your-qr-id 
   
     --accelerator-type=your-accelerator-type 
   
     --runtime-version=tpu-vm-runtime-version 
   
     --node-count=node-count 
   
     --node-prefix=your-qr-id 
   
     [--reserved|--spot]
     

   Descrizioni flag di comando

   your-qr-id

   Una stringa definita dall'utente che identifica la richiesta QR.

   accelerator-type

   Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi Versioni TPU.

   runtime-version

   La versione software di Cloud TPU.

   node-count

   Il numero di sezioni da creare.

   node-prefix

   Il prefisso utilizzato per generare i nomi di ogni sezione. Viene aggiunto un numero al prefisso di ogni sezione. Ad esempio, se imposti node-prefix a mySlice, le sezioni sono denominate: mySlice-0, mySlice-1 e così via.

   reserved

   Utilizza la quota prenotata durante la creazione delle sezioni.

   best-effort

   Utilizza la quota best effort durante la creazione delle sezioni [predefinita].

6. Potrebbero essere necessari fino a cinque minuti dall'avvio del provisioning del codice QR a seconda delle dimensioni del QR. Attendi l'arrivo della risorsa in coda (QR) lo stato ACTIVE. Puoi controllare lo stato di una richiesta QR utilizzando il seguente comando:

   $ gcloud compute tpus queued-resources list \
   --filter=your-qr-id
   

   Dovrebbe essere generato un output simile al seguente:

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v4-16             ACTIVE
   ...
   

   Contatta il rappresentante del tuo account Google Cloud se lo stato QR è in WAITING_FOR_RESOURCES o PROVISIONING per più di 15 minuti.

7. Installare le dipendenze.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="bash setup.sh"
   

8. Esegui shardings.py su ogni worker utilizzando multihost_runner.py.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="python3 pedagogical_examples/shardings.py \
     --dcn_data_parallelism $SLICE_COUNT \
     --ici_fsdp_parallelism 8 \
     --batch_size 131072 \
     --embedding_dimension 2048"
   

   Vedrai circa 230 TFLOP al secondo di prestazioni nel log. .

9. Al termine, pulisci le TPU e il QR. L'eliminazione richiederà da due a cinque minuti per il completamento e può essere eseguito in background con --async flag.

Scalabilità di un carico di lavoro a più sezioni

Prima di eseguire il modello in un ambiente con più sezioni, le seguenti modifiche al codice:

• Utilizza jax.experimental.mesh_utils.create_hybrid_device_mesh invece di jax.experimental.mesh_utils.create_device_mesh durante la creazione della rete mesh.

Queste dovrebbero essere le uniche modifiche al codice necessarie per il passaggio a più sezioni. Per ottenere prestazioni elevate, la rete DCN deve essere mappata su dati paralleli, con sharding dei dati paralleli o assi paralleli della pipeline. Considerazioni sulle prestazioni e le strategie dello sharding sono illustrate più dettagliatamente nel Sharding con più sezioni per ottenere le massime prestazioni

Per verificare che il tuo codice possa accedere a tutti i dispositivi, puoi affermare che len(jax.devices()) è uguale al numero di chip nel tuo multisezione completamente gestito di Google Cloud. Ad esempio, se utilizzi quattro sezioni di v4-16, hai otto chip per sezione * 4 sezioni, quindi len(jax.devices()) dovrebbe restituire 32.

Scelta delle dimensioni delle sezioni per gli ambienti con più sezioni

Per aumentare la velocità lineare, aggiungi nuove sezioni delle stesse dimensioni di quelle esistenti sezione. Ad esempio, se utilizzi una sezione v4-512, Multisezione ottenere circa il doppio delle prestazioni aggiungendo una seconda sezione v4-512 e raddoppiando la dimensione del batch globale. Per ulteriori informazioni, vedi Sharding con più sezioni per ottenere le massime prestazioni

Esecuzione del job su più sezioni

Esistono tre diversi approcci per l'esecuzione di un carico di lavoro personalizzato Ambiente su più sezioni:

1. Utilizzando lo script di esecuzione della sperimentazione, multihost_runner.py
2. Utilizzando lo script runner di produzione, multihost_job.py
3. Adottare un approccio manuale

Script runner della sperimentazione

La multihost_runner.py lo script distribuisce il codice a un ambiente Multislice esistente ed esegue il tuo comando su ciascun host, copia i log e tiene traccia dell'errore di ciascun comando . Lo script multihost_runner.py è documentato in File README di MaxText.

Poiché multihost_runner.py mantiene connessioni SSH permanenti, è adatte a sperimentazioni di dimensioni modeste e relativamente brevi. Puoi adatta i passaggi nel tutorial multihost_runner.py al carico di lavoro e alla configurazione hardware.

Script runner di produzione

Per i job di produzione che richiedono resilienza contro guasti hardware e altre prerilascio, è preferibile integrare direttamente la risorsa Crea risorsa in coda API. Come esempio pratico, forniamo multihost_job.py, che attiva la chiamata API Created Queued Resource con l'avvio appropriato per eseguire l'addestramento e riprendere con il prerilascio. multihost_job.py sia documentato nel File README di MaxText.

Poiché multihost_job.py deve eseguire il provisioning delle risorse per ogni esecuzione, offrono un ciclo di iterazione rapido quanto multihost_runner.py.

Approccio manuale

Ti consigliamo di utilizzare o adattare multihost_runner.py oppure multihost_job.py per eseguire il tuo carico di lavoro personalizzato la configurazione multisezione. Tuttavia, se preferisci eseguire il provisioning e gestire l'ambiente direttamente con i comandi QR, vedi Gestire un ambiente con più sezioni.

Gestire un ambiente Multislice

Per eseguire manualmente il provisioning e la gestione dei codici QR senza utilizzare gli strumenti fornito nel repository MaxText, leggi le sezioni seguenti.

Crea codici QR

Imposta le seguenti variabili di ambiente prima di eseguire il provisioning della capacità:

  $ export your-qr-id=your-queued-resource-id
  $ export PROJECT=your-project-name
  $ export ZONE=us-central2-b
  $ export NETWORK_NAME=your-network-name
  $ export SUBNETWORK_NAME=your-subnetwork-name
  $ export RUNTIME_VERSION=tpu-ubuntu2204-base
  $ export ACCELERATOR_TYPE=v4-16
  $ export SLICE_COUNT=4
  $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
  $ gcloud config set project project-name
  $ gcloud config set compute/zone zone

Input	Descrizione
your-qr-id	L'ID del QR assegnato dall'utente.
PROGETTO	Nome progetto Google Cloud
ZONA	us-central2-b
NETWORK_NAME	Nome delle reti VPC.
SUBNETWORK_NAME	Nome della subnet nelle reti VPC
RUNTIME_VERSION	TPU-ubuntu2204-base
ACCELERATOR_TYPE	v4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2...	Tag utilizzati per identificare origini o destinazioni valide per i firewall di rete
SLICE_COUNT	Numero di sezioni. Limitato a un massimo di 256 sezioni.
STARTUP_SCRIPT	Se aggiunto alla richiesta di creazione, uno script di avvio può essere eseguito ogni volta che viene eseguito il provisioning o il riavvio di una sezione TPU e se la sezione TPU viene riparata o reimpostata.

Crea una richiesta QR utilizzando gcloud

$ gcloud alpha compute tpus queued-resources \
  create ${your-qr-id} \
  --project your-project-id \
  --zone your-zone \
  --node-count ${SLICE_COUNT} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --network ${NETWORK_NAME} \
  --subnetwork ${SUBNETWORK_NAME} \
  --tags ${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \ --metadata=startup-script='${STARTUP_SCRIPT}'
  [--reserved|--spot]
  

Descrizioni flag di comando

your-qr-id

Una stringa definita dall'utente che identifica la richiesta QR.

project

Una stringa definita dall'utente che identifica la richiesta QR.

zone

La zona Google Cloud in cui creare il QR.

node-count

Il numero di sezioni da creare.

accelerator-type

Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi Versioni TPU.

runtime-version

La versione software di Cloud TPU.

network

Il nome di una rete VPC a cui collegare la risorsa TPU.

subnetwork

Il nome di una subnet VPC a cui collegare la risorsa TPU.

reserved

Utilizza la quota prenotata durante la creazione delle sezioni.

spot

Utilizza la quota VM spot durante la creazione delle sezioni.

Assicurati di disporre della quota corrispondente prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quota, consulta i Criteri per le quote.

Crea una richiesta QR utilizzando curl

Crea un file denominato queued-resource-req.json e copia al suo interno il seguente JSON.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
    {
      "parent": "projects/your-project-number/locations/your-zone",
        "node": {
          "accelerator_type": "accelerator-type",
          "runtime_version": "tpu-vm-runtime-version",
          "network_config": {
            "network": "your-network-name",
            "subnetwork": "your-subnetwork-name",
            "enable_external_ips": true
          },
          "tags" : ["example-tag-1"]
          "metadata": {
            "startup-script": "your-startup-script"
          }
      },
      "multi_node_params": {
        "node_count": slice-count,
        "node_id_prefix": "your-queued-resource-id"
      }
    }
    ]
  }
}

Nota :sostituisci i seguenti valori:

• your-project-number: il numero del tuo progetto Google Cloud
• your-zone - La zona in cui vuoi creare il codice QR
• accelerator-type: la versione e la dimensione di una singola sezione
• tpu-vm-runtime-version: le versioni runtime VM TPU
• your-network-name - Facoltativo, una rete a cui verrà collegato il QR
• your-subnetwork-name - Facoltativo, una subnet a cui verrà collegato il QR
• example-tag-1 - Facoltativo, una stringa tag arbitraria
• your-startup-script: uno script di avvio che verrà eseguito quando viene allocato il QR
• slice-count: il numero di sezioni TPU nel tuo ambiente multisezione
• your-qr-id: l'ID fornito dall'utente per il codice QR

Per ulteriori informazioni, consulta la sezione sull'API REST Queued Resource. documentazione per tutte le opzioni disponibili.

Nota: anche se node_spec è un elenco, è supportata una sola voce node_spec.

Per utilizzare la capacità Spot, sostituisci:

"guaranteed": { "reserved": true } con "spot": {}

Rimuovi la riga per utilizzare la capacità on demand predefinita.

Invia la richiesta di creazione QR con il payload JSON:

  $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d @queuedresourcereq.json https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-qr-id

Nota :sostituisci i seguenti valori:

• your-project-id: l'ID del tuo progetto Google Cloud
• your-zone - La zona in cui vuoi creare il codice QR
• your-qr-id: l'ID fornito dall'utente per il codice QR

La risposta dovrebbe essere simile alla seguente:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Utilizza il valore GUID alla fine del valore della stringa per l'attributo name per ottenere informazioni sulla richiesta QR.

Recupera lo stato di un QR

Per ottenere lo stato della richiesta QR, utilizza il seguente comando:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/operations/operation-your-qr-guid

Nota: nel comando precedente sostituisci quanto segue:

• your-project-id: l'ID del tuo progetto Google Cloud
• your-zone: la zona in cui creare il QR
• your-qr-guid - Il GUID che segue name nell'output del Richiesta di creazione di codice QR.

La risposta di questo comando contiene lo stato dell'operazione:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qa-guid>,
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Se il QR viene creato correttamente ("done = true"), lo stato all'interno del Il campo response può essere WAITING_FOR_RESOURCES o FAILED. Se il codice QR è nello stato WAITING_FOR_RESOURCES, è stato in coda e inizierà il provisioning quando le risorse saranno sufficienti. Se il codice QR è nello stato FAILED, il motivo dell'errore sarà nell'output. Per ulteriori informazioni informazioni su altri possibili stati, vedi le Guida dell'utente sulle risorse in coda.

Terminata l'operazione, usa la descrizione dei codici QR. per monitorare le fasi del QR.

In uno scenario raro, potresti trovare il codice QR nello stato FAILED mentre alcune le sezioni sono ACTIVE. In questo caso, elimina le risorse create, e riprova tra qualche minuto o contattaci al team di Cloud TPU per risolvere il problema.

SSH e installazione delle dipendenze

Esegui codice JAX sulle pod di TPU TPU descrive come connetterti alle VM TPU tramite SSH in una singola sezione. A connettersi a tutte le VM TPU nel tuo ambiente Multislice tramite SSH installa le dipendenze, usa questo comando gcloud:

  $ gcloud compute tpus queued-resources ssh ${your-qr-id} \
    --zone your-zone \
    --node=all \
    --worker=all \
    --command="command-to-run"
    --batch-size=4

Questo comando gcloud invia il comando specificato a tutti i worker e i nodi in QR tramite SSH. Il comando viene suddiviso in gruppi di quattro e viene inviato contemporaneamente. Il batch successivo di comandi viene inviato quando il batch corrente completa l'esecuzione. In caso di errore con uno dei comandi, l'elaborazione e non vengono inviati ulteriori batch. Per ulteriori informazioni, consulta Riferimento API delle risorse in coda. Se il numero di sezioni che utilizzi supera quello dei thread del tuo computer locale (detto anche limite di batching) si imbatterà in un deadlock. Ad esempio, supponiamo che il limite per il batch sulla tua macchina locale sia 64. Se provi a eseguire un su più di 64 sezioni, ad esempio 100, il comando SSH le sezioni in batch. Lo script di addestramento verrà eseguito sul primo batch di 64 e attendi il completamento degli script prima di eseguirlo il batch rimanente di 36 sezioni. Tuttavia, il primo batch di 64 sezioni non può completato finché le 36 sezioni rimanenti non iniziano a eseguire lo script, causando un una situazione di stallo.

Per evitare questo scenario, puoi eseguire lo script di addestramento in background ogni VM aggiungendo una e commerciale (&) al comando di script specificato con il flag --command. Dopo aver avviato lo script di addestramento, sul primo batch di sezioni, il controllo tornerà immediatamente il comando SSH. Il comando SSH può quindi iniziare a eseguire lo script di addestramento il batch rimanente di 36 sezioni. Dovrai inviare le app stdout e stderr i flussi di dati in modo appropriato quando vengono eseguiti in background. Per aumentare parallelismo nello stesso QR, puoi selezionare sezioni specifiche usando lo --node .

Configurazione della rete

Assicurati che le sezioni TPU possano comunicare tra loro eseguendo questi passaggi. Installa JAX su ciascuna delle sezioni. Per ulteriori informazioni, vedi Esegui il codice JAX sulle pod di TPU TPU. Dichiarare len(jax.devices()) è uguale al numero di chip nel tuo multisezione completamente gestito di Google Cloud. Per farlo, su ogni sezione, esegui:

  $ python3 -c 'import jax; print(jax.devices())'

Se esegui questo codice su quattro sezioni delle versioni 4-16, ci sono otto chip per una sezione e quattro sezioni, per un totale di 32 chip (dispositivi) devono essere restituiti. entro il giorno jax.devices().

Elenca QR

Puoi visualizzare lo stato dei tuoi codici QR utilizzando il comando queued-resources list:

$ gcloud compute tpus queued-resources list

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v4-16             ACTIVE
...

Descrivi i QR

Per visualizzare la configurazione dettagliata e lo stato di un QR, utilizza il "Descrivi l'API QR". Puoi chiamare questa API utilizzando gcloud o curl.

Utilizzo di gcloud:

$ gcloud compute tpus queued-resources describe ${your-qr-id}
...state:
 state: ACTIVE
...

Utilizzo di curl:

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${your-qr-id}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
}

state rappresenta lo stato di un QR. Per ulteriori informazioni sui possibili degli stati dei QR, consulta la sezione Risorse in coda.

Avvia il job in un ambiente di cui è stato eseguito il provisioning

Puoi eseguire manualmente i carichi di lavoro connettendoti a tutti gli host in ogni sezione su SSH ed eseguendo questo comando su tutti gli host.

$ gcloud compute tpus tpu-vm ssh your-qr-id \
  --zone=your-zone \
  --worker=all \
  --node=all \
  --command="command-to-run"

Reimpostazione dei codici QR

È possibile utilizzare l'API ResetQueuedResource per reimpostare tutte le VM in un codice QR ACTIVE. La reimpostazione della VM cancella forzatamente la memoria la macchina e reimposta la VM allo stato iniziale. Tutti i dati archiviati localmente rimangono intatti e lo script di avvio verrà richiamato dopo un ripristino. La L'API ResetQueuedResource può essere utile quando vuoi riavviare tutte le TPU. Per Ad esempio, quando l'addestramento è bloccato e la reimpostazione di tutte le VM è più semplice del debug.

Le reimpostazioni di tutte le VM vengono eseguite in parallelo e un ResetQueuedResource il completamento dell'operazione richiede uno o due minuti. Per richiamare l'API, utilizza il seguente codice :

$ gcloud compute tpus queued-resources reset your-qr-id

Eliminazione dei codici QR in corso...

Per svincolare le risorse al termine della sessione di formazione, elimina le risorse in coda risorsa con il flag --force. L'eliminazione richiederà da due a cinque minuti completata e può essere eseguita in background con il flag facoltativo --async.

$ gcloud compute tpus queued-resources \
delete your-qr-id --force (--async)

Ripristino automatico da errori

In caso di interruzione, Multislice offre la riparazione della sezione interessata e il ripristino di tutte le sezioni in seguito. Le persone interessate la sezione viene sostituita con una nuova sezione e le altre sezioni altrimenti integre sono reimpostati. Se non è disponibile alcuna capacità disponibile per allocare una sezione sostitutiva, l'addestramento si interrompe.

Per riprendere automaticamente l'addestramento dopo un'interruzione, devi specificare script di avvio che verifica la presenza di e carica gli ultimi checkpoint salvati. Lo script di avvio viene eseguito automaticamente ogni volta che una sezione viene riallocata o viene reimpostata una VM. Specifica una startup nel payload JSON che invii all'API di richiesta QR di creazione.

Il seguente script di avvio (utilizzato in Creare QR) consente di ripristinare automaticamente gli errori e di riprendere l'addestramento checkpoint archiviati in un bucket Cloud Storage durante l'addestramento di MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Prima di provare questa opzione, clona il repository MaxText fuori.

Profilazione e debug

La profilazione è la stessa negli ambienti a sezione singola e a più sezioni. Per Per ulteriori informazioni, consulta la sezione Profilazione dei programmi JAX.

Ottimizza la formazione

Sharding con più sezioni per ottenere le massime prestazioni

Il raggiungimento delle massime prestazioni in ambienti multisettore richiede considerando come eseguire lo sharding sulle più sezioni. Di solito ci sono tre (parallelismo dei dati, parallelismo dei dati completamente partizionato e parallelismo delle pipeline). Sconsigliamo di partizionare le attivazioni per tutte le dimensioni del modello (a volte chiamato parallelismo tensore) perché richiede troppa larghezza di banda tra le sezioni. Per tutte queste strategie, puoi mantenere la stessa strategia dello sharding all'interno di una sezione. che ha funzionato per te in passato.

Consigliamo di iniziare con il parallelismo puro dei dati. Utilizzo di dati completamente suddivisi il parallelismo è utile per liberare memoria. Lo svantaggio è che la comunicazione tra le sezioni utilizza la rete DCN e rallenterà le tue carico di lavoro. Usa il parallelismo della pipeline solo quando necessario in base alle dimensioni del batch (come analizzato di seguito).

Quando utilizzare il parallelismo dei dati

Il parallelismo dei dati puro funziona bene nei casi in cui hai un carico di lavoro ma vorresti migliorarne le prestazioni scalando più sezioni.

Per ottenere una solida scalabilità su più sezioni, la quantità di tempo richiesta per eseguire la riduzione completa su DCN deve essere inferiore al tempo richiesto per eseguire un passaggio all'indietro. La rete DCN viene utilizzata per la comunicazione tra sezioni e è un fattore limitante della velocità effettiva del carico di lavoro.

Ogni chip TPU v4 ha un picco di 275 * 10¹² FLOPS al secondo.

Sono presenti quattro chip per host TPU e ogni host ha una larghezza di banda di rete massima di 50 Gbps.

Ciò significa che l'intensità aritmetica è pari a 4 * 275 * 10¹² FLOPS / 50 Gbps = 22000 FLOPS / bit.

Il modello utilizzerà da 32 a 64 bit di larghezza di banda DCN per ogni parametro per passaggio. Se utilizzi due sezioni, il modello utilizzerà 32 bit di larghezza di banda DCN. Se Se usi più di due sezioni, il compilatore eseguirà uno shuffling completo operativa e utilizzerai fino a 64 bit di larghezza di banda DCN per ciascun parametro passaggio. La quantità di FLOPS necessaria per ogni parametro varia a seconda del tuo model. Nello specifico, per i modelli linguistici basati su Transformer, il numero di FLOPS necessari per un passaggio in avanti e all'indietro sono circa 6 * B * P dove:

• B è la dimensione del batch in token
• P è il numero di parametri

Il numero di FLOPS per parametro è 6 * B e il numero di FLOPS per parametro durante il passaggio a ritroso è 4 * B.

Per garantire una scalabilità elevata su più sezioni, assicurati che l'ambiente operativo supera l'intensità aritmetica dell'hardware TPU. Per calcolare il intensità operativa, dividi il numero di FLOPS per parametro durante a ritroso della larghezza di banda della rete (in bit) per parametro per passaggio: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Pertanto, per un modello linguistico basato su Transformer, se utilizzi due sezioni: Operational intensity = 4 * B / 32

Se utilizzi più di due sezioni: Operational intensity = 4 * B/64

Questo suggerisce una dimensione minima del batch compresa tra 176.000 e 352.000 per Transformer basati su modelli linguistici. Poiché la rete DCN può rilasciare rapidamente pacchetti, la soluzione migliore per mantenere un margine di errore significativo, implementando solo il parallelismo dei dati se la dimensione del batch per pod è compresa tra 350.000 (due pod) e 700.000 (molti pod).

Per altre architetture del modello, dovrai stimare il runtime passare all'indietro per sezione (scadendola con un profiler o contando FLOPS). Quindi puoi confrontarlo con il tempo di esecuzione previsto per ridurre DCN e capire se il parallelismo dei dati ha senso per te.

Quando utilizzare il parallelismo dei dati con sharding completo (FSDP)

Il parallelismo dei dati con sharding completo (FSDP) combina il parallelismo dei dati (suddividendo tra nodi) con lo sharding dei pesi tra i nodi. Per ogni operazione in passi avanti e indietro, i pesi vengono tutti raccolti in modo che la sezione ha i pesi necessari. Invece di sincronizzare i gradienti con all-Reduce, i gradienti vengono ridotti man mano che vengono prodotti. In questo modo, ciascuna sezione ottiene i gradienti solo per le ponderazioni di cui è responsabile.

Analogamente al parallelismo dei dati, FSDP richiederà di scalare la dimensione globale del batch in modo lineare con il numero di sezioni. FSDP ridurrà la pressione di memoria aumenti il numero di sezioni. Questo perché il numero di pesi e lo stato di ottimizzatore per sezione diminuisce, ma allo stesso tempo aumenta traffico di rete e maggiore possibilità di bloccare a causa di un ritardo collettivo.

In pratica, la funzionalità FSDP tra le sezioni è la migliore se si aumenta il batch per archiviando più attivazioni per ridurre al minimo la rimaterializzazione durante passare all'indietro o aumentare il numero di parametri nella rete neurale.

Le operazioni all-gather e all-Reduce in FSDP funzionano in modo simile a quelle in DP. per consentirti di determinare se il carico di lavoro FSDP è limitato dalle prestazioni della rete DCN in allo stesso modo descritto nella sezione precedente.

Quando utilizzare il parallelismo della pipeline

Il parallelismo della pipeline diventa pertinente quando si raggiunge prestazioni elevate con altri strategie di parallelismo che richiedono una dimensione del batch globale maggiore di quella dimensione massima preferita per il batch. Il parallelismo della pipeline consente le sezioni costituito da una pipeline per "condividere" un batch. Tuttavia, il parallelismo della pipeline ha due svantaggi significativi:

1. Viene visualizzata la "bolla della pipeline" in cui i chip sono inattivi perché sono in attesa per i dati.
2. Richiede micro-batching che riduce la dimensione effettiva del batch, intensità aritmetica e, infine, modellare l'utilizzo di FLOP.

Il parallelismo della pipeline deve essere usato solo se le altre strategie di parallelismo non è richiesta una dimensione del batch globale troppo grande. Prima di provare il parallelismo della pipeline, vale la pena sperimentare per vedere concretamente se la convergenza per campione rallenta la dimensione del batch necessaria per ottenere un FSDP ad alte prestazioni. Il FSDP tende a raggiungere maggiore utilizzo FLOP del modello, ma se la convergenza per campione rallenta la dimensione del batch aumenta, il parallelismo della pipeline potrebbe comunque essere la scelta migliore. Più alta carichi di lavoro sono in grado di tollerare dimensioni batch sufficientemente grandi per non trarre vantaggio parallelismo della pipeline, ma il tuo carico di lavoro potrebbe essere diverso.

Se è necessario il parallelismo della pipeline, ti consigliamo di combinarlo con i dati il parallelismo o FSDP. Ciò ti consentirà di ridurre al minimo la profondità della pipeline aumentando la dimensione del batch per pipeline fino a quando la latenza DCN non diventa un fattore determinante nella velocità effettiva. Concretamente, se hai N sezioni, considera le pipeline di repliche di parallelismo dei dati e di profondità 2 e N/2, poi pipeline di profondità 4 e N/4 di repliche di parallelismo dei dati e così via, fino a quando le dimensioni del batch per pipeline non diventano al punto tale che i collettivi della DCN possano essere nascosti dietro l'aritmetica nella all'indietro. Ciò ridurrà al minimo il rallentamento introdotto dalla pipeline il parallelismo, permettendo allo stesso tempo di scalare oltre il limite globale della dimensione del batch.

Best practice per più sezioni

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente i batch da un set di dati da inserire nel feed model. Disporre di un caricatore dati asincrono efficiente che esegue lo sharding del batch è importante per evitare di esaurire le TPU di lavoro. Il caricatore di dati attuale in MaxText ha ogni carico dell'host un sottoinsieme uguale di esempi. Questa soluzione adeguato per il testo, ma richiede un controllo all'interno del modello. Inoltre, MaxText non offre ancora snapshot deterministici che consentirebbe all'iteratore di dati gli stessi dati prima e dopo il prerilascio.

Checkpoint

La libreria di checkpoint Orbax fornisce per il checkpoint di JAX PyTree nello spazio di archiviazione locale o in Google Cloud Storage. Forniamo un'integrazione di riferimento con checkpoint sincrono in MaxText a checkpointing.py.

Configurazioni supportate

Forme

Tutte le sezioni devono avere la stessa forma (ad esempio, lo stesso AcceleratorType). Le forme di sezioni eterogenee non sono supportate.

Orchestrazione

L'orchestrazione è supportata con GKE. Per ulteriori informazioni, consulta TPU in GKE.

Framework

Multislice supporta solo i carichi di lavoro JAX e PyTorch.

Parallelismo

Consigliamo agli utenti di testare più sezioni con parallelismo dei dati. Per scoprire di più sull'implementazione del parallelismo della pipeline con Multislice, contatta Rappresentante dell'account Google Cloud.

Assistenza e feedback

Tutti i feedback sono ben accetti. Contattaci per condividere un feedback o richiedere assistenza utilizzando il modulo di feedback o assistenza per Cloud TPU.