Addestra un modello utilizzando TPU7x (Ironwood)

Anteprima

Questo prodotto o funzionalità è soggetto ai "Termini delle offerte pre-GA" nella sezione Termini di servizio generali dei Termini specifici dei servizi. I prodotti e le funzionalità pre-GA sono disponibili "così come sono" e potrebbero avere un supporto limitato. Per ulteriori informazioni, consulta le descrizioni delle fasi di lancio.

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

TPU7x è la prima release della famiglia Ironwood,la TPU di settima generazione di Google. Google CloudLa 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.

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 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 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 saperne di più, consulta 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 di XPK.

Installa XPK e le dipendenze

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

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

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'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_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
   

Crea un cluster a singola NIC e singola sezione

Scegli una delle seguenti opzioni per creare il cluster. Per ottenere prestazioni ottimali,è consigliabile utilizzare una rete personalizzata con MTU di 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 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.

2. Crea una rete personalizzata con un'unità 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 tua 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 elevata, 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 node pool predefinito disponga di CPU sufficiente per i pod di sistema, ad esempio il webhook JobSet, impedendo 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.

Crea o carica l'immagine Docker MaxText

Nota: MaxText supporta solo Python 3.12. Crea il tuo ambiente virtuale con 3.12 per installare le dipendenze corrette.

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

Creare build 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

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

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 il logging 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à.

Passaggi successivi

• Utilizza la piattaforma Google Cloud ML Diagnostics per ottimizzare ed eseguire la diagnostica dei tuoi carichi di lavoro
• Esegui un carico di lavoro di addestramento utilizzando una ricetta ottimizzata per TPU7x
• Esegui un microbenchmark TPU7x