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 riscontra 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ò riavviarsi automaticamente senza l'intervento dell'utente, oltre che durante il caricamento e il 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 di 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.
- Prenotazione: ha come target la quota prenotata. Per utilizzare la quota richiesta, devi avere un
contratto di prenotazione con Google. Utilizza il flag
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.
Configura l'ambiente:
$ gcloud auth login $ gcloud config set project your-project-id $ gcloud config set compute/zone your-zone
Crea chiavi SSH per
gcloud
. Ti consigliamo di lasciare una password vuota (premi inseriscilo due volte dopo aver eseguito questo comando). Se ti viene richiestogoogle_compute_engine
file esiste già. Sostituisci la versione esistente.$ ssh-keygen -f ~/.ssh/google_compute_engine
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.
Attendi che il QR sia nello stato
ACTIVE
, il che significa che i nodi worker sono nello statoREADY
. 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
Una sezione v4-8 ha una singola VM TPU. Connettiti alla VM TPU tramite SSH:
$ gcloud compute tpus tpu-vm ssh your-qr-id
Clona MaxText (che include
shardings.py
) sulla VM TPU.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 alcune minuti per l'esecuzione.
$ bash setup.sh
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 FLOP di oltre il 90%! In questo caso, abbiamo è stato selezionato approssimativamente il batch massimo che rientra nell'intervallo Larghezza di banda (HBM).
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
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
.$ 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.
Con l'aumento del 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
multihost_runner.py
su un worker non è supportato.
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
Clona MaxText sul tuo runner in una macchina virtuale.
Vai alla directory del repository.
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 filegoogle_compute_engine
esiste già, scegli di non mantenere una versione esistente.$ ssh-keygen -f ~/.ssh/google_compute_engine
Aggiungi una variabile di ambiente per impostare il conteggio delle sezioni TPU su
2
.$ export SLICE_COUNT=2
Crea un ambiente con più sezioni utilizzando
queued-resources create
.Il seguente comando mostra come creare una TPU Multislice v4. Per utilizzare v5e, specifica un valore
accelerator-type
v5e (ad esempiov5litepod-16
) e la Versione 5eruntime-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
amySlice
, 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].
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
oPROVISIONING
per più di 15 minuti.Installare le dipendenze.
$ python3 multihost_runner.py \ --TPU_PREFIX=your-qr-id \ --COMMAND="bash setup.sh"
Esegui
shardings.py
su ogni worker utilizzandomultihost_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. .
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 in parallelo, con sharding dei dati e 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:
- Utilizzando lo script di esecuzione della sperimentazione,
multihost_runner.py
- Utilizzando lo script runner di produzione,
multihost_job.py
- 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
tramite Google Cloud CLI
o tramite l'API Compute Engine. 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" } } ] } }
- 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.
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
- 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
- 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 maggiori 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 la raccolta in batch sulla 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 la
"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 con più sezioni richiede considerando come eseguire lo sharding su 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 * 1012 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 * 1012 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 un modello di machine learning. 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:
- Viene visualizzata la "bolla della pipeline" in cui i chip sono inattivi perché sono in attesa per i dati.
- 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 un modello di machine learning. 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
nel seguente paese: 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.