Cloud TPU v5e-Training

Cloud TPU v5e ist der KI-Beschleuniger der neuesten Generation von Google Cloud. Mit einem weniger 256 Chips pro Pod, V5e ist für den höchsten Wert optimiert Produkt für Transformer, Text-to-Image und Convolutional Neural Network (CNN) Training, Feinabstimmung und Bereitstellung. Weitere Informationen zur Verwendung von Cloud TPU v5e für die Bereitstellung finden Sie unter Inferenz mit v5e.

Weitere Informationen zur Hardware und Konfiguration von Cloud TPU v5e-TPUs finden Sie unter TPU v5e:

Mehr erfahren

In den folgenden Abschnitten werden die ersten Schritte mit TPU v5e beschrieben.

Anfragekontingent

Sie benötigen ein Kontingent, um TPU v5e für das Training nutzen zu können. Es gibt verschiedene Kontingenttypen für On-Demand-TPUs, reservierte TPUs und TPU-Spot-VMs. Es gibt separate Erforderliche Kontingente, wenn Sie TPU v5e für Inferenz. Weitere Informationen zu Kontingenten finden Sie unter Kontingente: Wenn Sie ein TPU v5e-Kontingent anfordern möchten, wenden Sie sich an Cloud Verkäufe.

Google Cloud-Konto und -Projekt erstellen

Sie benötigen ein Google Cloud-Konto und ein Projekt, um Cloud TPU verwenden zu können. Weitere Informationen Weitere Informationen finden Sie unter Cloud TPU einrichten Umgebung.

Hinweis: Zur Verwendung von TPU v5e müssen Sie das gcloud alpha-Paket installieren.

Cloud TPU erstellen

Es empfiehlt sich, Cloud TPU v5es in der Warteschlange bereitzustellen. Ressourcen mithilfe der queued-resource create . Weitere Informationen finden Sie unter In der Warteschlange verwalten Ressourcen

Sie können auch die Create Node API (gcloud alpha compute tpus tpu-vm create) verwenden zum Bereitstellen von Cloud TPU v5es. Weitere Informationen finden Sie unter TPU verwalten Ressourcen

Weitere Informationen zu verfügbaren v5e-Konfigurationen für das Training siehe Cloud TPU v5e-Typen für Schulung.

Framework-Einrichtung

In diesem Abschnitt wird der allgemeine Einrichtungsprozess für das Training benutzerdefinierter Modelle mit JAX oder PyTorch mit TPU v5e. TensorFlow-Unterstützung ist verfügbar in der tpu-vm-tf-2.17.0-pjrt und tpu-vm-tf-2.17.0-pod-pjrt TPU Laufzeitversionen.

Eine Anleitung zur Einrichtung der Inferenz finden Sie unter Inferenz der Version 5e. .

Einrichtung für JAX

Wenn Sie Segmentformen mit mehr als 8 Chips haben, haben Sie mehrere VMs in einem Segment. In diesem Fall müssen Sie das Flag --worker=all verwenden, um den Installation auf allen TPU-VMs in einem einzigen Schritt, ohne SSH für die Anmeldung in den einzelnen VMs zu verwenden. separat:

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'

Mit dem folgenden Befehl können Sie die Anzahl der Geräte (die angezeigten Ausgaben) prüfen. hier mit einem v5litepod-16-Slice produziert wurden. Mit diesem Code wird getestet, durch Überprüfung, ob JAX die Cloud TPU TensorCores erkennt, und kann grundlegende Vorgänge ausführen:

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())"'

Die Ausgabe sollte in etwa so aussehen:

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() zeigt die Gesamtzahl der Chips im jeweiligen Segment an. jax.local_device_count() gibt die Anzahl der Chips an, auf die für einen einzelnen Nutzer zugegriffen werden kann. VM in diesem Slice.

# 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))"'

Die Ausgabe sollte in etwa so aussehen:

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

Probieren Sie die JAX-Anleitungen in diesem Dokument aus, um v5e-Training mit JAX.

PyTorch einrichten

Beachten Sie, dass Version 5e nur die PJRT-Laufzeit unterstützt. und PyTorch 2.1+ verwenden PJRT als Standardlaufzeit für alle TPU-Versionen.

In diesem Abschnitt wird beschrieben, wie Sie PJRT in Version 5e mit PyTorch/XLA mit für alle Worker ausführen.

Abhängigkeiten installieren

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'

Wenn beim Einbau der Räder für torch, torch_xla oder torchvision „Gefällt mir“-Angabe pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, Führen Sie mit dem folgenden Befehl ein Downgrade Ihrer Version durch:

pip3 install setuptools==62.1.0

Skript mit PJRT ausführen

Hinweis: Verwenden Sie für Modelle mit umfangreichen, häufigen Zuweisungen tcmalloc. die Trainingszeit im Vergleich zum Standard-malloc erheblich verkürzen kann Implementierung, sodass der Standard-malloc auf einer TPU-VM tcmalloc Abhängig von Ihrer Arbeitslast (z. B. mit DLRM mit sehr großen Zuweisungen für seine Einbettungstabellen) tcmalloc kann zu einer Verlangsamung führen. In diesem Fall könnten Sie versuchen, die Festlegung der folgenden Variablen Verwenden Sie stattdessen die Standardeinstellung malloc:

unset LD_PRELOAD

Im Folgenden finden Sie ein Beispiel, in dem ein Python-Skript für eine Berechnung verwendet wird auf einer v5e-VM:

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)"'

Dadurch wird eine Ausgabe generiert, die etwa so aussieht:

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')

Probieren Sie die PyTorch-Anleitungen in diesem Dokument aus, um zu beginnen, v5e-Training mit PyTorch.

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung. So löschen Sie in die Warteschlange gestellt, löschen Sie das Segment und dann die Ressource in der Warteschlange in zwei Schritten:

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

Mit diesen beiden Schritten können auch Ressourcenanfragen in der Warteschlange entfernt werden, die sich in FAILED.

JAX/FLAX-Beispiele

In den folgenden Abschnitten werden Beispiele für das Trainieren von JAX- und FLAX-Modellen auf TPU v5e.

ImageNet in Version 5e trainieren

In diesem Tutorial wird beschrieben, wie ImageNet in v5e mit fiktiven Eingabedaten trainiert wird. Wenn Wenn Sie echte Daten verwenden möchten, finden Sie entsprechende Informationen in der README-Datei auf GitHub.

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald sich die Ressource in der Warteschlange befindet. dem Status ACTIVE:

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

   Wenn sich die QueuedResource im Status ACTIVE befindet, wird die Ausgabe etwa so aussehen:

    state: ACTIVE
   

3. Installieren Sie die neueste Version von JAX und jaxlib:

   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'
   

   Hinweis: Wenn Sie in einer Fehlermeldung darauf hingewiesen werden, jax>=0.4.8 zu verwenden, können Sie diese der Nachricht.

4. Klonen Sie das ImageNet-Modell und installieren Sie die entsprechenden Anforderungen:

   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'
   

5. Um fiktive Daten zu generieren, benötigt das Modell Informationen zu den Dimensionen des Datasets. Sie können sie aus den Metadaten des ImageNet-Datasets ermitteln:

   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'
   

Modell trainieren

Sobald alle vorherigen Schritte abgeschlossen sind, können Sie das Modell trainieren.

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'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

Hugging Face FLAX-Modelle

Hugging Face-Modelle, die in FLAX implementiert sind, in Cloud TPU v5e. Dieser Abschnitt enthält Anweisungen zum Ausführen beliebten Modellen.

ViT auf Imagenette trainieren

In dieser Anleitung erfahren Sie, wie Sie Vision Transformer trainieren. (ViT)-Modell von HuggingFace mit der schnellen KI-Technologie von Imagenette Dataset in Cloud TPU v5e.

Das ViT-Modell war das erste, bei dem ein Transformer-Encoder erfolgreich trainiert wurde auf ImageNet mit hervorragenden Ergebnissen im Vergleich zu Convolutional Networks. Weitere Informationen finden Sie in den folgenden Ressourcen:

• ViT-Übersicht
• ViT-Coderessource

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald die Ressource in der Warteschlange steht. befindet sich im Status ACTIVE:

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

   Wenn sich die Ressource in der Warteschlange im Status ACTIVE befindet, sieht die Ausgabe ähnlich aus. in:

    state: ACTIVE
   

3. Installieren Sie JAX und die zugehörige Bibliothek:

   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'
   

4. Hugging Face-Repository herunterladen und Installationsvoraussetzungen:

   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'
   

5. Laden Sie das Imagenette-Dataset herunter:

   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'
   

Modell trainieren

Trainieren Sie das Modell mit einem vorab zugeordneten Zwischenspeicher mit 4 GB.

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'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

ViT-Benchmarking-Ergebnisse

Das Trainingsskript wurde auf v5litepod-4, v5litepod-16 und v5litepod-64 ausgeführt. Die In der folgenden Tabelle sind die Durchsatzraten bei verschiedenen Beschleunigertypen aufgeführt.

Beschleunigertyp	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
Globale Batchgröße	32	128	512
Durchsatz (Beispiele/s)	263,40	429,34	470,71

Trainiere die Streuung von Pokémon

In dieser Anleitung erfahren Sie, wie Sie das Modell HuggingFace nutzt Pokémon Dataset in Cloud TPU v5e.

Das Stable Diffusion-Modell ist ein latentes Text-zu-Bild-Modell, das fotorealistische Bilder aus jeder Texteingabe. Weitere Informationen finden Sie in den folgenden Ressourcen:

• Stabile Diffusion – Übersicht
• Codequelle für stabile Diffusion

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald sich die Ressource in der Warteschlange im Status ACTIVE befindet:

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

   Wenn die Ressource in der Warteschlange den Status ACTIVE hat, sieht die Ausgabe in etwa so aus:

    state: ACTIVE
   

3. Installieren Sie JAX und die zugehörige Bibliothek.

   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'
   

4. HuggingFace-Repository herunterladen und Installationsvoraussetzungen.

   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'
   

Modell trainieren

Trainieren Sie das Modell mit einem vorab zugeordneten Zwischenspeicher mit 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

Benchmarking-Ergebnisse für Diffusion

Das Trainingsskript wurde auf v5litepod-4, v5litepod-16 und v5litepod-64 ausgeführt. In der folgenden Tabelle sind die Durchsatzraten aufgeführt.

Beschleunigertyp	v5litepod-4	v5litepod-16	v5litepod-64
Trainingsschritt	1500	1500	1500
Globale Batchgröße	32	64	128
Durchsatz (Beispiele/s)	36,53	43,71	49,36

GPT2 im OSCAR-Dataset trainieren

In dieser Anleitung erfahren Sie, wie Sie das GPT2 aus HuggingFace mit dem Dataset OSCAR in Cloud TPU v5e trainieren.

GPT2 ist ein Transformer-Modell, das mit Rohtexten ohne Beschriftungen durch Menschen. Es wurde darauf trainiert, das nächste Wort in Sätzen vorherzusagen. Weitere Informationen finden Sie in den folgenden Ressourcen:

• GPT2-Übersicht
• GPT2-Codequelle

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald die Ressource in der Warteschlange steht. hat den Status ACTIVE:

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

   Wenn sich die Ressource in der Warteschlange im Status ACTIVE befindet, sieht die Ausgabe ähnlich aus. in:

    state: ACTIVE
   

3. Installieren Sie JAX und die zugehörige Bibliothek.

   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'
   

4. HuggingFace-Repository herunterladen und Installationsvoraussetzungen.

   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'
   

5. Konfigurationen zum Trainieren des Modells herunterladen.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
   

Modell trainieren

Trainieren Sie das Modell mit einem vorab zugeordneten Zwischenspeicher mit 4 GB.

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'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

Benchmarking-Ergebnisse für GPT2

Das Trainingsskript wurde auf v5litepod-4, v5litepod-16 und v5litepod-64 ausgeführt. In der folgenden Tabelle sind die Durchsatzraten aufgeführt.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
Globale Batchgröße	64	64	64
Durchsatz (Beispiele/s)	74,60	72,97	72,62

PyTorch/XLA

In den folgenden Abschnitten wird anhand von Beispielen beschrieben, wie PyTorch/XLA-Modelle auf TPU v5e.

ResNet mit der PJRT-Laufzeit trainieren

PyTorch/XLA wird von PyTorch 2.0+ von XRT zu PjRt migriert. Hier sind die aktualisierte Anleitung zum Einrichten von v5e für PyTorch/XLA-Trainingsarbeitslasten.

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald Ihre QueuedResource befindet sich im Status ACTIVE:

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

   Wenn die Ressource in der Warteschlange den Status ACTIVE hat, sieht die Ausgabe in etwa so aus:

    state: ACTIVE
   

3. Torch-/XLA-spezifische Abhängigkeiten installieren

   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'
   

ResNet-Modell trainieren

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'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

Benchmarkergebnis

In der folgenden Tabelle sind die Benchmark-Durchsätze aufgeführt.

Beschleunigertyp	Durchsatz (Beispiele/Sekunde)
v5litepod-4	4240 Ex/s
v5litepod-16	10.810 ex/s
v5litepod-64	46.154 ex/s

GPT2 in Version 5e trainieren

In dieser Anleitung erfahren Sie, wie Sie GPT2 in Version 5e mit dem HuggingFace-Repository ausführen. zu PyTorch/XLA mit dem Wikitext-Dataset

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich Ihre QueuedResource im ACTIVE-Bundesstaat:

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

   Wenn die Ressource in der Warteschlange den Status ACTIVE hat, lautet die Ausgabe: etwa so aussehen:

   state: ACTIVE
   

3. PyTorch/XLA-Abhängigkeiten installieren

   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'
   

4. HuggingFace-Repository herunterladen und Installationsvoraussetzungen.

   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
         '
   

5. Konfigurationen des vortrainierten Modells herunterladen.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
         gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
   

Modell trainieren

Trainieren Sie das 2B-Modell mit einer Batchgröße von 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'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

Benchmarkergebnis

Das Trainingsskript wurde auf v5litepod-4, v5litepod-16 und v5litepod-64 ausgeführt. Die Die folgende Tabelle zeigt die Benchmark-Durchsätze für verschiedene Beschleunigertypen.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
config	600 Mio.	2 Milliarden	16 Mrd.
Globale Batchgröße	64	128	256
Durchsatz (Beispiele/s)	66	77	31

ViT auf Version 5e trainieren

In dieser Anleitung erfahren Sie, wie Sie VIT in v5e mit dem HuggingFace-Repository ausführen. mit PyTorch/XLA im cifar10-Dataset.

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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
   

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich Ihre QueuedResource im dem Status ACTIVE:

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

   Wenn die Ressource in der Warteschlange den Status ACTIVE hat, lautet die Ausgabe: etwa so aussehen:

    state: ACTIVE
   

3. PyTorch-/XLA-Abhängigkeiten installieren

   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'
   

4. HuggingFace-Repository herunterladen und Installationsvoraussetzungen.

      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"
   

Modell trainieren

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'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die in die Warteschlange gestellten Ressourcen am Ende der Sitzung.

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

Benchmarkergebnis

Die folgende Tabelle zeigt die Benchmark-Durchsätze für verschiedene Beschleunigertypen.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
Globale Batchgröße	32	128	512
Durchsatz (Beispiele/s)	201	657	2.844

TensorFlow 2.x

In den folgenden Abschnitten werden Beispiele für das Training von TensorFlow 2.x beschrieben. auf TPU v5e-Modellen.

Resnet auf einem v5e-Host mit einem Host trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet mit v5litepod-4 oder v5litepod-8 trainieren mithilfe eines fiktiven Datasets. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie Dataset vorbereiten

Einrichten

1. Erstellen Sie Umgebungsvariablen:

   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 kann entweder v5litepod-4 oder v5litepod-8 sein.

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald sich die Ressource in der Warteschlange befindet. ACTIVE. So prüfen Sie den Status der Ressource in der Warteschlange: verwenden Sie den folgenden Befehl:

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

3. Über SSH eine Verbindung zur TPU herstellen

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

4. Umgebungsvariablen festlegen

   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. Wechseln Sie zum Verzeichnis des Modell-Repositorys und installieren Sie die Installationsvoraussetzungen.

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

Modell trainieren

Führen Sie das Trainingsskript aus.

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"

TPU und Ressource in der Warteschlange löschen

1. TPU löschen

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

2. Ressourcenanfrage in der Warteschlange löschen

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

Resnet auf einem v5e mit mehreren Hosts trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet auf v5litepod-16 oder größer mit ein fiktives Dataset. Wenn Sie ein anderes Dataset verwenden möchten, finden Sie weitere Informationen unter Dataset vorbereiten.

1. Erstellen Sie Umgebungsvariablen:

   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 kann entweder v5litepod-16 oder größer sein.

2. TPU-Ressource erstellen:

   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}
   

   Hinweis: Das Flag QUOTA_TYPE kann reserved sein. oder best-effort. Wenn keine dieser Optionen angegeben ist, wird der Standardwert on-demand Siehe Kontingente finden Sie Informationen zu den verschiedenen Arten von Kontingenten, die von Cloud TPU unterstützt werden.

   Sie können eine SSH-Verbindung zur TPU-VM herstellen, sobald sich die Ressource in der Warteschlange befindet. ACTIVE. Prüfen Sie den Status Ihrer in die Warteschlange gestellten Ressource mithilfe des folgenden Befehl:

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

3. Über SSH eine Verbindung zu Ihrer TPU (Worker Zero) herstellen

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

4. Umgebungsvariablen festlegen

   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
   

5. Wechseln Sie zum Verzeichnis des Modell-Repositorys und installieren Sie die Installationsvoraussetzungen.

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

Modell trainieren

Führen Sie das Trainingsskript aus.

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

TPU und Ressource in der Warteschlange löschen

1. TPU löschen

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

2. Ressourcenanfrage in der Warteschlange löschen

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