Entraîner un modèle à l'aide de TPU7x (Ironwood)

Ce document explique comment provisionner des ressources TPU7x avec Cluster Toolkit et déployer une charge de travail d'entraînement à l'aide de MaxText (JAX) ou de TorchTitan sur TPU (PyTorch).

Le TPU7x est la première version de la famille Ironwood,le TPU de septième génération de Google Cloud. La génération Ironwood est conçue pour l'entraînement et l'inférence d'IA à grande échelle. Pour en savoir plus, consultez TPU7x.

Pour obtenir d'autres exemples optimisés pour TPU7x, consultez les dépôts suivants sur GitHub :

Avant de commencer

Avant de commencer, assurez-vous de remplir les conditions suivantes :

Ensuite, procédez comme suit :

  1. Définissez les variables d'environnement pour votre projet et votre 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"

    Remplacez les éléments suivants :

    • YOUR_PROJECT_ID : ID de votre projet Google Cloud .
    • YOUR_ZONE : zone dans laquelle créer le cluster.
    • YOUR_CLUSTER_NAME : nom du nouveau cluster.
    • YOUR_BUCKET_NAME : nom de votre bucket Cloud Storage pour la sortie d'entraînement.

  2. Si vous n'avez pas de bucket Cloud Storage, créez-en un :

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

Déployer le cluster TPU7x (Ironwood)

Déployez un cluster TPU7x Google Kubernetes Engine (GKE) en suivant les instructions de la section Déployer un cluster TPU7x GKE. Assurez-vous que la topologie du pool de nœuds de votre cluster correspond aux exigences de votre charge de travail d'entraînement, par exemple 4x4x8 pour l'exemple MaxText ou 2x2x4 pour l'exemple TorchTitan.

Déployer une charge de travail d'entraînement avec MaxText

Les sections suivantes expliquent comment créer une image de conteneur et envoyer un job d'entraînement MaxText avec Cluster Toolkit.

Créer et importer l'image Docker MaxText

Créez et transférez l'image 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}

Définir la commande d'entraînement MaxText

Préparez la commande pour exécuter votre script d'entraînement dans le conteneur Docker.

La configuration du modèle MaxText 1B entraîne un modèle de langage avec environ 1 milliard de paramètres. Utilisez ce modèle pour tester de petites échelles de puces. Les performances ne sont pas optimisées.

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"

Envoyer le job d'entraînement MaxText

Envoyez votre job d'entraînement à l'aide de gcluster job submit. Pour en savoir plus sur les options d'envoi de jobs, consultez le guide des jobs 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}"

Les noms des charges de travail doivent être uniques dans le cluster. Cet exemple ajoute $(date +%H%M) au nom de la charge de travail.

Déployer une charge de travail d'entraînement avec TorchTitan

TorchTitan utilise TorchTPU pour exécuter un entraînement PyTorch distribué sur Cloud TPU. Pour en savoir plus, consultez le fichier README de TorchTitan sur TPU.

Étapes suivantes