Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Entraîner un modèle à l'aide de TPU v6e

Ce document vous guide dans l'entraînement de modèles sur Cloud TPU v6e (également appelé Trillium). Il couvre la configuration de l'environnement, l'optimisation des performances et des exemples pratiques d'entraînement à l'aide de JAX et PyTorch.

Les TPU v6e, également appelés Trillium, sont la sixième génération de TPU de Google. Sur toutes les surfaces techniques, telles que l'API et les journaux, et tout au long de ce document, Trillium sera appelé v6e. Avec 256 puces par pod, l'architecture de TPU v6e présente de nombreuses similitudes avec v5e. Le TPU v6e est optimisé pour l'entraînement, l'ajustement précis et la mise en services des réseaux de neurones convolutifs (CNN), des transformeurs et des modèles de conversion de texte en image. Pour en savoir plus sur l'architecture et les configurations du système TPU v6e, consultez TPU v6e.

Pour savoir comment exécuter l'inférence sur Cloud TPU v6e, consultez les tutoriels suivants :

Avant de commencer

Avant de commencer, vous devez :

Configurer un projet Google Cloud pour les TPU
Installer Google Cloud CLI
Activer l'API Compute Engine
Créer un compte de service Compute Engine et accorder des autorisations

Pour en savoir plus, consultez Configurer l'environnement Cloud TPU.

Vérifier le quota et les autorisations

Vérifiez que votre projet dispose des quotas suivants :

Quota d'accélérateurs TPU Compute Engine
Quota de VM Compute Engine
Adresses IP en cours d'utilisation
Quota pour Hyperdisk Balanced et pour tout autre type de disque que vous souhaitez utiliser

Remarque : v6e est compatible avec Hyperdisk Balanced et Hyperdisk ML, mais pas avec les disques persistants. Pour en savoir plus, consultez Options de stockage pour les données Cloud TPU.

Si vous utilisez Google Kubernetes Engine (GKE) avec XPK (Accelerated Processing Kit), vous avez besoin d'autorisations supplémentaires dans la console Google Cloud . Pour en savoir plus, consultez Autorisations requises dans la console Google Cloud .

Options de provisionnement

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

Compute Engine (Google Cloud CLI) : vous pouvez utiliser l'outil Google Cloud CLI pour gérer les VM TPU directement en tant que ressources Compute Engine. Vous pouvez provisionner des VM TPU individuelles ou utiliser des groupes d'instances gérés (MIG) pour gérer des groupes de VM TPU. Pour en savoir plus, consultez Présentation de la création de TPU.
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. 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. Pour en savoir plus, consultez le dépôt GitHub XPK.

Provisionner des TPU v6e avec GKE et XPK

Si vous utilisez GKE avec des TPU v6e, vous pouvez utiliser des commandes Kubernetes ou XPK pour provisionner des TPU et entraîner ou diffuser des modèles. Pour en savoir plus sur l'utilisation de GKE avec des TPU, consultez À propos des TPU dans GKE.

Cloud TPU v6e est compatible avec les configurations de carte d'interface réseau (NIC) qui vous permettent de faire évoluer le débit sur plusieurs réseaux. Les sections suivantes fournissent des commandes permettant de créer un cluster GKE avec prise en charge des cartes d'interface réseau simples ou multiples à l'aide de XPK. Pour la plupart des charges de travail à tranche unique, une carte d'interface réseau unique offre des performances suffisantes avec moins de configuration. Pour les charges de travail Multislice et celles qui nécessitent des vitesses d'ingestion de données élevées, utilisez plusieurs cartes d'interface réseau.

Créer un cluster compatible avec une seule carte d'interface réseau à l'aide de XPK

Pour la plupart des charges de travail à tranche unique, une seule carte d'interface réseau offre des performances suffisantes avec moins de configuration. Pour les charges de travail Multislice et celles qui nécessitent des vitesses d'ingestion de données élevées, utilisez plusieurs cartes d'interface réseau.

Les sections suivantes expliquent comment créer un cluster GKE avec prise en charge des cartes d'interface réseau simples à l'aide de XPK.

Installer XPK et configurer les variables d'environnement

Installez XPK. Suivez les instructions du dépôt GitHub XPK.
Configurez des variables d'environnement pour votre cluster :
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=1
```
Définissez les variables d'environnement suivantes :
- CLUSTER_NAME : nom de votre cluster.
- ZONE : zone dans laquelle le cluster TPU sera créé. Pour en savoir plus sur les zones compatibles, consultez Régions et zones.
- PROJECT_ID : ID de votre projet Google Cloud .
- ACCELERATOR_TYPE : le type de TPU, également appelé type d'accélérateur, spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Exemple : v6e-256. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
- NUM_SLICES : nombre de tranches de TPU pour votre cluster. Chaque tranche comporte le nombre de puces spécifié dans ACCELERATOR_TYPE. Pour un cluster à tranche unique, définissez NUM_SLICES sur 1. Pour un cluster Multislice, spécifiez le nombre de tranches en fonction des exigences d'évolutivité de votre charge de travail. Le nombre total de puces dans le cluster est le nombre de puces dans ACCELERATOR_TYPE multiplié par NUM_SLICES.

Créer le cluster

Choisissez l'une des options suivantes pour créer votre cluster. Nous vous recommandons d'utiliser un réseau personnalisé avec une MTU de 8 896 pour des performances optimales. Pour en savoir plus, consultez Configurer la MTU.

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 de réseau et de 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=e2-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=e2-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 un cluster compatible avec plusieurs cartes d'interface réseau à l'aide de XPK

Pour les charges de travail Multislice ou d'autres charges de travail nécessitant une bande passante réseau élevée, par exemple pour l'ingestion de données, vous pouvez utiliser plusieurs cartes d'interface réseau pour améliorer les performances. Lorsque vous utilisez plusieurs cartes d'interface réseau, chaque VM TPU se voit attribuer des interfaces réseau supplémentaires, chacune connectée à un réseau VPC unique, ce qui augmente le débit global du réseau. Pour la plupart des charges de travail à tranche unique, une carte d'interface réseau unique offre des performances suffisantes avec moins de configuration.

Les sections suivantes expliquent comment créer un cluster GKE avec prise en charge de plusieurs cartes d'interface réseau à l'aide de XPK.

Installer XPK et configurer les variables d'environnement

Installez XPK. Suivez les instructions du dépôt GitHub XPK.
Configurez des variables d'environnement pour votre cluster et votre réseau principal :
```
export CLUSTER_NAME=XPK_CLUSTER_NAME
export REGION=REGION
export ZONE=us-east1-d
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export FIREWALL_RULE_NAME_1=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME_1=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG_1=${CLUSTER_NAME}-natconfig-1-${ZONE}

export NETWORK_NAME_2=${CLUSTER_NAME}-mtu9k-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME_2=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME_2=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG_2=${CLUSTER_NAME}-natconfig-2-${ZONE}
```
Définissez les variables d'environnement suivantes :
- CLUSTER_NAME : nom de votre cluster.
- REGION : région dans laquelle votre cluster TPU sera créé.
- ZONE : zone dans laquelle le cluster TPU sera créé. Pour en savoir plus sur les zones compatibles, consultez Régions et zones.
- PROJECT_ID : ID de votre projet Google Cloud .
- ACCELERATOR_TYPE : le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Exemple :v6e-256 Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
- NUM_SLICES : nombre de tranches de TPU pour votre cluster. Chaque tranche comporte le nombre de puces spécifié dans ACCELERATOR_TYPE. Pour un cluster à tranche unique, définissez NUM_SLICES sur 1. Pour un cluster Multislice, spécifiez le nombre de tranches en fonction des exigences d'évolutivité de votre charge de travail. Le nombre total de puces dans le cluster correspond au nombre de puces dans ACCELERATOR_TYPE multiplié par NUM_SLICES.

Créer les ressources réseau principales

Créez le réseau principal avec une unité de transmission maximale (MTU) de 8 896 :
```
gcloud compute networks create ${NETWORK_NAME_1} \
    --mtu=8896 \
    --bgp-routing-mode=regional \
    --subnet-mode=custom \
    --project=${PROJECT_ID}
```
L'utilisation d'un réseau personnalisé avec une MTU de 8 896 offre de meilleures performances. Pour en savoir plus, consultez Configurer la MTU.

Créez le sous-réseau principal :

gcloud compute networks subnets create ${SUBNET_NAME_1} \
    --network=${NETWORK_NAME_1} \
    --range=10.11.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

Créez une règle de pare-feu pour le réseau principal qui autorise le trafic tcp, icmp et udp sur le réseau principal :

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

Créez un routeur Cloud Router pour le réseau principal :

gcloud compute routers create ${ROUTER_NAME_1} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_1} \
    --region=${REGION}

Configurez la NAT pour le réseau principal. La commande suivante permet au trafic de votre cluster d'accéder à Internet :

gcloud compute routers nats create ${NAT_CONFIG_1} \
    --router=${ROUTER_NAME_1} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

Créer les ressources réseau secondaires

Créez le réseau secondaire :

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

Créez un sous-réseau pour le réseau secondaire :

gcloud compute networks subnets create ${SUBNET_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --range=10.10.0.0/18 \
    --region=${REGION} \
    --project=${PROJECT_ID}

Créez une règle de pare-feu pour autoriser le trafic dans le nouveau réseau :

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME_2} \
    --network=${NETWORK_NAME_2} \
    --allow tcp,icmp,udp \
    --source-ranges 10.10.0.0/18 \
    --project=${PROJECT_ID}

Créez un routeur Cloud Router pour le réseau secondaire :

gcloud compute routers create ${ROUTER_NAME_2} \
    --project=${PROJECT_ID} \
    --network=${NETWORK_NAME_2} \
    --region=${REGION}

Créez une configuration NAT pour le routeur Cloud Router :

gcloud compute routers nats create ${NAT_CONFIG_2} \
    --router=${ROUTER_NAME_2} \
    --region=${REGION} \
    --auto-allocate-nat-external-ips \
    --nat-all-subnet-ip-ranges \
    --project=${PROJECT_ID} \
    --enable-logging

Créer le cluster

Définissez une variable d'environnement pour les arguments de cluster et de pool de nœuds afin d'utiliser les réseaux et sous-réseaux que vous avez créés :
```
export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"
```
Ces arguments configurent le cluster pour qu'il utilise les deux réseaux que vous avez créés pour la prise en charge de plusieurs cartes d'interface réseau.
Créez le cluster. La commande suivante provisionne la capacité à la demande :
```
xpk cluster create \
    --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=e2-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE}  \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
    --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
    --create-vertex-tensorboard
```
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.

Remarque : L'option --cluster-cpu-machine-type spécifie le type de machine pour les VM hôtes de processeur du cluster. Bien que e2-standard-8 soit une valeur par défaut recommandée pour de nombreuses charges de travail, vous pouvez sélectionner un autre type de machine si votre charge de travail présente des exigences spécifiques en termes de processeur ou de mémoire pour des tâches telles que le prétraitement des données.

Valider la configuration de plusieurs cartes d'interface réseau

Après avoir créé un cluster compatible avec plusieurs cartes d'interface réseau, vous pouvez vérifier que les deux cartes d'interface réseau sont utilisées en créant une charge de travail XPK et en ajoutant le flag --command ifconfig.

Utilisez la commande suivante pour afficher la sortie de la commande ifconfig dans les journaux de la consoleGoogle Cloud . Vous devez spécifier le flag --base-docker-image maxtext_base_image pour utiliser l'image de base MaxText, comme dans l'exemple suivant, ou le flag --docker-image et l'image que vous souhaitez utiliser.
```
xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --num-slices=${NUM_SLICES}  \
    --on-demand \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --command "ifconfig"
```
Si vous souhaitez activer les journaux de débogage ou utiliser Vertex AI TensorBoard, ajoutez les arguments facultatifs suivants à la commande :
```
--enable-debug-logs \
--use-vertex-tensorboard
```
Vérifiez que la MTU est définie sur 8 896 pour eth0 et eth1 en consultant la sortie de la charge de travail XPK dans les journaux de la console Google Cloud .

Configurer JAX ou PyTorch

Les ressources suivantes expliquent comment configurer JAX ou PyTorch sur votre TPU, en fonction de la méthode de provisionnement et de gestion que vous utilisez :

GKE Autopilot : préparer votre application TPU
GKE Standard : préparer vos charges de travail
GKE et XPK : README XPK
Cloud TPU à hôte unique avec JAX : exécuter un calcul sur une VM Cloud TPU avec JAX
Cloud TPU multi-hôte avec JAX : Exécuter du code JAX sur des tranches de TPU
Cloud TPU avec PyTorch : PyTorch sur Cloud TPU

Pour configurer et exécuter XPK avec MaxText, consultez Exécuter MaxText à grande échelle avec XPK .

Utiliser SkyPilot

Vous pouvez utiliser Cloud TPU v6e avec SkyPilot. SkyPilot est un framework Open Source qui simplifie l'exécution, la gestion et le scaling des charges de travail d'IA. Vous pouvez ajouter des informations sur la localisation et les tarifs liés à la version 6e à SkyPilot. Pour en savoir plus, consultez l'exemple SkyPilot TPU v6e.

Exemples d'entraînement

Les sections suivantes fournissent des exemples d'entraînement de modèles MaxText, MaxDiffusion et PyTorch sur Cloud TPU v6e.

Ces exemples ont été testés avec les versions logicielles suivantes :

Python 3.10 ou une version ultérieure.
Versions logicielles courantes :
- JAX courant 0.4.32.dev20240912
- LibTPU courant 0.1.dev20240912+nightly
Versions logicielles stables :
- JAX + JAX Lib version 0.4.37

Entraîner MaxText et MaxDiffusion sur Cloud TPU v6e

Les sections suivantes couvrent le cycle de vie de l'entraînement des modèles MaxText et MaxDiffusion.

De manière générale, voici les grandes étapes à suivre :

Créez l'image de base de la charge de travail.
Exécutez votre charge de travail à l'aide de XPK.
1. Créez la commande d'entraînement pour la charge de travail.
2. Déployez la charge de travail.
Suivez la charge de travail et affichez les métriques.
Supprimez la charge de travail XPK si vous n'en avez pas besoin.
Supprimez le cluster lorsqu'il n'est plus nécessaire.

Créer l'image de base

Installez MaxText ou MaxDiffusion, puis créez l'image Docker :

Clonez le dépôt que vous souhaitez utiliser et accédez à son répertoire :

MaxText :

git clone https://github.com/google/maxtext.git
cd maxtext
git checkout maxtext-v0.2.0

MaxDiffusion :

git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103

Configurez Docker pour utiliser Google Cloud CLI :
```
gcloud auth configure-docker
```

Créez l'image Docker à l'aide de la commande suivante ou d'une image IA JAX. Pour en savoir plus sur les images JAX AI, consultez Images JAX AI.

MaxText :

bash dependencies/scripts/docker_build_dependency_image.sh MODE=stable

MaxDiffusion :

bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest

Définissez l'ID de votre projet dans votre configuration gcloud CLI active :
```
gcloud config set project ${PROJECT_ID}
```
Si vous lancez la charge de travail à partir d'une machine sur laquelle l'image n'a pas été créée localement, importez l'image.

Remarque : Assurez-vous d'utiliser bash version 4.2 ou ultérieure.
1. Définissez la variable d'environnement CLOUD_IMAGE_NAME :
```
export CLOUD_IMAGE_NAME=${USER}_runner
```
2. Importez l'image :
```
bash dependencies/scripts/docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
```

Exécuter votre charge de travail à l'aide de XPK

Définissez les variables d'environnement suivantes si vous n'utilisez pas les valeurs par défaut définies par MaxText ou MaxDiffusion :
```
export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
export PER_DEVICE_BATCH_SIZE=2
export NUM_STEPS=30
export MAX_TARGET_LENGTH=8192
```

Créez le script de votre modèle. Ce script sera copié en tant que commande d'entraînement lors d'une étape ultérieure.

N'exécutez pas encore le script du modèle.

MaxText

MaxText est un LLM Open Source hautes performances et hautement évolutif, écrit en Python et JAX purs, et ciblant les Google Cloud TPU et les GPU pour l'entraînement et l'inférence.

JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

Gemma 2

Gemma est une famille de LLM à pondération ouverte développée par Google DeepMind, basée sur la recherche et la technologie Gemini.

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

Mixtral 8x7b

Mixtral est un modèle d'IA de pointe développé par Mistral AI, qui utilise une architecture MoE (Mixture of Experts) éparse.

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

MaxDiffusion

MaxDiffusion est une collection d'implémentations de référence de divers modèles de diffusion latente écrits en Python et JAX purs, qui s'exécutent sur des appareils XLA, y compris les Cloud TPU et les GPU. Stable Diffusion est un modèle latent de texte vers image qui génère des images photoréalistes à partir de n'importe quelle entrée de texte.

Vous devez installer une branche Git spécifique pour exécuter MaxDiffusion, comme indiqué dans le script d'entraînement suivant.

git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

Exportez les variables suivantes :
```
export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT
```
Descriptions des variables d'environnement
- CLUSTER_NAME : le nom du cluster
- ACCELERATOR_TYPE : le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
- NUM_SLICES : nombre de tranches de TPU.
- YOUR_MODEL_SCRIPT : script de modèle à exécuter en tant que commande d'entraînement.
Exécutez le modèle à l'aide du script que vous avez créé à l'étape précédente. 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'ajouter les options facultatives suivantes :
- Vous pouvez activer la journalisation du débogage en incluant l'option --enable-debug-logs. Pour en savoir plus, consultez Déboguer JAX sur MaxText.
- Vous pouvez créer un Vertex AI Experiment pour importer des données dans Vertex AI TensorBoard en incluant le flag --use-vertex-tensorboard. Pour en savoir plus, consultez Surveiller JAX sur MaxText à l'aide de Vertex AI.
```
xpk workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"
```
Le résultat inclut un lien pour suivre votre charge de travail. Ouvrez le lien et cliquez sur l'onglet Journaux pour suivre votre charge de travail en temps réel.

Déboguer JAX sur MaxText

Utilisez des commandes XPK supplémentaires pour déterminer pourquoi le cluster ou la charge de travail ne s'exécutent pas :

Liste des charges de travail XPK
Inspecteur XPK
Activez la journalisation détaillée dans les journaux de charge de travail à l'aide du flag --enable-debug-logs lorsque vous créez la charge de travail XPK.

Surveiller JAX sur MaxText à l'aide de Vertex AI

Pour utiliser TensorBoard, votre compte utilisateur Google Cloud doit disposer du rôle aiplatform.user. Exécutez la commande suivante pour attribuer ce rôle :

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

Affichez les données scalaires et de profil via TensorBoard géré par Vertex AI.

Augmentez le nombre de requêtes de gestion des ressources (CRUD) pour la zone que vous utilisez, de 600 à 5 000. Cela ne devrait pas poser de problème pour les petites charges de travail utilisant moins de 16 VM.

Installez les dépendances telles que cloud-accelerator-diagnostics pour Vertex AI :

# xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

Créez votre cluster à l'aide de l'option --create-vertex-tensorboard, comme indiqué dans Créer Vertex AI TensorBoard. Vous pouvez également exécuter cette commande sur des clusters existants.
Créez votre Vertex AI Experiments lorsque vous exécutez votre charge de travail XPK à l'aide de l'indicateur --use-vertex-tensorboard et de l'indicateur facultatif --experiment-name. Pour obtenir la liste complète des étapes, consultez Créer un test Vertex AI pour importer des données dans Vertex AI TensorBoard.

Les journaux incluent un lien vers un Vertex AI TensorBoard, semblable à ce qui suit :

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

Vous pouvez également trouver le lien Vertex AI TensorBoard dans la console Google Cloud . Accédez à Vertex AI Experiments dans la console Google Cloud . Sélectionnez la région appropriée dans le menu déroulant.

Le répertoire TensorBoard est également écrit dans le bucket Cloud Storage que vous avez spécifié avec ${BASE_OUTPUT_DIR}.

Supprimer votre charge de travail XPK

Utilisez la commande xpk workload delete pour supprimer une ou plusieurs charges de travail en fonction du préfixe ou de l'état du job. Cette commande peut être utile si vous avez envoyé des charges de travail XPK qui n'ont plus besoin d'être exécutées ou si des jobs sont bloqués dans la file d'attente.

Supprimer votre cluster

Exécutez la commande xpk cluster delete pour supprimer votre cluster :

xpk cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

Résultats du benchmark MaxDiffusion

Nous avons exécuté le script d'entraînement pour MaxDiffusion sur un TPU v6e-4, un TPU v6e-16 et deux TPU v6e-16. Le tableau suivant présente les débits mesurés.

	v6e-4	v6e-16	Deux v6e-16
Étapes de l'entraînement	0,069	0,073	0,13
Taille du lot global	8	32	64
Débit (exemples/s)	115,9	438,4	492,3

Entraîner un modèle à l'aide de PyTorch sur Cloud TPU

Vous pouvez utiliser TorchTPU directement pour entraîner votre modèle ou réutiliser la bibliothèque de référence TorchTitan. Pour commencer à utiliser les charges de travail v6e, consultez les recettes TorchTitan Trillium. Pour en savoir plus sur l'écosystème PyTorch sur Cloud TPU, consultez PyTorch sur Cloud TPU.

Entraîner un modèle à l'aide de TPU v6e

Avant de commencer

Vérifier le quota et les autorisations

Options de provisionnement

Provisionner des TPU v6e avec GKE et XPK

Créer un cluster compatible avec une seule carte d'interface réseau à l'aide de XPK

Installer XPK et configurer les variables d'environnement

Créer le cluster

Réseau personnalisé

Réseau par défaut

Créer un cluster compatible avec plusieurs cartes d'interface réseau à l'aide de XPK

Installer XPK et configurer les variables d'environnement

Créer les ressources réseau principales

Créer les ressources réseau secondaires

Créer le cluster

Valider la configuration de plusieurs cartes d'interface réseau

Configurer JAX ou PyTorch

Utiliser SkyPilot

Exemples d'entraînement

Entraîner MaxText et MaxDiffusion sur Cloud TPU v6e

Créer l'image de base

Exécuter votre charge de travail à l'aide de XPK

MaxText

Gemma 2

Mixtral 8x7b

MaxDiffusion

Descriptions des variables d'environnement

Déboguer JAX sur MaxText

Surveiller JAX sur MaxText à l'aide de Vertex AI

Supprimer votre charge de travail XPK

Supprimer votre cluster

Résultats du benchmark MaxDiffusion

Entraîner un modèle à l'aide de PyTorch sur Cloud TPU