Addestra un modello utilizzando TPU7x (Ironwood)

Questo documento descrive come eseguire il provisioning delle risorse TPU7x e fornisce un esempio di deployment di un workload di addestramento utilizzando MaxText e Cluster Toolkit.

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ù, consulta TPU7x.

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

Esegui il deployment di un carico di lavoro di addestramento con MaxText e Cluster Toolkit

Utilizza Cluster Toolkit per creare cluster GKE pronti per la produzione ed eseguire carichi di lavoro di addestramento.

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

Prima di iniziare

Prima di iniziare, completa i seguenti passaggi:

  1. Imposta le seguenti variabili di ambiente:

    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 bucket Cloud Storage, che sarà la directory di output per l'addestramento del modello.

  2. 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

Esegui il deployment del cluster TPU7x (Ironwood)

Esegui il deployment di un cluster GKE TPU7x seguendo le istruzioni riportate in Esegui il deployment di un cluster GKE TPU 7x. Assicurati che la topologia del pool di nodi del cluster corrisponda ai requisiti del workload di addestramento (ad esempio 4x4x8 per il workload MaxText di esempio riportato di seguito).

Crea e carica l'immagine Docker MaxText

Crea ed esegui il push dell'immagine Docker MaxText:

# 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.

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_training \
        steps=30"

Esegui il deployment del workload di addestramento

Invia il job di addestramento utilizzando gcluster job submit. Per informazioni dettagliate 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. In questo esempio, $(date +%H%M) viene aggiunto al nome del workload per garantire l'unicità.

Passaggi successivi