Cette page a été traduite par l'API Cloud Translation.

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

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 de l'IA à grande échelle. Pour en savoir plus, consultez TPU v3.

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

Provisionner des TPU

Vous pouvez provisionner et gérer les TPU7x à l'aide des méthodes suivantes :

GKE : vous pouvez utiliser GKE pour provisionner et gérer les TPU en tant que pool d'accélérateurs pour vos charges de travail de machine learning conteneurisées. Utilisez la Google Cloud CLI pour créer manuellement votre instance de cluster GKE afin de personnaliser ou d'étendre précisément les environnements GKE de production existants. Pour en savoir plus, consultez À propos des TPU dans GKE.
GKE et XPK : XPK est un outil de ligne de commande qui simplifie la création de clusters et l'exécution de charges de travail sur GKE. Il est conçu pour permettre aux professionnels du ML de provisionner des TPU et d'exécuter des tâches d'entraînement sans avoir besoin d'une expertise approfondie de Kubernetes. Utilisez XPK pour créer rapidement des clusters GKE et exécuter des charges de travail à des fins de validation du concept et de test. Pour en savoir plus, consultez le dépôt GitHub XPK.
GKE et TPU Cluster Director : TPU Cluster Director est disponible via une réservation en mode "Toute capacité". Vous bénéficiez ainsi d'un accès complet à toute votre capacité réservée (sans retenue) et d'une visibilité totale sur la topologie matérielle, l'état d'utilisation et l'état de fonctionnement des TPU. Pour en savoir plus, consultez Présentation du mode "Toute capacité".

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

Utilisez Accelerated Processing Kit (XPK) pour créer des clusters GKE à des fins de validation du concept et de test. XPK est un outil de ligne de commande conçu pour simplifier le provisionnement, la gestion et l'exécution des charges de travail de machine learning.

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

Avant de commencer

Avant de commencer, procédez comme suit :

Assurez-vous de disposer d'un projet Google Cloud pour lequel la facturation est activée.
Accéder à TPU7x Pour en savoir plus, contactez l'équipe chargée de votre compte.
Assurez-vous que le compte que vous utilisez avec XPK dispose des rôles listés dans le dépôt GitHub XPK.

Installer XPK et les dépendances

Installez XPK. Suivez les instructions du dépôt GitHub XPK.
Installez Docker en suivant les instructions fournies par votre administrateur ou en suivant les instructions d'installation officielles. Une fois installé, exécutez les commandes suivantes pour configurer Docker et tester l'installation :
```
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
```
Définissez les variables d'environnement suivantes :
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_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. Pour la version Preview, seul us-central1-c est accepté.
- YOUR_CLUSTER_NAME : nom du nouveau cluster.
- YOUR_ACCELERATOR_TYPE : version et topologie du TPU. Exemple : tpu7x-4x4x8. Pour obtenir la liste des topologies compatibles, consultez Configurations compatibles.
- YOUR_RESERVATION_NAME : nom de votre réservation. Pour les réservations partagées, utilisez projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME : nom de votre bucket Cloud Storage, qui sera le répertoire de sortie pour l'entraînement du modèle.

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

Créer un cluster à une seule carte d'interface réseau et à un seul slice

Suivez les instructions de la section Configurer la MTU pour optimiser la configuration de votre réseau.
Renseignez la variable ${CLUSTER_ARGUMENTS}, que vous utiliserez dans la commande xpk cluster create :
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
Créez votre cluster GKE avec des pools de nœuds TPU7x à l'aide de la commande xpk cluster create :
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Définir l'indicateur --cluster-cpu-machine-type sur n1-standard-8 (ou une valeur supérieure) garantit que le pool de nœuds par défaut dispose de suffisamment de processeurs pour les pods système, par exemple le webhook JobSet, ce qui évite les erreurs. Par défaut, XPK utilise e2-standard-16. Certaines zones ne sont compatibles qu'avec des types de processeur spécifiques. Vous devrez peut-être donc choisir entre les types n1, n2 et e2. Sinon, vous risquez de rencontrer des erreurs de quota.
Ajoutez une exclusion de maintenance pour empêcher les mises à niveau du cluster :
```
gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"
```
Remplacez les éléments suivants :
- EXCLUSION_START_TIME : heure de début sélectionnée pour l'exclusion de maintenance au format YYYY-MM-DDTHH:MM:SSZ.
- EXCLUSION_END_TIME : heure de fin sélectionnée pour l'exclusion de maintenance au format YYYY-MM-DDTHH:MM:SSZ.

Créer ou importer l'image Docker MaxText

Vous pouvez créer une image Docker en local à l'aide des scripts fournis par MaxText ou utiliser une image prédéfinie.

Compiler en local

Les commandes suivantes copient votre répertoire local dans le conteneur :

# 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

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

Une fois les commandes exécutées, une image nommée maxtext_base_image doit être créée localement. Vous pouvez utiliser votre image locale directement dans la commande de charge de travail xpk.

Importer une image (facultatif)

Après avoir créé l'image Docker en local en suivant les instructions de la section précédente, vous pouvez l'importer dans le registre à l'aide de la commande suivante :

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Une fois cette commande exécutée, l'image MaxText doit s'afficher dans gcr.io sous le nom gcr.io/PROJECT_ID/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_xpk \
        steps=30"

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

Exécutez la commande xpk workload create pour déployer votre tâche d'entraînement. Vous devez spécifier le flag --base-docker-image pour utiliser l'image de base MaxText, ou le flag --docker-image et l'image que vous souhaitez utiliser. Vous pouvez choisir d'inclure l'option --enable-debug-logs pour activer la journalisation du débogage.

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]

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 son unicité.