Addestra un modello utilizzando TPU v6e

Questo documento ti guida nell'addestramento dei modelli su Cloud TPU v6e (chiamata anche Trillium), illustrando la configurazione dell'ambiente, l'ottimizzazione delle prestazioni ed esempi pratici di addestramento utilizzando JAX e PyTorch/XLA.

La TPU v6e, chiamata anche Trillium, è la TPU di sesta generazione di Google. Su tutte le piattaforme tecniche, come l'API e i log, e in tutto questo documento, Trillium verrà indicato come v6e. Con 256 chip per pod, l'architettura di TPU v6e condivide molte somiglianze con v5e. La TPU v6e è ottimizzata per l'addestramento, l'ottimizzazione e la pubblicazione di trasformatori, modelli di sintesi di immagini da testo e reti neurali convoluzionali (CNN). Per saperne di più sull'architettura e sulle configurazioni del sistema TPU v6e, consulta TPU v6e.

Per informazioni sull'esecuzione dell'inferenza su Cloud TPU v6e, consulta i seguenti tutorial:

Prima di iniziare

Prima di iniziare, devi:

Crea un Google Cloud account e un progetto con la fatturazione abilitata
Installa i componenti alpha di Google Cloud CLI
Abilita l'API Cloud TPU
Crea un service agent Cloud TPU
Crea un account di servizio Cloud TPU e concedi le autorizzazioni

Per saperne di più, vedi Configurare l'ambiente Cloud TPU.

Verificare la quota e le autorizzazioni

Verifica che il tuo progetto disponga delle seguenti quote:

Quota prerilasciabile o on demand di TPU v6e
Quota di indirizzi IP
Quota per Hyperdisk bilanciato e per qualsiasi altro tipo di disco che vuoi utilizzare

Nota: v6e supporta Hyperdisk bilanciato e Hyperdisk ML. v6e non supporta Persistent Disk (PD). Per maggiori informazioni, consulta Opzioni di archiviazione per i dati Cloud TPU.

Se utilizzi Google Kubernetes Engine (GKE) con XPK (Accelerated Processing Kit), devi disporre di autorizzazioni aggiuntive nella console Google Cloud . Per saperne di più, consulta Autorizzazioni necessarie nella console Google Cloud .

Opzioni di provisioning

Puoi eseguire il provisioning e gestire le TPU v6e utilizzando i seguenti metodi:

GKE: puoi utilizzare GKE per eseguire il provisioning e gestire le TPU come pool di acceleratori per i carichi di lavoro di machine learning containerizzati. Per maggiori informazioni, consulta Informazioni sulle TPU in GKE.
GKE e XPK: XPK è uno strumento a riga di comando che semplifica la creazione di cluster e l'esecuzione dei carichi di lavoro su GKE. È progettato per consentire ai professionisti di ML di eseguire il provisioning delle TPU ed eseguire job di addestramento senza richiedere una profonda conoscenza di Kubernetes. Per maggiori informazioni, consulta il repository GitHub di XPK.
Risorse Cloud TPU in coda: le risorse in coda ti consentono di richiedere capacità TPU di cui viene eseguito il provisioning quando diventa disponibile. È ideale per job batch e carichi di lavoro a tolleranza di errore che possono attendere in una coda. Puoi specificare un intervallo di tempo per la tua richiesta. Per saperne di più, vedi Gestire le risorse in coda.

Esegui il provisioning delle TPU v6e con GKE e XPK

Se utilizzi GKE con le TPU v6e, puoi utilizzare i comandi Kubernetes o XPK per eseguire il provisioning delle TPU e addestrare o pubblicare modelli. Per saperne di più sull'utilizzo di GKE con le TPU, consulta Informazioni sulle TPU in GKE.

Cloud TPU v6e supporta configurazioni di schede di interfaccia di rete (NIC) che consentono di scalare la velocità effettiva su più reti. Le sezioni seguenti forniscono comandi per creare un cluster GKE con supporto di una singola NIC o di più NIC utilizzando XPK. Per la maggior parte dei workload a una sola sezione, una sola NIC offre prestazioni sufficienti con meno configurazione. Per i workload Multislice e quelli che richiedono velocità di importazione dei dati elevate, utilizza più NIC.

Crea un cluster con supporto di una singola NIC utilizzando XPK

Per la maggior parte dei workload a singola sezione, una singola NIC offre prestazioni sufficienti con meno configurazione. Per i workload Multislice e quelli che richiedono velocità di importazione dei dati elevate, utilizza più NIC.

Le sezioni seguenti mostrano come creare un cluster GKE con supporto di una singola NIC utilizzando XPK.

Installa XPK e configura le variabili di ambiente

Installa XPK. Segui le istruzioni nel repository GitHub XPK.
Imposta le variabili di ambiente per il cluster:
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=1
```
Imposta le seguenti variabili di ambiente:
- CLUSTER_NAME: un nome per il cluster.
- ZONE: la zona in cui verrà creato il cluster TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone.
- PROJECT_ID: l'ID progetto Google Cloud .
- ACCELERATOR_TYPE: il tipo di TPU, chiamato anche tipo di acceleratore, specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Ad esempio, v6e-256. Per saperne di più sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
- NUM_SLICES: il numero di slice TPU per il cluster. Ogni fetta ha il numero di chip specificato in ACCELERATOR_TYPE. Per un cluster a singola fetta, imposta NUM_SLICES su 1. Per un cluster Multislice, specifica il numero di slice in base ai requisiti di scalabilità del tuo workload. Il numero totale di chip nel cluster è il numero di chip in ACCELERATOR_TYPE moltiplicato per NUM_SLICES.

Crea il cluster

Scegli una delle seguenti opzioni per creare il cluster. Per ottenere prestazioni ottimali,è consigliabile utilizzare una rete personalizzata con MTU di 8896. Per maggiori informazioni, vedi Configurare l'MTU.

Rete personalizzata

Per creare una rete personalizzata con MTU 8896 e utilizzarla per il cluster, segui questi passaggi:

Imposta le variabili di ambiente per i nomi della rete e del firewall:
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
Sostituisci quanto segue:
- NETWORK_NAME: un nome per la rete.
- FIREWALL_NAME: un nome per la regola firewall di rete.

Crea una rete personalizzata con un'unità MTU di 8896:

gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

Crea una regola firewall che consenta il traffico TCP, ICMP e UDP sulla tua rete:

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Imposta una variabile di ambiente per gli argomenti del cluster XPK in modo da utilizzare la rete che hai creato:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Crea il cluster XPK. Il seguente comando esegue il provisioning della capacità on demand:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

Per utilizzare la capacità riservata, sostituisci --on-demand con --reservation=RESERVATION_NAME. Per utilizzare le VM spot TPU, sostituisci --on-demand con --spot.

Rete predefinita

Se non hai bisogno di una rete con MTU elevata, puoi creare un cluster che utilizzi la rete VPC predefinita. Il seguente comando esegue il provisioning della capacità on demand:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

Per utilizzare la capacità riservata, sostituisci --on-demand con --reservation=RESERVATION_NAME. Per utilizzare le VM spot TPU, sostituisci --on-demand con --spot.

Crea un cluster con supporto multi-NIC utilizzando XPK

Per i carichi di lavoro multislice o altri carichi di lavoro che richiedono una larghezza di banda di rete elevata, ad esempio per l'importazione di dati, puoi utilizzare più NIC per migliorare le prestazioni. Quando utilizzi più NIC, a ogni VM TPU vengono allocate interfacce di rete aggiuntive, ognuna connessa a una rete VPC unica, aumentando il throughput di rete complessivo. Per la maggior parte dei workload a una sola sezione, una singola NIC offre prestazioni sufficienti con una configurazione minore.

Le sezioni seguenti mostrano come creare un cluster GKE con supporto multi-NIC utilizzando XPK.

Installa XPK e configura le variabili di ambiente

Installa XPK. Segui le istruzioni nel repository GitHub XPK.
Imposta le variabili di ambiente per il cluster e la rete principale:
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export REGION=REGION
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export FIREWALL_RULE_NAME_1=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME_1=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG_1=${CLUSTER_NAME}-natconfig-1-${ZONE}

export NETWORK_NAME_2=${CLUSTER_NAME}-mtu9k-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME_2=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME_2=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG_2=${CLUSTER_NAME}-natconfig-2-${ZONE}
```
Imposta le seguenti variabili di ambiente:
- CLUSTER_NAME: un nome per il cluster.
- REGION: la regione in cui verrà creato il cluster TPU.
- ZONE: la zona in cui verrà creato il cluster TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone.
- PROJECT_ID: l'ID progetto Google Cloud .
- ACCELERATOR_TYPE: il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Ad esempio, v6e-256. Per ulteriori informazioni sui tipi di acceleratore supportati per ogni versione di TPU, consulta Versioni di TPU.
- NUM_SLICES: il numero di slice TPU per il cluster. Ogni fetta ha il numero di chip specificato in ACCELERATOR_TYPE. Per un cluster a singola fetta, imposta NUM_SLICES su 1. Per un cluster Multislice, specifica il numero di slice in base ai requisiti di scalabilità del tuo workload. Il numero totale di chip nel cluster è il numero di chip in ACCELERATOR_TYPE moltiplicato per NUM_SLICES.

Crea le risorse di rete principali

Crea la rete principale con un'unità massima di trasmissione (MTU) di 8896:
```
gcloud compute networks create ${NETWORK_NAME_1} \
    --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}
```
L'utilizzo di una rete personalizzata con un'MTU di 8896 offre prestazioni migliori. Per maggiori informazioni, vedi Configurare l'MTU.

Crea la subnet principale:

gcloud compute networks subnets create ${SUBNET_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --range=10.11.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

Crea una regola firewall per la rete principale che consenta il traffico tcp, icmp e udp sulla rete principale:

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Crea un router Cloud per la rete principale:

gcloud compute routers create ${ROUTER_NAME_1} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_1} \
    --region=${REGION}

Configura NAT per la rete principale. Il seguente comando consente al traffico dal cluster di raggiungere internet:

gcloud compute routers nats create ${NAT_CONFIG_1} \
    --router=${ROUTER_NAME_1} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

Crea le risorse di rete secondarie

Crea la rete secondaria:

gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}

Crea una subnet per la rete secondaria:

gcloud compute networks subnets create ${SUBNET_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --range=10.10.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

Crea una regola firewall per consentire il traffico all'interno della nuova rete:

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --allow tcp,icmp,udp \
    --source-ranges 10.10.0.0/18 \
    --project=${PROJECT_ID}

Crea un router Cloud per la rete secondaria:

gcloud compute routers create ${ROUTER_NAME_2} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_2} \
    --region=${REGION}

Crea una configurazione NAT per il router Cloud:

gcloud compute routers nats create ${NAT_CONFIG_2} \
    --router=${ROUTER_NAME_2} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

Crea il cluster

Imposta una variabile di ambiente per gli argomenti del cluster e del node pool in modo da utilizzare le reti e le subnet che hai creato:

export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

Questi argomenti configurano il cluster in modo che utilizzi le due reti create per il supporto di più NIC.

Crea il cluster. Il seguente comando esegue il provisioning della capacità on demand:
```
xpk cluster create \
    --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE}  \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
    --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
    --create-vertex-tensorboard
```
Per utilizzare la capacità riservata, sostituisci --on-demand con --reservation=RESERVATION_NAME. Per utilizzare le VM spot TPU, sostituisci --on-demand con --spot.

Nota :il flag --cluster-cpu-machine-type specifica il tipo di macchina per le VM host CPU nel cluster. Sebbene e2-standard-8 sia un valore predefinito consigliato per molti workload, puoi selezionare un tipo di macchina diverso se il tuo workload ha requisiti specifici di CPU o memoria per attività come la preelaborazione dei dati.

Convalida la configurazione di più NIC

Dopo aver creato un cluster con supporto multi-NIC, puoi verificare che entrambe le NIC vengano utilizzate creando un workload XPK e aggiungendo il flag --command ifconfig.

Utilizza il seguente comando per visualizzare l'output del comando ifconfig nei log della consoleGoogle Cloud . Devi specificare il flag --base-docker-image maxtext_base_image per utilizzare l'immagine di base MaxText, come nel seguente esempio, oppure specificare il flag --docker-image e l'immagine che vuoi utilizzare.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --num-slices=${NUM_SLICES}  \
    --on-demand \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --command "ifconfig"

Se vuoi attivare i log di debug o utilizzare Vertex AI TensorBoard, aggiungi i seguenti argomenti facoltativi al comando:

--enable-debug-logs \
--use-vertex-tensorboard

Verifica che sia eth0 sia eth1 abbiano MTU impostato su 8896 controllando l'output del carico di lavoro XPK nei log della console Google Cloud .

Configurare JAX o PyTorch

Le seguenti risorse mostrano come configurare JAX o PyTorch sulla TPU, a seconda del metodo di provisioning e gestione che utilizzi:

GKE Autopilot: prepara l'applicazione TPU
GKE Standard: prepara i carichi di lavoro
GKE e XPK: README di XPK
Cloud TPU a singolo host utilizzando JAX: esegui un calcolo su una VM Cloud TPU utilizzando JAX
Cloud TPU multihost utilizzando JAX: esegui il codice JAX sulle sezioni TPU
Cloud TPU a singolo host utilizzando PyTorch: esegui un calcolo su una VM Cloud TPU utilizzando PyTorch
Cloud TPU multihost utilizzando PyTorch: esegui il codice PyTorch sugli slice TPU

Per configurare ed eseguire XPK con MaxText, consulta Esecuzione di MaxText su larga scala con XPK .

Migliorare le impostazioni TCP

Se hai eseguito il provisioning delle TPU v6e utilizzando le risorse in coda, puoi eseguire il comando seguente per migliorare le prestazioni di rete aumentando i limiti del buffer di ricezione TCP.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

Utilizzare SkyPilot

Puoi utilizzare Cloud TPU v6e con SkyPilot. SkyPilot è un framework open source che semplifica il processo di esecuzione, gestione e scalabilità dei carichi di lavoro AI. Puoi aggiungere a SkyPilot informazioni su prezzi e località correlate a v6e. Per ulteriori informazioni, consulta l'esempio di TPU v6e di SkyPilot.

Esempi di addestramento

Le sezioni seguenti forniscono esempi per l'addestramento di modelli MaxText, MaxDiffusion e PyTorch su Cloud TPU v6e.

Questi esempi sono stati testati con le seguenti versioni software:

Python 3.10 o versioni successive
Versioni software Nightly:
- JAX notturno 0.4.32.dev20240912
- LibTPU notturna 0.1.dev20240912+nightly
Versioni software stabili:
- JAX + JAX Lib v0.4.37

Addestra MaxText e MaxDiffusion su Cloud TPU v6e

Le sezioni seguenti illustrano il ciclo di vita dell'addestramento dei modelli MaxText e MaxDiffusion.

In generale, i passaggi di alto livello sono:

Crea l'immagine di base del workload.
Esegui il workload utilizzando XPK.
1. Crea il comando di addestramento per il workload.
2. Esegui il deployment del workload.
Segui il carico di lavoro e visualizza le metriche.
Elimina il workload XPK se non è necessario.
Elimina il cluster quando non è più necessario.

Crea l'immagine di base

Installa MaxText o MaxDiffusion e crea l'immagine Docker:

Clona il repository che vuoi utilizzare e passa alla directory del repository:

MaxText:

git clone https://github.com/google/maxtext.git && cd maxtext

MaxDiffusion:

git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103

Configura Docker in modo che utilizzi Google Cloud CLI:
```
gcloud auth configure-docker
```

Crea l'immagine Docker utilizzando il seguente comando o un'immagine JAX AI. Per saperne di più sulle immagini AI JAX, consulta Immagini AI JAX.

MaxText:

bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35

MaxDiffusion:

bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest

Imposta l'ID progetto nella configurazione dell'interfaccia a riga di comando gcloud attiva:
```
gcloud config set project ${PROJECT_ID}
```
Se avvii il workload da una macchina in cui l'immagine non è creata localmente, carica l'immagine.

Nota: assicurati di utilizzare bash versione 4.2 o successive.
1. Imposta la variabile di ambiente CLOUD_IMAGE_NAME:
```
export CLOUD_IMAGE_NAME=${USER}_runner
```
2. Carica l'immagine:
```
bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
```

Esegui il carico di lavoro utilizzando XPK

Imposta le seguenti variabili di ambiente se non utilizzi i valori predefiniti impostati da MaxText o MaxDiffusion:

export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
export PER_DEVICE_BATCH_SIZE=2
export NUM_STEPS=30
export MAX_TARGET_LENGTH=8192

Crea lo script del modello. Questo script verrà copiato come comando di addestramento in un passaggio successivo.

Non eseguire ancora lo script del modello.

MaxText

MaxText è un LLM open source ad alte prestazioni e altamente scalabile scritto in Python e JAX puri e destinato a TPU e GPU per l'addestramento e l'inferenza. Google Cloud

JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m MaxText.train MaxText/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

Gemma2

Gemma è una famiglia di LLM con pesi aperti sviluppati da Google DeepMind, basati sulla ricerca e sulla tecnologia Gemini.

python3 -m MaxText.train MaxText/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

Mixtral 8x7b

Mixtral è un modello di AI all'avanguardia sviluppato da Mistral AI, che utilizza un'architettura di tipo sparse mixture-of-experts (MoE).

python3 -m MaxText.train MaxText/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

Llama3-8b

Llama è una famiglia di LLM con pesi aperti sviluppati da Meta.

Per un esempio di come eseguire Llama3 su PyTorch, consulta i modelli torch_xla nel repository torchprime.

MaxDiffusion

MaxDiffusion è una raccolta di implementazioni di riferimento di vari modelli di diffusione latente scritti in Python e JAX puri che vengono eseguiti su dispositivi XLA, tra cui Cloud TPU e GPU. Stable Diffusion è un modello latente da testo a immagine che genera immagini fotorealistiche da qualsiasi input di testo.

Per eseguire MaxDiffusion, devi installare un ramo Git specifico, come mostrato nel seguente script di addestramento.

git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

Esporta le seguenti variabili:
```
export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT
```
Descrizioni delle variabili di ambiente
- CLUSTER_NAME: il nome del cluster.
- ACCELERATOR_TYPE: il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per saperne di più sui tipi di acceleratore supportati per ogni versione di TPU, consulta Versioni di TPU.
- NUM_SLICES: il numero di slice TPU.
- YOUR_MODEL_SCRIPT: lo script del modello da eseguire come comando di addestramento.
Esegui il modello utilizzando lo script creato nel passaggio precedente. Devi specificare il flag --base-docker-image per utilizzare l'immagine di base MaxText oppure specificare il flag --docker-image e l'immagine che vuoi utilizzare.

Puoi scegliere di aggiungere i seguenti flag facoltativi:
- Puoi attivare la registrazione di debug includendo il flag --enable-debug-logs. Per saperne di più, consulta Eseguire il debug di JAX su MaxText.
- Puoi creare un esperimento Vertex AI per caricare i dati in Vertex AI TensorBoard includendo il flag --use-vertex-tensorboard. Per saperne di più, consulta Monitorare JAX su MaxText utilizzando Vertex AI.
```
xpk workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"
```
L'output include un link per seguire il carico di lavoro. Apri il link e fai clic sulla scheda Log per monitorare il carico di lavoro in tempo reale.

Eseguire il debug di JAX su MaxText

Utilizza i comandi XPK supplementari per diagnosticare il motivo per cui il cluster o il workload non è in esecuzione:

Elenco dei workload XPK
XPK inspector
Attiva la registrazione dettagliata nei log del workload utilizzando il flag --enable-debug-logs quando crei il workload XPK

Monitorare JAX su MaxText utilizzando Vertex AI

Per utilizzare TensorBoard, il tuo account utente Google Cloud deve disporre del ruolo aiplatform.user. Esegui questo comando per concedere il ruolo:

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

Visualizza i dati scalari e di profilo tramite TensorBoard gestito da Vertex AI.

Aumenta le richieste di gestione delle risorse (CRUD) per la zona che utilizzi da 600 a 5000. Questo potrebbe non essere un problema per i piccoli carichi di lavoro che utilizzano meno di 16 VM.

Installa le dipendenze come cloud-accelerator-diagnostics per Vertex AI:

# xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

Crea il cluster utilizzando il flag --create-vertex-tensorboard, come documentato in Crea Vertex AI TensorBoard. Puoi eseguire questo comando anche sui cluster esistenti.
Crea l'esperimento Vertex AI durante l'esecuzione del carico di lavoro XPK utilizzando il flag --use-vertex-tensorboard e il flag facoltativo --experiment-name. Per l'elenco completo dei passaggi, vedi Crea Vertex AI Experiment per caricare i dati in Vertex AI TensorBoard.

I log includono un link a Vertex AI TensorBoard, simile al seguente:

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

Puoi trovare il link a Vertex AI TensorBoard anche nella console Google Cloud . Vai a Vertex AI Experiments nella console Google Cloud . Seleziona la regione appropriata dal menu a discesa.

La directory TensorBoard viene scritta anche nel bucket Cloud Storage specificato con ${BASE_OUTPUT_DIR}.

Elimina il workload XPK

Utilizza il comando xpk workload delete per eliminare uno o più workload in base al prefisso del job o allo stato del job. Questo comando potrebbe essere utile se hai inviato carichi di lavoro XPK che non devono più essere eseguiti o se hai job bloccati nella coda.

Elimina il cluster

Utilizza il comando xpk cluster delete per eliminare il cluster:

xpk cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

Risultati del benchmark di MaxDiffusion

Abbiamo eseguito lo script di addestramento per MaxDiffusion su v6e-4, v6e-16 e due v6e-16. La tabella seguente mostra i throughput misurati.

	v6e-4	v6e-16	Due v6e-16
Passaggi di addestramento	0,069	0,073	0,13
Dimensione batch globale	8	32	64
Throughput (esempi/sec)	115,9	438,4	492,3

Addestramento di modelli Llama utilizzando PyTorch/XLA su Cloud TPU v6e

Questa sezione descrive come addestrare i modelli Llama utilizzando PyTorch/XLA su Cloud TPU v6e utilizzando il set di dati WikiText.

Accedere a Hugging Face e al modello Llama 3

Per questo esempio è necessario un token di accesso utente Hugging Face. Per informazioni sulla creazione di token di accesso utente, consulta la documentazione di Hugging Face sui token di accesso utente.

Devi anche disporre dell'autorizzazione per accedere al modello Llama-3-8B su Hugging Face. Per ottenere l'accesso, vai al modello Meta-Llama-3-8B su HuggingFace e richiedi l'accesso.

Crea una VM Cloud TPU

Crea una Cloud TPU v6e con 8 chip per questo esempio.

Imposta le variabili di ambiente:
```
export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e
```
Descrizioni delle variabili di ambiente
- PROJECT_ID: il tuo ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
- TPU_NAME: il nome della TPU.
- ZONE: La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
- ACCELERATOR_TYPE: Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di acceleratore supportati per ogni versione di TPU, consulta Versioni di TPU.
- RUNTIME_VERSION: la versione software di Cloud TPU.

Crea una VM Cloud TPU:

gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

Installazione

Installa il fork pytorch-tpu/transformers di Hugging Face Transformers e le dipendenze. Questo esempio è stato testato con le seguenti versioni delle dipendenze:

torch: compatibile con la versione 2.5.0
torch_xla[tpu]: compatibile con la versione 2.5.0
jax: 0.4.33
jaxlib: 0.4.33

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

Configurare i file di configurazione del modello

Il comando di addestramento nella sezione successiva, Esegui il modello, utilizza due file di configurazione JSON per definire i parametri del modello e la configurazione Fully Sharded Data Parallel (FSDP). Lo sharding FSDP ti consente di utilizzare un batch più grande durante l'addestramento suddividendo i pesi del modello su più TPU. Quando esegui l'addestramento con modelli più piccoli, potrebbe essere sufficiente utilizzare il parallelismo dei dati e replicare i pesi su ogni dispositivo. Per saperne di più su come partizionare i tensori tra i dispositivi in PyTorch/XLA, consulta la Guida per l'utente di PyTorch/XLA SPMD.

Crea il file di configurazione dei parametri del modello. Di seguito è riportata la configurazione dei parametri del modello per Llama-3-8B. Per altri modelli, trova il file di configurazione su Hugging Face. Ad esempio, consulta la configurazione Llama-2-7B.

cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

Crea il file di configurazione FSDP:

cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

Per ulteriori informazioni su FSDP, consulta Parallelismo dei dati completamente partizionati utilizzando SPMD .

Carica i file di configurazione nelle tue VM Cloud TPU utilizzando il seguente comando:

gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

Esegui il modello

Utilizzando i file di configurazione creati nella sezione precedente, esegui lo script run_clm.py per addestrare il modello Llama-3-8B sul set di dati WikiText. L'esecuzione dello script di addestramento richiede circa 10 minuti su una Cloud TPU v6e-8.

Accedi a Hugging Face sulla tua Cloud TPU utilizzando il seguente comando:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

Esegui l'addestramento del modello:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

Risoluzione dei problemi di PyTorch/XLA

Se hai impostato le variabili facoltative per il debug nella sezione precedente, il profilo per il modello verrà archiviato nella posizione specificata dalla variabile PROFILE_LOGDIR. Puoi estrarre il file xplane.pb archiviato in questa posizione e utilizzare tensorboard per visualizzare i profili nel browser seguendo le istruzioni di TensorBoard.

Se PyTorch/XLA non funziona come previsto, consulta la Guida alla risoluzione dei problemi, che contiene suggerimenti per il debug, la profilazione e l'ottimizzazione del modello.

Addestra un modello utilizzando TPU v6e

Prima di iniziare

Verificare la quota e le autorizzazioni

Opzioni di provisioning

Esegui il provisioning delle TPU v6e con GKE e XPK

Crea un cluster con supporto di una singola NIC utilizzando XPK

Installa XPK e configura le variabili di ambiente

Crea il cluster

Rete personalizzata

Rete predefinita

Crea un cluster con supporto multi-NIC utilizzando XPK

Installa XPK e configura le variabili di ambiente

Crea le risorse di rete principali

Crea le risorse di rete secondarie

Crea il cluster

Convalida la configurazione di più NIC

Configurare JAX o PyTorch

Migliorare le impostazioni TCP

Utilizzare SkyPilot

Esempi di addestramento

Addestra MaxText e MaxDiffusion su Cloud TPU v6e

Crea l'immagine di base

Esegui il carico di lavoro utilizzando XPK

MaxText

Gemma2

Mixtral 8x7b

Llama3-8b

MaxDiffusion

Descrizioni delle variabili di ambiente

Eseguire il debug di JAX su MaxText

Monitorare JAX su MaxText utilizzando Vertex AI

Elimina il workload XPK

Elimina il cluster

Risultati del benchmark di MaxDiffusion

Addestramento di modelli Llama utilizzando PyTorch/XLA su Cloud TPU v6e

Accedere a Hugging Face e al modello Llama 3

Crea una VM Cloud TPU

Descrizioni delle variabili di ambiente

Installazione

Configurare i file di configurazione del modello

Esegui il modello

Risoluzione dei problemi di PyTorch/XLA