Entraînement Cloud TPU v5p

Cloud TPU v5p est le Cloud TPU de cinquième génération de Google Cloud, successeur du TPU v4. v5p est optimisé pour l'entraînement à grande échelle et d'être une plate-forme de référence pour le développement les LLM, les modèles de diffusion et l'IA générative. Globalement, v5p offre des performances jusqu'à deux fois supérieures à celles de la version 4, tout en multipliant par deux de TPU en plus dans un pod (6 000 tranche la plus grande contre 3 000 dans la v4), ce qui génère jusqu'à quatre fois plus de performances au niveau du pod. Il fonctionne également à une fréquence plus élevée fréquence d'horloge (1,75 GHz contre 1,05 GHz), ajoute SparseCore pour les représentations vectorielles continues à grande échelle et triple la capacité de la mémoire à haute bande passante (HBM).

Concepts Cloud TPU v5p

Si vous ne connaissez pas les Cloud TPU, consultez le Page d'accueil de la documentation TPU

Concepts Cloud TPU (tranches, hôtes, et TensorCores) et l'architecture du système Cloud TPU pour tous les Cloud TPU sont décrites dans le document Architecture du système Cloud TPU .

Chaque version de Cloud TPU nécessite des types d'accélérateurs spécifiques pour l'entraînement ou l'inférence. Ces types d'accélérateurs sont décrits Configurations v5p.

Gérer les ressources TPU

Toutes les commandes que vous pouvez utiliser pour gérer vos VM TPU sont décrites dans la section Gestion des TPU ou Guide de l'utilisateur des ressources en file d'attente pour la gestion en file d'attente.

Configuration du framework

Cette section décrit le processus général de configuration du modèle à l'aide de JAX ou PyTorch avec TPU v5p.

Configuration pour JAX

Si vous avez des formes de tranches supérieures à quatre chips, vous aurez plusieurs VM. sur une même tranche. Dans ce cas, vous devez utiliser l'option --worker=all. pour exécuter l'installation sur toutes les VM TPU à l'aide d'une seule commande:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les les résultats présentés ici ont été produits avec une tranche v5p-32). Ce code vérifie que tout est correctement installé en vérifiant que JAX voit les TensorCores de Cloud TPU et peut exécuter des opérations de base:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Le résultat ressemble à ce qui suit: 

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() indique le nombre total de chips dans pour le segment donné. jax.local_device_count() indique le nombre de chips accessibles par une seule VM dans cette tranche.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

Le résultat ressemble à ce qui suit: 

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Utilisez --node=all pour exécuter la commande sur tous les nœuds de calcul à plusieurs tranches.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Consultez les tutoriels JAX de ce document pour obtenir a commencé par l'entraînement v5p à l'aide de JAX.

Configuration pour PyTorch

L'environnement d'exécution PJRT est le seul environnement d'exécution compatible avec PyTorch 2.1 et versions ultérieures. PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU. Cette section décrit comment commencer à utiliser PJRT sur des pods v5p avec PyTorch/XLA 2.2.0 pour tous les nœuds de calcul.

Installer des dépendances

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Utiliser un script Python avec PJRT pour valider l'installation pour afficher les appareils TPU disponibles (les résultats présentés ici correspondent avec une tranche v5p-32).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Utilisez --node=all pour exécuter la commande sur tous les nœuds de calcul à plusieurs tranches.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Consultez les tutoriels PyTorch de ce document pour avec l'entraînement v5p à l'aide de PyTorch.

Surveiller et profiler

Cloud TPU v5p est compatible avec la surveillance et le profilage à l'aide du les mêmes méthodes que les générations précédentes de Cloud TPU. Vous pouvez Consultez la page Profiler votre modèle à l'aide des outils Cloud TPU. pour en savoir plus sur le profilage et la surveillance des VM Cloud TPU. pour en savoir plus sur la surveillance.

Tutoriels de formation

Cette section se concentre sur les tutoriels d'entraînement pour une seule tranche. Il est possible d'adapter ces tutoriels à l'entraînement à plusieurs tranches obtenu en ajoutant l'option --node=all aux commandes SSH. Pour en savoir plus et connaître les bonnes pratiques, consultez le Présentation des multitranches.

Tutoriels JAX

Train Diffusion 2.1

Ce tutoriel vous explique comment entraîner le modèle Stable Diffusion à partir de HuggingFace utilisant le Pokémon sur Cloud TPU v5p.

Le modèle Stable Diffusion est un modèle texte-image latent qui génère des images réalistes à partir de n'importe quelle entrée de texte. Pour en savoir plus, consultez les ressources suivantes :

Prérequis

  1. Créez des variables d'environnement: 

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

    Description des options de commande

    Variable Description
    PROJECT_ID Nom du projet Google Cloud
    ACCELERATOR_TYPE Consultez les versions TPU correspondant à votre version de TPU.
    ZONE Consultez le document Régions et zones TPU pour connaître les zones disponibles.
    RUNTIME_VERSION Pour la version v5p, utilisez v2-alpha-tpuv5 pour RUNTIME_VERSION.
    SERVICE_ACCOUNT Il s'agit de l'adresse de votre compte de service, dans la console Google Cloud -> IAM -> Comptes de service. Exemple: tpu-service-account@myprojectID.iam.gserviceaccount.com
    TPU_NAME ID textuel attribué par l'utilisateur au TPU, créé lors de la la demande de ressource en file d'attente est allouée.
    QUEUED_RESOURCE_ID ID du texte attribué par l'utilisateur dans la requête de ressource en file d'attente. Voir le document Queued Resources (Ressources en file d'attente) pour en savoir plus sur les ressources en file d'attente.
    QUOTA_TYPE Peut être reserved ou best-effort. Si aucune de ces options n'est spécifiée, le QUOTA_TYPE la valeur par défaut est on-demand. Voir les quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.
    VALID_UNTIL_DURATION Durée de validité de la requête. Voir <ph type="x-smartling-placeholder"></ph> Ressources en file d'attente pour connaître les différentes durées valides.

  2. Créez une ressource TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--valid-until-duration ${VALID_UNTIL_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Vous pourrez vous connecter en SSH à votre VM TPU une fois que la ressource en file d'attente sera dont l'état est ACTIVE. Vérifiez l'état de votre ressource en file d'attente en exécutant la commande suivante:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
--project ${PROJECT_ID} --zone ${ZONE}

    Lorsque la ressource en file d'attente est à l'état ACTIVE, la sortie ressemblera à ce qui suit:

    state: ACTIVE

  3. Installez JAX et ses dépendances.

    # compatible with v5p: only jax version 0.4.19 and later \
# jax 0.4.19 requires py 3.10 \

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
--command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Téléchargez le dépôt HuggingFace et configuration requise.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

  5. Entraîner le modèle

    Entraînez le modèle avec un tampon prémappé de 4 Go.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'

Effectuer un nettoyage

Supprimez votre TPU et la demande de ressource en file d'attente à la fin de votre session ou pour supprimer les demandes de ressources en file d'attente associées à l'état "FAILED" state. Pour supprimer une ressource en file d'attente, supprimez les tranches, puis la ressource en file d'attente en deux étapes:

   gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID}
   --zone=${ZONE} --quiet

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project ${PROJECT_ID} --zone ${ZONE} --quiet

Vous pouvez également utiliser --force pour supprimer la ou les tranches et la requête de ressource en file d'attente. en une seule étape:

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Résultats du benchmark

Le script d'entraînement de Stable Diffusion a été exécuté sur les versions v5p-8, v5p-32 et v5p-128. Le tableau suivant indique le débit.

V5p-8

V5p-32

V5P-128

Étape d'entraînement

150

150

150

Taille du lot global

32

64

64

Débit (exemples/s)

12.10

18,08

19,10

MaxText

Ce tutoriel vous montre comment entraîner le modèle MaxText à l'aide d'une requête ensemble de données sur Cloud TPU.

MaxText est un logiciel Open Source hautes performances, arbitrairement évolutif, LLM bien testé, écrit en Python/JAX pur et ciblant les Cloud TPU. MaxText offre aux chercheurs et aux développeurs un environnement et adaptable pour repousser les frontières de Natural Language la recherche et le développement (TLN).

Avant d'exécuter ce tutoriel, vous devez Configurez votre environnement Cloud TPU.

  1. Configurer des variables d'environnement

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name # user defined TPU name
export ACCELERATOR_TYPE=v5p-256
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export RUN_NAME=your_experiment_run_name # user defined name for this run
export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

    Description des options de commande

    Variable Description
    PROJECT_ID Nom du projet Google Cloud
    TPU_NAME Nom défini par l'utilisateur pour votre TPU.
    ACCELERATOR_TYPE Consultez les versions TPU correspondant à votre version de TPU.
    ZONE Consultez le document Régions et zones TPU pour connaître les zones disponibles.
    RUNTIME_VERSION Pour v5p, utilisez v2-alpha-tpuv5 pour la version d'exécution.
    RUN_NAME Nom de l'exécution du test fourni par l'utilisateur.

    Configuration facultative recommandée pour les tranches multitranches:

    export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

    Si vous exécutez des charges de travail multitranches et souhaitez des performances réseau optimales, envisagez de créer un réseau dédié avec une unité de transmission maximale (MTU) de 8 896 octets et la configuration des règles de pare-feu appropriées. Bien que facultative, cette étape peut améliorer considérablement les performances, en particulier lors d'une augmentation de tranches sur le réseau de centre de données (DCN). Création de notes un réseau nécessite l'autorisation compute.networks.create dans le projet. Les exemples suivants montrent comment créer un réseau et un pare-feu dédiés d'une règle.

    Créez un réseau dédié:

    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 :

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

  2. Cloner le dépôt MaxText

    git clone https://github.com/google/maxtext.git

  3. Entraîner le modèle

    Les sections suivantes décrivent deux options pour l'entraînement de MaxText.

    Option 1

    Si vous voulez qu'un script gère l'intégralité du workflow, du provisionnement Cloud TPU et installation de dépendances pour l'exécution de votre modèle et supprimer des ressources, vous pouvez utiliser multihost_job.py.

    cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
--NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
--VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
--BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
--COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

    Après avoir lancé le script, vous devriez voir un message semblable à à ce qui suit dans le journal. L'emplacement du journal est référencé dans le message de sortie. Cliquez sur le premier lien pour accéder les journaux de tous les nœuds de calcul une fois le provisionnement du TPU terminé.

    ------------------------------------

multihost_job finished running, TPUs are starting up to run your job remotely.

Logs for your job are displayed here:
https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID

To see the output of a single host, you may edit the slice and worker
number in the `log_file_path` property here:

https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID

When your job is finished, the main command log is in your Cloud Storage
bucket:
https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID

View the status of the created TPUs using:
gcloud compute tpus queued-resources list --filter=RUN_NAME
--zone=ZONE --project=PROJECT_ID
Option 2

Exécuter le script d'entraînement plusieurs fois sur une instance Cloud TPU, utilisez le script multihost_runner.py pour utiliser la ressource.

  1. Configurez des variables pour créer un TPU.

    export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export VALID_DURATION=1d
export QUOTA_TYPE=quota_type

    --node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

  2. Créez une ressource TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Vous pourrez vous connecter à vos VM TPU à l'aide de SSH une fois que votre QueuedResource possède l'état ACTIVE:

    Utilisez la commande describe pour interroger l'état de votre ressource en file d'attente.

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE}

    Lorsque votre ressource en file d'attente est à l'état ACTIVE, la sortie est semblable à ce qui suit:

     state: ACTIVE

  3. Se connecter à votre TPU à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  4. Installer des dépendances

    export TPU_NAME=your_tpu_name
export MAXTEXT_OUTPUT_PATH=output-path

    cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND='bash setup.sh'

  5. Exécutez le modèle avec différentes scripts de configuration tels que 32b.sh, 64b.sh. Si vous exécutez le script à partir d'une VM TPU, vous devez ajouter l'option --INTERNAL_IP=true.

    python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

Effectuer un nettoyage

Supprimez votre TPU et les ressources en file d'attente.

Résultats du benchmark

Le script d'entraînement MaxText a été exécuté de 32 à 1 160 milliards avec une précision en bf16. Les résultats de ces exécutions sont présentés dans le tableau suivant.

Nbre de paramètres

Type d'accélérateur

TFLOP/puce/s

Utilisation des FLOPS du modèle

(MFU)

32 Mrds

V5P-128

3,28E+02

71,47%

64 Mrds

V5P-128

3,23E+02

70,31%

128 Mrds

V5P-256

3,15E+02

68,68%

128 Mrds

V5P-512

3,15E+02

68,53%

256 Mrds

V5P-1024

3,16E+02

68,82%

512 Mrds

V5P-1024

2,94E+02

63,99%

1024B

V5P-2048

2,49E+02

64,05%

1024B

V5P-4096

2,97E+02

64,80%

1 160 Mrds

V5P-7680

2,95E+02

64,27%

1 160 Mrds

V5P-12288

3,04E + 02

66,23%

Le modèle des paramètres 256B a été testé sur les versions v5p-512 et v5p-1024. en utilisant les pondérations bf16 et int8. Le tableau suivant affiche les résultats de ces tests.

V5P-512

V5P-512

V5P-1024

V5P-1024

Taille du lot global

(jetons)

5,24E+05

5,24E+05

1,05E+06

1,05E+06

Précision

bf16

int8

bf16

int8

TFLOP/puce/s

307

408

308

414

Utilisation des FLOPS du modèle

(MFU)

66,98%

88,85%

67,09%

90,23%

Tutoriels TensorFlow

Entraîner ResNet sur un seul hôte v5p

Ce tutoriel explique comment entraîner ImageNet sur un TPU v5p-8 à l'aide d'un faux ensemble de données. Si vous souhaitez utiliser un autre jeu de données, reportez-vous à Préparer l'ensemble de données

Prérequis

  1. Créez des variables d'environnement :

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.17.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    Pour ce tutoriel, utilisez v5p-8 comme ACCELERATOR_TYPE.

  2. Créez une ressource TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --node-id ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --${QUOTA_TYPE}

    Vous pourrez vous connecter à votre VM TPU via SSH une fois que votre ressource en file d'attente se trouve l'état ACTIVE. Pour vérifier l'état de votre ressource en file d'attente, exécutez la commande suivante:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  3. Se connecter à votre TPU à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  4. Définir des variables d'environnement

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. Accédez au répertoire du dépôt de modèles et installez les conditions requises.

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

Entraîner le modèle

  1. Exécutez le script d'entraînement.

    python3 official/vision/train.py \
  --tpu=local \
  --experiment=resnet_imagenet \
  --mode=train_and_eval \
  --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
  --model_dir=${MODEL_DIR} \
  --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Effectuer un nettoyage

Supprimez votre TPU et les ressources en file d'attente.

Entraîner ResNet sur un serveur v5p à hôtes multiples

Ce tutoriel explique comment entraîner ImageNet sur v5p-16 ou version ultérieure à l'aide de un jeu de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez la section Préparer l'ensemble de données.

  1. Créez des variables d'environnement :

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name
export ZONE=us-east1-c
export ACCELERATOR_TYPE=v5p-16
export RUNTIME_VERSION=tpu-vm-tf-2.17.0-pod-pjrt
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE peut être de taille v5p-16 ou supérieure.

  2. Créez une ressource TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --node-id ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --${QUOTA_TYPE}

    Vous pourrez vous connecter à votre VM TPU via SSH une fois que votre ressource en file d'attente se trouve l'état ACTIVE.

    Utilisez la commande describe pour interroger l'état de votre ressource en file d'attente:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  3. Se connecter à votre TPU (nœud de calcul zéro) à l'aide de SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  4. Définir des variables d'environnement

    export TPU_NAME=your_tpu_name
export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0

  5. Accédez au répertoire du dépôt de modèles et installez les conditions requises.

    cd $MODELS_REPO && git checkout r2.15.0
pip install -r official/requirements.txt

Entraîner le modèle

  1. Exécutez le script d'entraînement.

    python3 official/vision/train.py \
  --tpu=${TPU_NAME} \
  --experiment=resnet_imagenet \
  --mode=train_and_eval \
  --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
  --model_dir=${MODEL_DIR} \
  --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Effectuer un nettoyage

Supprimez votre TPU et les ressources en file d'attente.

PyTorch/XLA

Llama 2

Ce tutoriel explique comment entraîner le modèle Llama 2 7B sur V5p à l'aide d'une fourchette du dépôt HuggingFace sur PyTorch/XLA avec les options "General and Scalable" Parallélisation pour les graphiques de calcul ML (GSPMD)

Configuration

  1. Créez des variables pour l'ID de projet, le type d'accélérateur, la zone, la version d'exécution, et le nom du TPU.

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_DURATION=1d

  2. Créer une ressource TPU

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Vous pourrez vous connecter à votre VM TPU via SSH une fois que votre QueuedResource est à l'état ACTIVE:

    Utilisez la commande describe pour interroger l'état de votre ressource en file d'attente.

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE}

    Lorsque votre ressource en file d'attente est à l'état ACTIVE, la sortie ressemblera à ce qui suit: 

     state: ACTIVE

  3. Installez Pytorch/XLA et les dépendances requises.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip3 install typing-extensions
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

  4. Téléchargez le dépôt HuggingFace et les conditions d'installation requises.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install git+file://$PWD
pip3 install datasets accelerate evaluate scikit-learn'

  5. Téléchargez la configuration du modèle 7B.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

  6. Entraîner le modèle

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU
export XLA_USE_BF16=1
export XLA_IR_DEBUG=1
export XLA_HLO_DEBUG=1

export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_gather=true
--xla_jf_spmd_threshold_for_windowed_einsum_mib=0"

export PROFILE_EPOCH=0
export PROFILE_STEP=3
export PROFILE_DURATION_MS=20000
export PROFILE_LOGDIR=/tmp/home/

cd transformers
python examples/pytorch/language-modeling/run_clm.py \
 --tokenizer_name hf-internal-testing/llama-tokenizer \
 --dataset_name wikitext \
 --dataset_config_name wikitext-2-raw-v1 \
 --per_device_train_batch_size 96 \
 --per_device_eval_batch_size 8 \
 --num_train_epochs 1 \
 --do_train \
 --output_dir /tmp/output \
 --overwrite_output_dir \
 --config_name ~/config.json \
 --save_strategy no \
 --logging_strategy no \
 --remove_unused_columns no \
 --optim adafactor \
 --torch_dtype bfloat16 \
 --dataloader_drop_last yes \
 --block_size 2048 \
 --spmd_2d_sharding 1 \
 --spmd_grad_chkpt
'

Si vous exécutez un environnement à plusieurs tranches, vous devez définissez l'option --spmd_dcn_parallelism sur le nombre de tranches.

SPMD_USER_GUIDE fournit un guide utilisateur plus détaillé qui explique tous les différents environnements les variables et les boutons d'activation/de désactivation du script HF. Il convient de noter que LIBTPU_INIT_ARGS sera intégré à PyTorch/XLA et activée par défaut dans les prochaines versions.

Effectuer un nettoyage

Supprimez votre TPU et les ressources en file d'attente.

Résultats du benchmark

Le débit des trois tailles de modèle Llama 2 est inclus dans tableau.

V5p-8

V5P-128

V5P-128

Taille du modèle

7 Mrds

13 Mrds

70 Mrds

Taille du lot global

96

1 024

128

Forme du maillage de segmentation

(4, 1)

(64, 1)

(16, 4)

Utilisation des FLOPS du modèle

(MFU)

56,67%

55,80%

51,85%

Assistance et commentaires

Tous vos commentaires sont les bienvenus. Pour partager vos commentaires ou demander de l'aide, remplissez le formulaire d'assistance ou de commentaires Cloud TPU.