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

Ce document explique comment provisionner des ressources TPU7x et fournit un exemple de déploiement d'une charge de travail d'entraînement à l'aide de MaxText et de Cluster Toolkit.

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 TPU v7.

Pour obtenir d'autres exemples optimisés pour TPU7x, consultez Training Recipes for Ironwood TPU sur GitHub.

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

Utilisez Cluster Toolkit pour créer des clusters GKE prêts pour la production et exécuter des charges de travail d'entraînement.

Les sections suivantes expliquent comment déployer une charge de travail d'entraînement à l'aide de MaxText et de Cluster Toolkit.

Avant de commencer

Avant de commencer, procédez comme suit :

  1. Définissez les variables d'environnement suivantes :

    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, qui sera le répertoire de sortie pour l'entraînement du modèle.

  2. Si vous n'avez pas de bucket Cloud Storage, créez-en un à l'aide de la commande suivante :

    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 GKE TPU7x en suivant les instructions de la section Déployer un cluster GKE TPU7x. 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 de charge de travail MaxText ci-dessous).

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 permettant d'exécuter votre script d'entraînement dans le conteneur Docker.

Le modèle MaxText 1B est une configuration du framework MaxText conçue pour entraîner un modèle de langage avec environ un milliard de paramètres. Utilisez ce modèle pour tester les petites échelles de chips. 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"

Déployer la charge de travail d'entraînement

Envoyez votre tâche d'entraînement à l'aide de gcluster job submit. Pour en savoir plus sur les options d'envoi de tâches, consultez le Guide des tâches 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 de charge de travail doivent être uniques dans le cluster. Dans cet exemple, $(date +%H%M) est ajouté au nom de la charge de travail pour garantir l'unicité.

Étapes suivantes