Questa pagina è stata tradotta dall'API Cloud Translation.

Addestra un modello utilizzando TPU7x (Ironwood)

Questo documento descrive come eseguire il provisioning delle risorse TPU7x e fornisce un esempio di deployment di un carico di lavoro di addestramento utilizzando MaxText e XPK.

TPU7x è la prima release della famiglia Ironwood, Google Cloudla TPU di settima generazione. La generazione Ironwood è progettata per l'addestramento e l'inferenza dell'AI su larga scala. Per saperne di più, vedi TPU7x.

Per altri esempi ottimizzati per TPU7x, consulta Training Recipes for Ironwood TPU su GitHub.

Provisioning delle TPU

Puoi eseguire il provisioning e gestire TPU7x 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. Utilizza Google Cloud CLI per creare manualmente l'istanza del cluster GKE per una personalizzazione precisa o per l'espansione degli ambienti GKE di produzione esistenti. 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 i professionisti del machine learning per eseguire il provisioning delle TPU ed eseguire job di addestramento senza bisogno di competenze approfondite su Kubernetes. Utilizza XPK per creare rapidamente cluster GKE ed eseguire carichi di lavoro per la prova di fattibilità e i test. Per maggiori informazioni, consulta il repository GitHub di XPK.
GKE e TPU Cluster Director: TPU Cluster Director è disponibile tramite una prenotazione in modalità All Capacity, che ti offre l'accesso completo a tutta la capacità prenotata (senza riserve) e la piena visibilità della topologia hardware TPU, dello stato di utilizzo e dello stato di integrità. Per ulteriori informazioni, vedi la panoramica della modalità Tutta la capacità.

Esegui il deployment di un carico di lavoro di addestramento con MaxText e XPK

Utilizza Accelerated Processing Kit (XPK) per creare cluster GKE per la prova di fattibilità e i test. XPK è uno strumento a riga di comando progettato per semplificare il provisioning, la gestione e l'esecuzione dei workload di machine learning.

Le sezioni seguenti mostrano come eseguire il deployment di un carico di lavoro di addestramento utilizzando MaxText e XPK.

Prima di iniziare

Prima di iniziare, completa i seguenti passaggi:

Assicurati di avere un progetto Google Cloud con la fatturazione abilitata.
Accedere a TPU7x. Per ulteriori informazioni, contatta il team dedicato al tuo account.
Assicurati che l'account che utilizzi con XPK disponga dei ruoli elencati nel repository GitHub XPK.

Installa XPK e le dipendenze

Installa XPK. Segui le istruzioni nel repository GitHub XPK.
Installa Docker seguendo le istruzioni fornite dall'amministratore o segui le istruzioni di installazione ufficiali. Una volta installato, esegui i seguenti comandi per configurare Docker e testare l'installazione:
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
Imposta le seguenti variabili di ambiente:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
Sostituisci quanto segue:
- YOUR_PROJECT_ID: l'ID progetto Google Cloud .
- YOUR_ZONE: la zona in cui creare il cluster. Per l'anteprima, è supportato solo us-central1-c.
- YOUR_CLUSTER_NAME: il nome del nuovo cluster.
- YOUR_ACCELERATOR_TYPE: la versione e la topologia della TPU. Ad esempio, tpu7x-4x4x8. Per un elenco delle topologie supportate, consulta Configurazioni supportate.
- YOUR_RESERVATION_NAME: il nome della prenotazione. Per le prenotazioni condivise, utilizza projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: il nome del bucket Cloud Storage, che sarà la directory di output per l'addestramento del modello.

Se non hai un bucket Cloud Storage esistente, creane uno utilizzando il seguente comando:

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

Crea un cluster a singola NIC e singola sezione

Segui le istruzioni nella sezione Configurare MTU per ottimizzare la configurazione della rete.
Compila la variabile ${CLUSTER_ARGUMENTS}, che utilizzerai nel comando xpk cluster create:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
Crea il cluster GKE con i node pool TPU7x utilizzando il comando xpk cluster create:
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Se imposti il flag --cluster-cpu-machine-type su n1-standard-8 (o un valore superiore), il pool di nodi predefinito dispone di CPU sufficiente per i pod di sistema, ad esempio il webhook JobSet, evitando errori. Per impostazione predefinita, XPK utilizza e2-standard-16. Alcune zone supportano solo tipi di CPU specifici, quindi potresti dover passare dai tipi n1, n2 e e2. In caso contrario, potresti riscontrare errori di quota.

Aggiungi un'esclusione dalla manutenzione per impedire gli upgrade per il cluster:

gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"

Sostituisci quanto segue:

EXCLUSION_START_TIME: l'ora di inizio selezionata per l'esclusione della manutenzione nel formato YYYY-MM-DDTHH:MM:SSZ.
EXCLUSION_END_TIME: l'ora di fine selezionata per l'esclusione della manutenzione nel formato YYYY-MM-DDTHH:MM:SSZ.

Crea o carica l'immagine Docker MaxText

Puoi creare un'immagine Docker localmente utilizzando gli script forniti da MaxText o utilizzare un'immagine predefinita.

Compila in locale

I seguenti comandi copiano la directory locale nel container:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

Dopo l'esecuzione corretta dei comandi, dovresti visualizzare un'immagine denominata maxtext_base_image creata localmente. Puoi utilizzare l'immagine locale direttamente nel comando del workload xpk.

Carica immagine (facoltativo)

Dopo aver creato l'immagine Docker localmente seguendo le istruzioni della sezione precedente, puoi caricarla nel registro utilizzando il seguente comando:

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Dopo l'esecuzione corretta di questo comando, dovresti visualizzare l'immagine MaxText in gcr.io con il nome gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Definisci il comando di addestramento MaxText

Prepara il comando per eseguire lo script di addestramento all'interno del container Docker.

Il modello MaxText 1B è una configurazione all'interno del framework MaxText progettata per l'addestramento di un modello linguistico con circa 1 miliardo di parametri. Utilizza questo modello per sperimentare con scale di chip di piccole dimensioni. Le prestazioni non sono ottimizzate.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

Esegui il deployment del carico di lavoro di addestramento

Esegui il comando xpk workload create per eseguire il deployment del job di addestramento. 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 includere il flag --enable-debug-logs per attivare la registrazione di debug.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

I nomi dei workload devono essere univoci all'interno del cluster. In questo esempio, $(date +%H%M) viene aggiunto al nome del workload per garantire l'unicità.