Google utilizza la tecnologia AI per tradurre i contenuti nella tua lingua preferita. Le traduzioni generate dall'AI potrebbero contenere errori.

Addestra un modello utilizzando TPU v6e

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

TPU v6e, chiamata anche Trillium, è la sesta generazione di TPU 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. TPU v6e è ottimizzato per l'addestramento, l'ottimizzazione e la distribuzione di Transformer, modelli da testo a immagine 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:

Configura un progetto Google Cloud per le TPU
Installa Google Cloud CLI
Abilita l'API Compute Engine
Crea un account di servizio Compute Engine e concedi le autorizzazioni

Per saperne di più, consulta Configura l'ambiente Cloud TPU.

Verifica la quota e le autorizzazioni

Verifica che il tuo progetto abbia le seguenti quote:

Quota di acceleratori TPU di Compute Engine
Quota di VM di Compute Engine
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 (DP). Per saperne di più, consulta Opzioni di archiviazione per i dati di Cloud TPU.

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

Opzioni di provisioning

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

Compute Engine (Google Cloud CLI): puoi utilizzare lo strumento Google Cloud CLI per gestire le VM TPU direttamente come risorse di Compute Engine. Puoi eseguire il provisioning di singole VM TPU o utilizzare i gruppi di istanze gestite (MIG) per gestire gruppi di VM TPU. Per saperne di più, consulta la panoramica della creazione di TPU.
GKE: puoi utilizzare GKE per eseguire il provisioning e gestire le TPU come un pool di acceleratori per i tuoi workload di machine learning containerizzati. Per saperne di più, 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 di workload su GKE. È progettato per consentire ai professionisti del ML di eseguire il provisioning delle TPU ed eseguire job di addestramento senza bisogno di una conoscenza approfondita di Kubernetes. Per saperne di più, consulta il repository GitHub di XPK.

Esegui il provisioning delle TPU v6e con GKE e XPK

Se utilizzi GKE con TPU v6e, puoi utilizzare i comandi Kubernetes o XPK per eseguire il provisioning delle TPU e addestrare o gestire i 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 il throughput su più reti. Le sezioni seguenti forniscono i comandi per creare un cluster GKE con supporto per una singola NIC o per multi-NIC utilizzando XPK. Per la maggior parte dei workload a singola sezione, la NIC singola offre prestazioni sufficienti con una configurazione più semplice. Per i workload multislice e i workload che richiedono velocità di importazione dati elevate, utilizza multi-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 una configurazione più semplice. Per i workload multislice e i workload che richiedono velocità di importazione dati elevate, utilizza multi-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 tuo 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 acceleratore supportati per ogni versione di TPU, consulta Versioni di TPU.
- NUM_SLICES: il numero di sezioni TPU per il cluster. Ogni sezione ha il numero di chip specificato in ACCELERATOR_TYPE. Per un cluster a singola sezione, imposta NUM_SLICES su 1. Per un cluster multislice, specifica il numero di slice in base ai requisiti di scalabilità del 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 prestazioni ottimali,ti consigliamo di utilizzare una rete personalizzata con MTU di 8896. Per saperne di più, consulta Configura l'MTU.

Rete personalizzata

Per creare una rete personalizzata con MTU 8896 e utilizzarla per il tuo 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 una 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 utilizza 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 workload Multislice o altri workload che richiedono una larghezza di banda di rete elevata, ad esempio per l'importazione dati, puoi utilizzare multi-NIC per migliorare le prestazioni. Quando utilizzi più NIC, a ogni VM TPU vengono allocate interfacce di rete aggiuntive, ciascuna connessa a una rete VPC univoca, aumentando il throughput complessivo della rete. Per la maggior parte dei workload a singola sezione, una singola NIC offre prestazioni sufficienti con una configurazione ridotta.

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.
Configura 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 sezioni TPU per il cluster. Ogni sezione ha il numero di chip specificato in ACCELERATOR_TYPE. Per un cluster a singola sezione, imposta NUM_SLICES su 1. Per un cluster multislice, specifica il numero di slice in base ai requisiti di scalabilità del 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 una MTU di 8896 offre prestazioni migliori. Per saperne di più, consulta Configura 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 il NAT per la rete principale. Il seguente comando consente al traffico del tuo 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 pool di nodi 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 multi-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 della CPU nel cluster. Sebbene e2-standard-8 sia un'impostazione predefinita consigliata per molti workload, puoi selezionare un tipo di macchina diverso se il tuo workload ha requisiti di CPU o memoria specifici per attività come la pre-elaborazione dei dati.

Convalida la configurazione multi-NIC

Dopo aver creato un cluster con supporto multi-NIC, puoi verificare che entrambe le NIC siano in uso 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 nell'esempio seguente, o 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 che eth1 abbiano l'MTU impostata su 8896 controllando l'output del workload XPK nei log della console Google Cloud .

Configura 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 workload
GKE e XPK: README di XPK
Cloud TPU single-host utilizzando JAX: esegui un calcolo su una VM Cloud TPU utilizzando JAX
Cloud TPU multi-host utilizzando JAX: esegui il codice JAX su sezioni TPU
Cloud TPU con PyTorch: PyTorch su Cloud TPU

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

Utilizza SkyPilot

Puoi utilizzare Cloud TPU v6e con SkyPilot. SkyPilot è un framework open source che semplifica il processo di esecuzione, gestione e scalabilità dei workload di AI. Puoi aggiungere a SkyPilot informazioni su località e prezzi relative alla serie v6e. Per saperne di più, consulta l'esempio di SkyPilot TPU v6e.

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 notturne:
- 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 workload 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
git checkout maxtext-v0.2.0

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 comando seguente o un'immagine AI. Per saperne di più sulle immagini AI di JAX, consulta Immagini AI di JAX.

MaxText:

bash dependencies/scripts/docker_build_dependency_image.sh MODE=stable

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 gcloud CLI attiva:
```
gcloud config set project ${PROJECT_ID}
```
Se avvii il workload da una macchina su cui l'immagine non è stata 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 dependencies/scripts/docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
```

Esegui il workload 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 alle TPU e alle GPU Google Cloud per l'addestramento e l'inferenza.

JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m maxtext.trainers.pre_train.train src/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 open-weights sviluppati da Google DeepMind, basati sulla ricerca e sulla tecnologia di Gemini.

python3 -m maxtext.trainers.pre_train.train src/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 allo stato dell'arte sviluppato da Mistral AI, che utilizza un'architettura Mixture of Experts (MoE) sparsa.

python3 -m maxtext.trainers.pre_train.train src/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

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.

Devi installare un ramo Git specifico per eseguire MaxDiffusion, 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 tuo 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 acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
- NUM_SLICES: il numero di sezioni 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 o specificare il flag --docker-image e l'immagine che vuoi utilizzare.

Puoi scegliere di aggiungere i seguenti flag facoltativi:
- Puoi attivare il logging di debug includendo il flag --enable-debug-logs. Per saperne di più, consulta Esegui il debug di JAX su MaxText.
- Puoi creare un esperimento Vertex AI per caricare i dati su Vertex AI TensorBoard includendo il flag --use-vertex-tensorboard. Per ulteriori informazioni, consulta Monitora 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 tuo workload. Apri il link e fai clic sulla scheda Log per monitorare il workload in tempo reale.

Esegui 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
Abilita il logging dettagliato nei log del workload utilizzando il flag --enable-debug-logs quando crei il workload XPK

Monitora 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 stai utilizzando da 600 a 5000. Questo potrebbe non essere un problema per i workload di piccole dimensioni 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 anche eseguire questo comando sui cluster esistenti.
Crea l'esperimento Vertex AI durante l'esecuzione del workload XPK utilizzando il flag --use-vertex-tensorboard e il flag facoltativo --experiment-name. Per l'elenco completo dei passaggi, consulta Crea un esperimento Vertex AI per caricare i dati su Vertex AI TensorBoard.

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

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

Puoi anche trovare il link a Vertex AI TensorBoard 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 che hai 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 o allo stato del job. Questo comando potrebbe essere utile se hai inviato workload 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 un v6e-4, un 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 del batch globale	8	32	64
Throughput (esempi/sec)	115,9	438,4	492,3

Addestra un modello utilizzando PyTorch su Cloud TPU

Puoi utilizzare TorchTPU direttamente per addestrare il tuo modello o riutilizzare la libreria di riferimento TorchTitan. Per iniziare a utilizzare i workload v6e, consulta le ricette TorchTitan Trillium. Per saperne di più sull'ecosistema PyTorch su Cloud TPU, consulta PyTorch su Cloud TPU.

Addestra un modello utilizzando TPU v6e

Prima di iniziare

Verifica 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 multi-NIC

Configura JAX o PyTorch

Utilizza SkyPilot

Esempi di addestramento

Addestra MaxText e MaxDiffusion su Cloud TPU v6e

Crea l'immagine di base

Esegui il workload utilizzando XPK

MaxText

Gemma2

Mixtral 8x7b

MaxDiffusion

Descrizioni delle variabili di ambiente

Esegui il debug di JAX su MaxText

Monitora JAX su MaxText utilizzando Vertex AI

Elimina il workload XPK

Elimina il cluster

Risultati del benchmark di MaxDiffusion

Addestra un modello utilizzando PyTorch su Cloud TPU