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/XLA.

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 le serving 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 :

  • Créez un compte Google Cloud et un projet avec la facturation activée.
  • Installer les composants alpha de Google Cloud CLI
  • Activer l'API Cloud TPU
  • Créer un agent de service Cloud TPU
  • Créer un compte de service Cloud TPU 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 :

Si vous utilisez GKE avec XPK, vous avez besoin d'autorisations supplémentaires dans la console Google Cloud . Pour en savoir plus, consultez Autorisations requises dans la consoleGoogle Cloud .

Provisionner des TPU

Vous pouvez provisionner et gérer les TPU v6e à 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. 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.
  • Ressources Cloud TPU en file d'attente : les ressources en file d'attente vous permettent de demander une capacité de TPU qui est provisionnée lorsqu'elle devient disponible. Cette solution est idéale pour les tâches par lot et les charges de travail tolérantes aux pannes qui peuvent attendre dans une file d'attente. Vous pouvez spécifier une période pour votre demande. Pour en savoir plus, consultez Gérer les ressources mises en file d'attente.

Provisionner des Cloud TPU v6e avec GKE et XPK

Si vous utilisez des commandes GKE avec v6e, vous pouvez utiliser des commandes Kubernetes ou XPK pour provisionner des Cloud TPU et entraîner ou diffuser des modèles. Consultez Planifier des Cloud TPU dans GKE pour savoir comment planifier vos configurations Cloud TPU dans les clusters GKE. Les sections suivantes fournissent des commandes permettant de créer un cluster XPK avec prise en charge d'une seule interface réseau et de plusieurs interfaces réseau.

Créer un cluster XPK compatible avec une seule interface réseau

export CLUSTER_NAME=xpk-cluster-name
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME=${CLUSTER_NAME}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw
 
gcloud compute networks create ${NETWORK_NAME} \
   --mtu=8896 \
   --project=${PROJECT_ID} \
   --subnet-mode=auto \
   --bgp-routing-mode=regional
 
gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
   --network=${NETWORK_NAME} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}
 
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

python3 xpk.py cluster create --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --create-vertex-tensorboard

Description des options de commande

Variable Description
CLUSTER_NAME Nom attribué par l'utilisateur au cluster XPK.
PROJECT_ID Nom du projet :Google Cloud Utilisez un projet existant ou créez-en un. Pour en savoir plus, consultez Configurer votre projet Google Cloud .
ZONE Pour connaître les zones compatibles, consultez le document Régions et zones Cloud TPU.
TPU_TYPE Consultez la section Types d'accélérateurs.
NUM_SLICES Nombre de segments à créer
CLUSTER_ARGUMENTS Réseau et sous-réseau à utiliser.

Par exemple : --network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}

NUM_SLICES Nombre de tranches à créer.
NETWORK_NAME Nom d'un réseau secondaire à utiliser.
NETWORK_FW_NAME Nom d'un pare-feu réseau secondaire à utiliser.

Créer un cluster XPK compatible avec plusieurs cartes d'interface réseau

export CLUSTER_NAME=xpk-cluster-name
export REGION=your-region
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}
 
gcloud compute networks create ${NETWORK_NAME_1} \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=${PROJECT_ID}
 
gcloud compute networks subnets create ${SUBNET_NAME_1} \
   --network=${NETWORK_NAME_1} \
   --range=10.11.0.0/18 \
   --region=${REGION} \
   --project=${PROJECT_ID}
 
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_1} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}
 
gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_1} \
   --region=${REGION}
 
gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging
 
# Secondary subnet for multi-nic experience.
# Need custom IP routing to be different from the first network's subnet.

export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}

gcloud compute networks create ${NETWORK_NAME_2} \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=${PROJECT_ID}
 
gcloud compute networks subnets create ${SUBNET_NAME_2} \
   --network=${NETWORK_NAME_2} \
   --range=10.10.0.0/18 \
   --region=${REGION} \
   --project=${PROJECT_ID}
 
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_2} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}
 
gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_2} \
   --region=${REGION}
 
gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging
 
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}"

python3 xpk.py cluster create \
   --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE}  \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
   --create-vertex-tensorboard

Description des options de commande

Variable Description
CLUSTER_NAME Nom attribué par l'utilisateur au cluster XPK.
PROJECT_ID Nom du projet :Google Cloud Utilisez un projet existant ou créez-en un. Pour en savoir plus, consultez Configurer votre projet Google Cloud .
ZONE Pour connaître les zones compatibles, consultez le document Régions et zones Cloud TPU.
TPU_TYPE Consultez la section Types d'accélérateurs.
NUM_SLICES Nombre de segments à créer
CLUSTER_ARGUMENTS Réseau et sous-réseau à utiliser.

Par exemple : --enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}

NODE_POOL_ARGUMENTS Réseau de nœuds supplémentaires à utiliser.

Par exemple : --additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}

NUM_SLICES Nombre de segments à créer (nécessaire uniquement pour Multislice).
NETWORK_NAME Nom d'un réseau secondaire à utiliser.
NETWORK_FW_NAME Nom d'un pare-feu réseau secondaire à utiliser.

Configurer JAX ou PyTorch

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

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

Optimiser les performances du réseau

Cette section explique comment optimiser les performances de votre réseau en configurant l'unité de transmission maximale (MTU), en utilisant plusieurs cartes d'interface réseau pour les environnements multislices et en améliorant les paramètres TCP.

Configurer la MTU

Pour obtenir les meilleures performances réseau, utilisez un réseau avec une MTU (unité de transmission maximale) de 8 896.

Par défaut, un cloud privé virtuel (VPC) ne fournit qu'une MTU de 1 460 octets, ce qui offre des performances réseau sous-optimales. Vous pouvez définir la MTU d'un réseau VPC sur n'importe quelle valeur comprise entre 1 300 octets et 8 896 octets (inclus). Les tailles de MTU personnalisées les plus courantes sont de 1 500 octets (le standard Ethernet) ou de 8 896 octets (le maximum possible). Pour en savoir plus, consultez Tailles de MTU valides pour le réseau VPC.

Pour en savoir plus sur la modification du paramètre de MTU d'un réseau existant ou par défaut, consultez Modifier le paramètre de MTU d'un réseau VPC.

L'exemple suivant crée un réseau avec une MTU de 8 896 et une règle de pare-feu correspondante qui autorise le trafic TCP, ICMP et UDP au sein du réseau.

export RESOURCE_NAME=your-resource-name
export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork
export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall
gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \
    --subnet-mode=auto --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp --project=${PROJECT_ID}

Remplacez your-resource-name par un nom de base pour le réseau et le pare-feu.

Utiliser l'option multi-NIC pour Multislice

Si vous utilisez un environnement multislices, définissez les variables d'environnement suivantes, qui sont requises pour un sous-réseau secondaire :

export NETWORK_NAME_2=${RESOURCE_NAME}
export SUBNET_NAME_2=${RESOURCE_NAME}
export FIREWALL_RULE_NAME=${RESOURCE_NAME}
export ROUTER_NAME=${RESOURCE_NAME}-network-2
export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2
export REGION=your-region

Utilisez les commandes suivantes pour créer un routage IP personnalisé pour le réseau et le sous-réseau.

  1. 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}

  2. 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}

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

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

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

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

  5. Créez une configuration NAT pour Cloud Router.

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

Après avoir créé un slice réseau multiréseau, vous pouvez vérifier que les deux cartes d'interface réseau (NIC) sont utilisées en configurant un cluster XPK et en ajoutant l'indicateur --command ifconfig à la commande de création de charge de travail XPK.

  1. Utilisez la commande workload create suivante pour afficher la sortie de la commande ifconfig dans les journaux de la console Google Cloud et vérifiez que la MTU est définie sur 8 896 pour eth0 et eth1.

    python3 xpk.py workload create \
    --cluster CLUSTER_NAME \
    {--base-docker-image maxtext_base_image | --docker-image your-cloud-image-name} \
    --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

  2. Vérifiez que la MTU d'eth0 et d'eth1 est définie sur 8 896 en consultant la sortie de la charge de travail XPK dans les journaux de la console Google Cloud .

Améliorer les paramètres TCP

Si vous avez provisionné vos Cloud TPU à l'aide de ressources en file d'attente, vous pouvez exécuter la commande suivante pour améliorer les performances du réseau en augmentant les limites du tampon de réception TCP.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

Optimiser les performances de l'allocation de mémoire

La bibliothèque tcmalloc est utilisée par défaut sur les VM Cloud TPU pour améliorer les performances des modèles avec des allocations de mémoire volumineuses et fréquentes. Cette valeur est configurée à l'aide de la variable d'environnement LD_PRELOAD.

Toutefois, pour certaines charges de travail (par exemple, DLRM avec des allocations de tables d'intégration très volumineuses), tcmalloc peut entraîner un ralentissement. Dans ce cas, vous pouvez revenir à la fonction malloc standard en annulant la définition de la variable LD_PRELOAD dans votre session shell avant d'exécuter votre script d'entraînement :

unset LD_PRELOAD

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 version ultérieure
  • Versions logicielles Nightly :
    • JAX quotidien 0.4.32.dev20240912
    • LibTPU nightly 0.1.dev20240912+nightly
  • Versions logicielles stables :
    • JAX + JAX Lib de la 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 étapes à suivre :

  1. Créez l'image de base de la charge de travail.
  2. 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.
  3. Suivez la charge de travail et affichez les métriques.
  4. Supprimez la charge de travail XPK si vous n'en avez pas besoin.
  5. Supprimez le cluster XPK lorsqu'il n'est plus nécessaire.

Créer l'image de base

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

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

    MaxDiffusion :

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

  2. Configurez Docker pour utiliser Google Cloud CLI :

    gcloud auth configure-docker

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

    MaxText :

    bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35

    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

  4. Définissez l'ID de votre projet dans votre configuration gcloud CLI active :

    gcloud config set project ${PROJECT_ID}

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

    1. Définissez la variable d'environnement CLOUD_IMAGE_NAME :

      export CLOUD_IMAGE_NAME=${USER}_runner

    2. Importez l'image :

      bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}

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

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

  2. 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.train 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'"

    Gemma2

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

    python3 -m MaxText.train 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.train 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

    Llama3-8b

    Llama est une famille de LLM à pondération ouverte développée par Meta.

    Pour obtenir un exemple d'exécution de Llama3 sur PyTorch, consultez les modèles torch_xla dans le dépôt torchprime.

    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

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

    Variable Description
    CLUSTER_NAME Nom de votre cluster XPK.
    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.

  4. Exécutez le modèle à l'aide du script que vous avez créé à l'étape précédente. Vous devez spécifier l'indicateur --base-docker-image pour utiliser l'image de base MaxText, ou l'indicateur --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 test Vertex AI pour importer des données dans Vertex AI TensorBoard en incluant l'indicateur --use-vertex-tensorboard. Pour en savoir plus, consultez Surveiller JAX sur MaxText à l'aide de Vertex AI.
    python3 xpk.py 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 :

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.

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

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

  3. Créez votre cluster XPK à l'aide de l'indicateur --create-vertex-tensorboard, comme indiqué dans Créer Vertex AI TensorBoard. Vous pouvez également exécuter cette commande sur des clusters existants.

  4. Créez votre test Vertex AI 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 XPK

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

python3 xpk.py 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 des modèles Llama à l'aide de PyTorch/XLA sur Cloud TPU v6e

Cette section explique comment entraîner des modèles Llama à l'aide de PyTorch/XLA sur Cloud TPU v6e à l'aide de l'ensemble de données WikiText.

Accéder à Hugging Face et au modèle Llama 3

Pour cet exemple, vous avez besoin d'un jeton d'accès utilisateur Hugging Face. Pour savoir comment créer des jetons d'accès utilisateur, consultez la documentation Hugging Face sur les jetons d'accès utilisateur.

Vous devez également disposer de l'autorisation d'accéder au modèle Llama-3-8B sur Hugging Face. Pour y accéder, accédez au modèle Meta-Llama-3-8B sur HuggingFace et demandez-y l'accès.

Créer une VM Cloud TPU

Pour cet exemple, créez un Cloud TPU v6e avec huit puces.

  1. Configurez des variables d'environnement :

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e

    Descriptions des variables d'environnement

    Variable Description
    PROJECT_ID L'ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
    TPU_NAME Nom du TPU.
    ZONE Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones de TPU.
    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.
    RUNTIME_VERSION La version logicielle de Cloud TPU.

  2. Créez une VM Cloud TPU :

    gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

Installation

Installez le fork pytorch-tpu/transformers des transformateurs Hugging Face et des dépendances. Cet exemple a été testé avec les versions de dépendances suivantes :

  • torch : compatible avec la version 2.5.0
  • torch_xla[tpu] : compatible avec la version 2.5.0
  • jax : 0.4.33
  • jaxlib : 0.4.33
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

Configurer les fichiers de configuration du modèle

La commande d'entraînement de la section suivante, Exécuter le modèle, utilise deux fichiers de configuration JSON pour définir les paramètres du modèle et la configuration FSDP (Fully Sharded Data Parallel). Le sharding FSDP vous permet d'utiliser une taille de lot plus importante lors de l'entraînement en shardant les pondérations de votre modèle sur plusieurs TPU. Lorsque vous entraînez des modèles plus petits, il peut suffire d'utiliser le parallélisme des données et de répliquer les pondérations sur chaque appareil. Pour en savoir plus sur le partitionnement des Tensors sur plusieurs appareils dans PyTorch/XLA, consultez le guide de l'utilisateur SPMD PyTorch/XLA.

  1. Créez le fichier de configuration des paramètres du modèle. Voici la configuration des paramètres du modèle pour Llama-3-8B. Pour les autres modèles, recherchez le fichier de configuration sur Hugging Face. Par exemple, consultez la configuration Llama-2-7B.

    cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

  2. Créez le fichier de configuration FSDP :

    cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

    Pour en savoir plus sur FSDP, consultez Fully Sharded Data Parallel using SPMD .

  3. Importez les fichiers de configuration dans vos VM Cloud TPU à l'aide de la commande suivante :

    gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

Exécuter le modèle

À l'aide des fichiers de configuration que vous avez créés dans la section précédente, exécutez le script run_clm.py pour entraîner le modèle Llama-3-8B sur l'ensemble de données WikiText. Le script d'entraînement prend environ 10 minutes sur un Cloud TPU v6e-8.

  1. Connectez-vous à Hugging Face sur votre Cloud TPU à l'aide de la commande suivante :

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

  2. Exécutez l'entraînement du modèle :

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

Résoudre les problèmes liés à PyTorch/XLA

Si vous avez défini les variables facultatives pour le débogage dans la section précédente, le profil du modèle sera stocké à l'emplacement spécifié par la variable PROFILE_LOGDIR. Vous pouvez extraire le fichier xplane.pb stocké à cet emplacement et utiliser tensorboard pour afficher les profils dans votre navigateur en suivant les instructions TensorBoard.

Si PyTorch/XLA ne fonctionne pas comme prévu, consultez le guide de dépannage, qui contient des suggestions pour déboguer, profiler et optimiser votre modèle.