Entraînement Cloud TPU v5e
Cloud TPU v5e est l'accélérateur d'IA de dernière génération de Google Cloud. Avec un d'encombrement de 256 puces par pod, la version v5e est optimisée pour offrir produit pour les modèles Transformer, Text-to-Image et de réseaux de neurones convolutifs (CNN) l'entraînement, l'affinage et l'inférence. Pour en savoir plus sur l'utilisation de Cloud TPU Pour l'inférence, consultez la page Inférence avec v5e.
Pour en savoir plus sur le matériel et les configurations des TPU Cloud TPU v5e, consultez la page TPU v5e.
Commencer
Les sections suivantes décrivent comment commencer à utiliser TPU v5e.
Quota de requêtes
Vous avez besoin d'un quota d'utilisation de TPU v5e pour l'entraînement. Il existe différents types de quotas les TPU à la demande, les TPU réservés et les VM TPU Spot. Il existe des les quotas requis si vous utilisez votre TPU v5e pour inférence. Pour en savoir plus sur les quotas, consultez Quotas. Pour demander un quota de TPU v5e, contactez Cloud Ventes.
Créer un projet et un compte Google Cloud
Vous devez disposer d'un compte et d'un projet Google Cloud pour utiliser Cloud TPU. Pour plus consultez la page Configurer une instance Cloud TPU environnement.
Créer une instance Cloud TPU
Il est recommandé de provisionner les Cloud TPU v5es en file d'attente
des ressources à l'aide de l'queued-resource create
. Pour en savoir plus, consultez la section Gérer les fichiers
ressources.
Vous pouvez également utiliser l'API Create Node (gcloud compute tpus tpu-vm create
).
pour provisionner Cloud TPU
v5. Pour en savoir plus, consultez la page Gérer les ressources
ressources.
Pour en savoir plus sur les métriques pour les configurations v5e pour l'entraînement, consultez la section Types de Cloud TPU v5e pour entraînement.
Configuration du framework
Cette section décrit le processus général de configuration de l'entraînement de modèle personnalisé à l'aide de
JAX ou PyTorch avec TPU v5e. TensorFlow est compatible avec
TPU tpu-vm-tf-2.17.0-pjrt
et tpu-vm-tf-2.17.0-pod-pjrt
versions d'exécution.
Pour obtenir des instructions sur la configuration de l'inférence, consultez la page Inférence v5e présentation.
Configuration pour JAX
Si vous avez des formes de tranches supérieures à huit chips, vous aurez plusieurs VM en une.
tranche. Dans ce cas, vous devez utiliser l'option --worker=all
pour exécuter
l'installation sur toutes les VM TPU en une seule étape, sans utiliser SSH pour se connecter à chacune
séparément:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Description des options de commande
Variable | Description |
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. |
PROJECT_ID | Nom du projet Google Cloud. Utilisez un projet existant ou créez-en un. à <ph type="x-smartling-placeholder"></ph> Configurer votre projet Google Cloud |
ZONE | Voir les régions et zones TPU pour les zones compatibles. |
Worker [class name, see definition and ref site provided] | VM TPU ayant accès aux TPU sous-jacents. |
Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les résultats ont été produits ici, avec un segment "v5litepod-16"). Ce code vérifie que tout est installé correctement 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 doit ressembler à 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 la tranche donnée.
jax.local_device_count()
indique le nombre de chips accessibles par un seul
VM de 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 doit ressembler à 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.]
Essayez les tutoriels JAX de ce document pour commencer à Entraînement v5e à l'aide de JAX.
Configuration pour PyTorch
Notez que la version v5e n'est compatible qu'avec l'environnement d'exécution PJRT. et PyTorch 2.1+ utilisera PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU.
Cette section explique comment commencer à utiliser PJRT sur v5e avec PyTorch/XLA avec 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 -y
sudo apt-get install libomp5 -y
pip3 install mkl mkl-include
pip3 install tf-nightly tb-nightly tbp-nightly
pip3 install numpy
sudo apt-get install libopenblas-dev -y
pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Si vous recevez un message d'erreur lors de l'installation des roues pour torch
, torch_xla
ou
torchvision
"J'aime"
pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end
or semicolon (after name and no valid version specifier) torch==nightly+20230222
,
rétrograder votre version à l'aide de la commande suivante:
pip3 install setuptools==62.1.0
Exécuter un script avec PJRT
unset LD_PRELOAD
Voici un exemple utilisant un script Python pour effectuer un calcul sur une VM v5e:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker all \
--command='
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
unset LD_PRELOAD
export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'
Un résultat semblable à celui-ci doit s'afficher :
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
Essayez les tutoriels PyTorch de ce document pour commencer à Entraînement v5e à l'aide de PyTorch.
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session. Pour supprimer un ressource en file d'attente, supprimez la tranche, 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
Ces deux étapes peuvent également être utilisées pour
supprimer les demandes de ressources en file d’attente qui se trouvent dans
l'état FAILED
.
Exemples JAX/FLAX
Les sections suivantes décrivent des exemples d'entraînement des modèles JAX et FLAX sur TPU v5e.
Entraîner ImageNet sur v5e
Ce tutoriel explique comment entraîner ImageNet sur v5e à l'aide de fausses données d'entrée. Si Si vous souhaitez utiliser des données réelles, consultez le fichier README sur GitHub.
Prérequis
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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_UNTIL_DURATION=1d
-
gcloud 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 se trouvera l'état
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Lorsque l'état de QueuedResource est
ACTIVE
, la sortie doit ressembler à ce qui suit:state: ACTIVE
Installez la dernière version de JAX et de jaxlib:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Clonez le modèle ImageNet et installez la configuration requise correspondante:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'
Pour générer de fausses données, le modèle a besoin d'informations sur les dimensions le jeu de données. Ces informations peuvent être collectées à partir des métadonnées de l'ensemble de données ImageNet:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
Entraîner le modèle
Une fois toutes les étapes précédentes terminées, vous pouvez entraîner le modèle.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Modèles FLAX Hugging Face
Les modèles Hugging Face qui sont implémentés en FLAX à partir de sur Cloud TPU v5e. Cette section fournit des instructions pour exécuter modèles populaires.
Entraîner ViT sur Imagenette
Ce tutoriel vous montre comment entraîner le modèle Vision Transformer (ViT) de HuggingFace à l'aide de l'outil Fast AI Imagenette sur Cloud TPU v5e.
Le modèle ViT a été le premier à avoir entraîné avec succès un encodeur Transformer sur ImageNet, avec d'excellents résultats par rapport aux réseaux convolutifs. Pour plus consultez les ressources suivantes:
Prérequis
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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_UNTIL_DURATION=1d
-
gcloud 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 votre ressource en file d'attente possède l'état
ACTIVE
: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
, le résultat est semblable comme suit:state: ACTIVE
Installez JAX et sa bibliothèque:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Télécharger le dépôt Hugging Face et les conditions d'installation requises:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.17.0 && pip install -r examples/flax/vision/requirements.txt'
Téléchargez l'ensemble de données Imagenette:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
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='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Résultats de l'analyse comparative ViT
Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. La Le tableau suivant indique les débits avec différents types d'accélérateurs.
Type d'accélérateur | v5litepod-4 | V5litepod-16 | V5litepod-64 |
Époque | 3 | 3 | 3 |
Taille du lot global | 32 | 128 | 512 |
Débit (exemples/s) | 263,40 | 429,34 | 470,71 |
Entraînement Diffusion sur Pokémon
Ce tutoriel vous explique comment entraîner le modèle Stable Diffusion à partir de HuggingFace utilisant le Pokémon sur Cloud TPU v5e.
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
Configurez un bucket de stockage pour la sortie du modèle.
gcloud storage buckets create gs://your_bucket
--project=your_project
--location=us-west1
export GCS_BUCKET_NAME=your_bucketCréez des variables d'environnement
export GCS_BUCKET_NAME=your_bucket export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west1-c export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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 GCS_BUCKET_NAME Affichage dans la console Google Cloud -> Cloud Storage -> Buckets PROJECT_ID Nom du projet Google Cloud. Utilisez un projet existant ou créez-en un. à <ph type="x-smartling-placeholder"></ph> Configurer votre projet Google Cloud ACCELERATOR_TYPE Consultez les versions TPU correspondant à votre version de TPU. ZONE Voir les régions et zones TPU pour les zones compatibles. RUNTIME_VERSION 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. Par 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
ouspot
. Si aucune de ces options n'est spécifiée, le QUOTA_TYPE la valeur par défaut eston-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. -
gcloud 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 passe à l'état
ACTIVE
: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
, le résultat ressemble à ce qui suit:state: ACTIVE
Installez JAX et sa bibliothèque.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Télécharger le dépôt HuggingFace et les exigences d'installation.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.17.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
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} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
git clone https://github.com/google/maxdiffusion
cd maxdiffusion
git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip3 install -r requirements.txt
pip3 install .
export LIBTPU_INIT_ARGS=""
python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Résultats de l'analyse comparative de la diffusion
Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits.
Type d'accélérateur | v5litepod-4 | V5litepod-16 | V5litepod-64 |
Étape d'entraînement | 1500 | 1500 | 1500 |
Taille du lot global | 32 | 64 | 128 |
Débit (exemples/s) | 36,53 | 43,71 | 49,36 |
Entraîner GPT2 sur l'ensemble de données OSCAR
Ce tutoriel explique comment entraîner le modèle GPT2 depuis HuggingFace en utilisant l'ensemble de données OSCAR sur Cloud TPU v5e.
GPT2 est un modèle Transformer pré-entraîné sur du texte brut, l'étiquetage manuel. Il a été entraîné pour prédire le mot suivant dans des phrases. Pour en savoir plus, consultez les ressources suivantes :
Prérequis
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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
-
gcloud 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 votre ressource en file d'attente possède l'état
ACTIVE
: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
, le résultat est semblable comme suit:state: ACTIVE
Installez JAX et sa bibliothèque.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Télécharger le dépôt HuggingFace et les exigences d'installation.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'
Téléchargez les configurations pour entraîner le modèle.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'
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='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Résultats de l'analyse comparative pour GPT2
Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits.
v5litepod-4 | V5litepod-16 | V5litepod-64 | |
Époque | 3 | 3 | 3 |
Taille du lot global | 64 | 64 | 64 |
Débit (exemples/s) | 74,60 | 72,97 | 72,62 |
PyTorch/XLA
Les sections suivantes décrivent des exemples d'entraînement des modèles PyTorch/XLA sur TPU v5e.
Entraîner ResNet à l'aide de l'environnement d'exécution PJRT
PyTorch/XLA migre de XRT vers PjRt depuis PyTorch 2.0 et versions ultérieures. Voici les Mise à jour des instructions de configuration de v5e pour les charges de travail d'entraînement PyTorch/XLA.
Prérequis
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite export SERVICE_ACCOUNT=your_service_account export TPU_NAME=tpu-name export QUEUED_RESOURCE_ID=queued_resource_id export QUOTA_TYPE=quota_type export VALID_UNTIL_DURATION=1d
-
gcloud 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 votre QueuedResource possède l'état
ACTIVE
: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
, le résultat ressemble à ce qui suit:state: ACTIVE
Installer des dépendances spécifiques à Torch/XLA
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' sudo apt-get update -y sudo apt-get install libomp5 -y pip3 install mkl mkl-include pip3 install tf-nightly tb-nightly tbp-nightly pip3 install numpy sudo apt-get install libopenblas-dev -y pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Entraîner le modèle ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
date
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
export XLA_USE_BF16=1
export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
git clone https://github.com/pytorch/xla.git
cd xla/
git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
python3 test/test_train_mp_imagenet.py --model=resnet50 --fake_data --num_epochs=1 —num_workers=16 --log_steps=300 --batch_size=64 --profile'
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Résultat de l'analyse comparative
Le tableau suivant présente les débits de benchmark.
Type d'accélérateur | Débit (exemples/seconde) |
v5litepod-4 | 4240 ex/s |
V5litepod-16 | 10 810 ex/s |
V5litepod-64 | 46 154 ex/s |
Entraîner GPT2 sur v5e
Ce tutoriel explique comment exécuter GPT2 dans la version 5e à l'aide du dépôt HuggingFace sur PyTorch/XLA à l'aide de l'ensemble de données wikitext.
Prérequis
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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
-
gcloud 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 votre QueuedResource sera
ACTIVE
: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 doit ressembler à ce qui suit:state: ACTIVE
Installez les dépendances PyTorch/XLA.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' sudo apt-get -y update sudo apt install -y libopenblas-base pip3 install torchvision pip3 uninstall -y torch pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Télécharger le dépôt HuggingFace et les exigences d'installation.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' git clone https://github.com/pytorch/xla.git pip install --upgrade accelerate git clone https://github.com/huggingface/transformers.git cd transformers git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2 pip install . git log -1 pip install datasets evaluate scikit-learn '
Télécharger les configurations du modèle pré-entraîné
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
Entraîner le modèle
Entraînez le modèle 2B en utilisant une taille de lot de 16.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU_C_API
cd transformers/
export LD_LIBRARY_PATH=/usr/local/lib/
export PT_XLA_DEBUG=0
export USE_TORCH=ON
python3 examples/pytorch/xla_spawn.py \
--num_cores=4 \
examples/pytorch/language-modeling/run_clm.py \
--num_train_epochs=3 \
--dataset_name=wikitext \
--dataset_config_name=wikitext-2-raw-v1 \
--per_device_train_batch_size=16 \
--per_device_eval_batch_size=16 \
--do_train \
--do_eval \
--logging_dir=./tensorboard-metrics \
--cache_dir=./cache_dir \
--output_dir=/tmp/test-clm \
--overwrite_output_dir \
--cache_dir=/tmp \
--config_name=examples/pytorch/language-modeling/my_config_2.json \
--tokenizer_name=gpt2 \
--block_size=1024 \
--optim=adafactor \
--adafactor=true \
--save_strategy=no \
--logging_strategy=no \
--fsdp=full_shard \
--fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Résultat de l'analyse comparative
Le script d'entraînement a été exécuté sur les versions v5litepod-4, v5litepod-16 et v5litepod-64. La Le tableau suivant montre les débits de benchmark pour les différents types d'accélérateurs.
v5litepod-4 | V5litepod-16 | V5litepod-64 | |
Époque | 3 | 3 | 3 |
config | 600M | 2 milliards | 16 Mrds |
Taille du lot global | 64 | 128 | 256 |
Débit (exemples/s) | 66 | 77 | 31 |
Entraîner ViT sur v5e
Ce tutoriel explique comment exécuter VIT sur la version 5e à l'aide du dépôt HuggingFace. sur PyTorch/XLA sur l'ensemble de données cifar10.
Prérequis
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite export SERVICE_ACCOUNT=your_service_account export TPU_NAME=tpu-name export QUEUED_RESOURCE_ID=queued_resource_id export QUOTA_TYPE=quota_type export VALID_UNTIL_DURATION=1d
-
gcloud 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 votre QueuedResource sera l'état
ACTIVE
: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 doit ressembler à ce qui suit:state: ACTIVE
Installer des dépendances PyTorch/XLA
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all --command=' sudo apt-get update -y sudo apt-get install libomp5 -y pip3 install mkl mkl-include pip3 install tf-nightly tb-nightly tbp-nightly pip3 install numpy sudo apt-get install libopenblas-dev -y pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Télécharger le dépôt HuggingFace et les exigences d'installation.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=" git clone https://github.com/suexu1025/transformers.git vittransformers; \ cd vittransformers; \ pip3 install .; \ pip3 install datasets; \ wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
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_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
export TF_CPP_MIN_LOG_LEVEL=0
export XLA_USE_BF16=1
export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
cd vittransformers
python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
--remove_unused_columns=False \
--label_names=pixel_values \
--mask_ratio=0.75 \
--norm_pix_loss=True \
--do_train=true \
--do_eval=true \
--base_learning_rate=1.5e-4 \
--lr_scheduler_type=cosine \
--weight_decay=0.05 \
--num_train_epochs=3 \
--warmup_ratio=0.05 \
--per_device_train_batch_size=8 \
--per_device_eval_batch_size=8 \
--logging_strategy=steps \
--logging_steps=30 \
--evaluation_strategy=epoch \
--save_strategy=epoch \
--load_best_model_at_end=True \
--save_total_limit=3 \
--seed=1337 \
--output_dir=MAE \
--overwrite_output_dir=true \
--logging_dir=./tensorboard-metrics \
--tpu_metrics_debug=true'
Supprimer le TPU et la ressource en file d'attente
Supprimez votre TPU et la ressource en file d'attente à la fin de votre session.
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
Résultat de l'analyse comparative
Le tableau suivant indique les débits de benchmark pour les différents types d'accélérateurs.
v5litepod-4 | V5litepod-16 | V5litepod-64 | |
Époque | 3 | 3 | 3 |
Taille du lot global | 32 | 128 | 512 |
Débit (exemples/s) | 201 | 657 | 2 844 |
TensorFlow 2.x
Les sections suivantes décrivent des exemples d'entraînement de TensorFlow 2.x sur les TPU v5e.
Entraîner Resnet sur une version v5e à hôte unique
Ce tutoriel explique comment entraîner ImageNet sur v5litepod-4
ou v5litepod-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
Créez des variables d'environnement :
export PROJECT_ID=your-project-ID export ACCELERATOR_TYPE=v5litepod-4 export ZONE=us-east1-c export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt export TPU_NAME=your-tpu-name export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPE
peut êtrev5litepod-4
ouv5litepod-8
.-
gcloud 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 en SSH à votre VM TPU une fois que la ressource en file d'attente se trouvera 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}
Se connecter à votre TPU à l'aide de SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
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
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
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"
Supprimer le TPU et la ressource en file d'attente
Supprimer votre TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Supprimer votre requête de ressource en file d'attente
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Entraîner Resnet sur une v5e à hôtes multiples
Ce tutoriel explique comment entraîner ImageNet sur v5litepod-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.
Créez des variables d'environnement :
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-east1-c export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt export TPU_NAME=your_tpu_name export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPE
peut être de taillev5litepod-16
ou supérieure.-
gcloud 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 en SSH à votre VM TPU une fois que la ressource en file d'attente se trouvera l'état
ACTIVE
. Pour vérifier l'état d'une ressource en file d'attente, utilisez la méthode la commande suivante:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
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}
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 TPU_LOAD_LIBRARY=0 export TPU_NAME=your_tpu_name
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
Exécutez le script d'entraînement.
python3 official/vision/train.py \
--tpu=${TPU_NAME} \
--experiment=resnet_imagenet \
--mode=train_and_eval \
--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*"
Supprimer le TPU et la ressource en file d'attente
Supprimer votre TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Supprimer votre requête de ressource en file d'attente
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet