Addestra un modello utilizzando TPU7x (Ironwood)

Questo documento descrive come eseguire il provisioning delle risorse TPU7x con Cluster Toolkit ed eseguire il deployment di un workload di addestramento utilizzando MaxText (JAX) o TorchTitan su TPU (PyTorch).

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

Per altri esempi ottimizzati per TPU7x, consulta i seguenti repository su GitHub:

Prima di iniziare

Prima di iniziare, assicurati di soddisfare i seguenti requisiti:

Poi completa i seguenti passaggi di configurazione:

  1. Imposta le variabili di ambiente per il progetto e il cluster:

    export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_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.
    • YOUR_CLUSTER_NAME: il nome del nuovo cluster.
    • YOUR_BUCKET_NAME: il nome del tuo bucket Cloud Storage per l'output di addestramento.

  2. Se non hai un bucket Cloud Storage esistente, creane uno:

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

Esegui il deployment del cluster TPU7x (Ironwood)

Esegui il deployment di un cluster TPU7x di Google Kubernetes Engine (GKE) seguendo le istruzioni riportate in Esegui il deployment di un cluster TPU 7x GKE. Assicurati che la topologia del pool di nodi del cluster corrisponda ai requisiti del workload di addestramento, ad esempio 4x4x8 per l'esempio MaxText o 2x2x4 per l'esempio TorchTitan.

Esegui il deployment di un workload di addestramento con MaxText

Le sezioni seguenti mostrano come creare un'immagine container e inviare un job di addestramento MaxText con Cluster Toolkit.

Crea e carica l'immagine Docker MaxText

Crea l'immagine Docker MaxText ed eseguine il push:

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

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

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

Definisci il comando di addestramento MaxText

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

La configurazione del modello MaxText 1B addestra un modello linguistico con circa 1 miliardo di parametri. Utilizza questo modello per testare 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_training \
        steps=30"

Invia il job di addestramento MaxText

Invia il job di addestramento utilizzando gcluster job submit. Per maggiori dettagli sulle opzioni di invio dei job, consulta la guida ai job di Cluster Toolkit:

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

I nomi dei workload devono essere univoci all'interno del cluster. Questo esempio aggiunge $(date +%H%M) al nome del workload.

Esegui il deployment di un workload di addestramento con TorchTitan

TorchTitan utilizza TorchTPU per eseguire l'addestramento PyTorch distribuito su Cloud TPU. Per saperne di più, consulta il file README di TorchTitan su TPU.

Passaggi successivi