Addestrare 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 di Google. 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.

Eseguire il 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 un pool di acceleratori per i carichi di lavoro di machine learning containerizzati. Utilizza la 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 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 carichi di lavoro su GKE. È progettato per i professionisti del machine learning per eseguire il provisioning delle TPU ed eseguire job di addestramento senza richiedere una conoscenza approfondita di Kubernetes. Utilizza XPK per creare rapidamente cluster GKE ed eseguire carichi di lavoro per la prova di fattibilità e i test. Per saperne di più, 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à riservata (senza riserve) e la visibilità completa della topologia hardware TPU, dello stato di utilizzo e dello stato di integrità. Per saperne di più, consulta la panoramica della modalità All Capacity overview.

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

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 Google Cloud progetto con la fatturazione abilitata.
• Ottieni l'accesso a TPU7x. Per saperne di più, contatta il team dedicato all'account.
• Assicurati che l'account che utilizzi con XPK abbia i ruoli elencati nel repository GitHub di XPK.

Installare XPK e le dipendenze

1. Installa XPK. Segui le istruzioni nel repository GitHub di XPK.

2. Installa Docker utilizzando 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
   

3. 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 BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"

   Sostituisci quanto segue:

   • YOUR_PROJECT_ID: l' Google Cloud ID progetto.
   • YOUR_ZONE: la zona in cui creare il cluster.
   • 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_BUCKET_NAME: il nome del bucket Cloud Storage, che sarà la directory di output per l'addestramento del modello.

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

Creare un cluster a singola NIC e singola slice

Scegli una delle seguenti opzioni per creare il cluster. Per prestazioni ottimali,è consigliabile utilizzare una rete personalizzata con MTU 8896.

Rete personalizzata

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

1. Imposta le variabili di ambiente per i nomi di rete e 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.

2. Crea una rete personalizzata con un MTU di 8896:

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

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

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

4. 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}"

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

   xpk cluster create --cluster=${CLUSTER_NAME} \
       --cluster-cpu-machine-type=n1-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 elevato, 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=n1-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.

Nota: l'impostazione del flag --cluster-cpu-machine-type su n1-standard-8 (o superiore) garantisce che il pool di nodi predefinito abbia una 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 ed e2. In caso contrario, potresti riscontrare errori di quota.

Creare o caricare l'immagine Docker MaxText

Nota: MaxText supporta solo Python 3.12. Crea l'ambiente virtuale con la versione 3.12 per installare le dipendenze corrette.

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

Creare localmente

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

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

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

Caricare l'immagine (facoltativo)

Dopo aver creato l'immagine Docker localmente seguendo le istruzioni nella sezione precedente, puoi caricare l'immagine Docker MaxText 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 vedere l'immagine MaxText in gcr.io con il nome gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Definire 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 ridotte. 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"

Eseguire 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 o 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 carichi di lavoro devono essere univoci all'interno del cluster. In questo esempio, $(date +%H%M) viene aggiunto al nome del carico di lavoro per garantirne l'unicità.

Passaggi successivi

• Utilizzare la piattaforma Google Cloud ML Diagnostics per ottimizzare e diagnosticare i carichi di lavoro
• Eseguire un carico di lavoro di addestramento utilizzando una ricetta ottimizzata per TPU7x
• Eseguire un microbenchmark TPU7x