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 le kit de traitement accéléré (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 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_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

Choisissez l'une des options suivantes pour créer votre cluster. Pour des performances optimales,nous vous recommandons d'utiliser un réseau personnalisé avec une MTU de 8 896.

Réseau personnalisé

Pour créer un réseau personnalisé avec une MTU de 8 896 et l'utiliser pour votre cluster, procédez comme suit :

Définissez des variables d'environnement pour les noms du réseau et du pare-feu :
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
Remplacez les éléments suivants :
- NETWORK_NAME : nom du réseau.
- FIREWALL_NAME : nom de la règle de pare-feu réseau.

Créez un réseau personnalisé avec une MTU de 8 896 :

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

Créez une règle de pare-feu qui autorise le trafic TCP, ICMP et UDP sur votre réseau :

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

Définissez une variable d'environnement pour les arguments du cluster XPK afin d'utiliser le réseau que vous avez créé :
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Créez le cluster XPK. La commande suivante provisionne la capacité à la demande :

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

Pour utiliser la capacité réservée, remplacez --on-demand par --reservation=RESERVATION_NAME. Pour utiliser des VM Spot TPU, remplacez --on-demand par --spot.

Réseau par défaut

Si vous n'avez pas besoin d'un réseau à MTU élevée, vous pouvez créer un cluster qui utilise le réseau VPC par défaut. La commande suivante provisionne la capacité à la demande :

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

Pour utiliser la capacité réservée, remplacez --on-demand par --reservation=RESERVATION_NAME. Pour utiliser des VM Spot TPU, remplacez --on-demand par --spot.

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

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

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