Cloud TPU v5e-Training

Cloud TPU v5e ist der KI-Beschleuniger der neuesten Generation von Google Cloud. Mit einem geringeren Platzbedarf von 256 Chips pro Pod ist ein v5e als Produkt mit dem höchsten Wert für das Training, die Feinabstimmung und die Bereitstellung von Transformern, Text-zu-Bild- und Convolutional Neural Network (CNN) optimiert.

Konzepte, Systemarchitektur und Konfigurationen von Cloud TPU v5e

Wenn Sie Cloud TPUs noch nicht kennen, finden Sie weitere Informationen auf der Startseite der TPU-Dokumentation.

Allgemeine Cloud TPU-Konzepte (z. B. Slices, Hosts, Chips und TensorCores) sowie die Cloud TPU-Systemarchitektur werden auf der Seite Cloud TPU-Systemarchitektur beschrieben.

Jede Cloud TPU-Version erfordert bestimmte Beschleunigertypen für Training und Inferenz. Diese Beschleunigertypen werden im Abschnitt v5e-Konfigurationen beschrieben.

Inferenz

Inferenz ist der Vorgang, bei dem ein trainiertes Modell verwendet wird, um Vorhersagen für neue Daten zu treffen. Sie wird im Bereitstellungsprozess verwendet.

Slices

Ein Slice stellt eine Sammlung von Chips dar, die sich alle im selben Pod befinden und über Highspeed-Inter-Chip Interconnect-Verbindungen (ICI) verbunden sind. v5e hat 2D-Slice-Formen. Unterstützte Segmentformen finden Sie in der Tabelle im Abschnitt V5e-Konfigurationen.

Chipform und Chiptopologie beziehen sich auch auf Segmentformen.

Bereitstellung

Bereitstellung ist der Prozess der Bereitstellung eines trainierten Modells für maschinelles Lernen in einer Produktionsumgebung, in der es Vorhersagen oder Entscheidungen treffen kann. Latenz und Verfügbarkeit auf Dienstebene sind wichtig für die Bereitstellung.

Einzelner Host oder mehrere Hosts

Ein Host ist ein physischer Computer (CPU), auf dem VMs ausgeführt werden. Ein Host kann mehrere VMs gleichzeitig ausführen.

Slices mit weniger als 8 Chips verwenden höchstens einen Host. Slices mit mehr als 8 Chips haben Zugriff auf mehr als einen Host und können verteiltes Training mit mehreren Hosts ausführen. Weitere Informationen zu Slices und Chips finden Sie auf der Seite TPU-Systemarchitektur.

v5e unterstützt das Training mit mehreren Hosts und die Inferenz mehrerer Hosts (mit SAX).

TPU-VM

Eine virtuelle Maschine, auf der Linux ausgeführt wird und die Zugriff auf die zugrunde liegenden TPUs hat. Bei v5e-TPUs hat jede TPU-VM je nach vom Nutzer angegebenen Beschleunigertyp direkten Zugriff auf 1, 4 oder 8 Chips. Eine TPU-VM wird auch als Worker bezeichnet.

Worker

Siehe TPU-VM.

Jetzt starten

Informationen zur v5e-TPU-Hardware finden Sie unter Systemarchitektur.

Kapazität sichern

Wenden Sie sich an den Cloud-Vertrieb, um mit der Verwendung von Cloud TPU v5e für Ihre KI-Arbeitslasten zu beginnen.

Google Cloud-Projekt vorbereiten

  1. Melden Sie sich bei Ihrem Google-Konto an. Falls noch nicht geschehen, registrieren Sie sich für ein neues Konto.
  2. In der Google Cloud Console können Sie auf der Seite für die Projektauswahl ein Google Cloud-Projekt auswählen oder erstellen.

  3. Für die gesamte Google Cloud-Nutzung ist eine Abrechnungseinrichtung erforderlich. Achten Sie also darauf, dass die Abrechnung für Ihr Projekt aktiviert ist.

  4. Installieren Sie gcloud alpha components.

  5. Wenn Sie ein TPU-Nutzer sind, der vorhandene gcloud alpha-Komponenten wiederverwendet, aktualisieren Sie diese, damit die relevanten Befehle und Flags unterstützt werden:

    gcloud components update

  6. Aktivieren Sie die TPU API über den folgenden gcloud-Befehl in Cloud Shell. Sie können sie auch über die Google Cloud Console aktivieren.

    gcloud services enable tpu.googleapis.com

  7. Aktivieren Sie das TPU-Dienstkonto.

    Dienstkonten ermöglichen dem Cloud TPU-Dienst den Zugriff auf andere Google Cloud-Dienste. Ein nutzerverwaltetes Dienstkonto wird für Google Cloud empfohlen. Folgen Sie diesen Leitfäden, um Rollen zu erstellen und zu gewähren. Die folgenden Rollen sind erforderlich:

    • TPU-Administrator
    • Storage-Administrator: für den Zugriff auf Cloud Storage erforderlich
    • Logautor: erforderlich zum Schreiben von Logs mit der Logging API
    • Monitoring-Messwert-Autor: erforderlich, um Messwerte in Cloud Monitoring zu schreiben

  8. Konfigurieren Sie das Projekt und die Zone.

    Die Projekt-ID ist der Name Ihres Projekts, der in der Cloud Console angezeigt wird.

    export PROJECT_ID=your-project-id
export ZONE=us-west4-a

gcloud compute tpus tpu-vm service-identity create --zone=${ZONE}

gcloud auth login
gcloud config set project ${PROJECT_ID}
gcloud config set compute/zone ${ZONE}

Cloud TPU-Umgebung bereitstellen

Die Best Practice besteht darin, Cloud TPU v5es als Ressourcen in der Warteschlange mit dem Befehl queued-resource create bereitzustellen. Sie können jedoch auch die Create Node API (gcloud alpha compute tpus tpu-vm create) verwenden, um Cloud TPU v5es bereitzustellen.

Umgebungsvariablen erstellen

Legen Sie die erforderlichen Umgebungsvariablen für die TPU-Erstellung fest.

Ersetzen Sie die Variablen (in rot) in der folgenden Liste durch Werte, die Sie für Ihren Trainings- oder Inferenzjob verwenden.

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
Variable Beschreibung
PROJECT_ID Name des Google Cloud-Projekts
ACCELERATOR_TYPE Unterstützte Beschleunigertypen finden Sie im Abschnitt Beschleunigertypen.
ZONE Die gesamte Kapazität ist in us-west4-a angegeben.
RUNTIME_VERSION Verwenden Sie v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT Dies ist die Adresse Ihres Dienstkontos. Sie finden sie in der Google Cloud Console unter „IAM“ -> „Dienstkonten“. Beispiel: [email protected]
TPU_NAME Die vom Nutzer zugewiesene Text-ID der TPU, die bei der Zuweisung der Ressourcenanfrage in der Warteschlange erstellt wird.
QUEUED_RESOURCE_ID Die vom Nutzer zugewiesene Text-ID der Ressourcenanfrage in der Warteschlange. Informationen zu Ressourcen in der Warteschlange finden Sie unter Ressourcen in der Warteschlange.
QUOTA_TYPE Kann reserved oder best-effort sein. Wenn keine dieser Optionen angegeben ist, wird standardmäßig on-demand als QUOTA_TYPE festgelegt. Informationen zu den verschiedenen Kontingenttypen, die von Cloud TPU unterstützt werden, finden Sie unter quotas.
VALID_UNTIL_DURATION Die Dauer, für die die Anfrage gültig ist. Informationen zu den verschiedenen gültigen Daueren finden Sie unter Ressourcen in der Warteschlange.

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}

Wenn die Ressource in der Warteschlange erfolgreich erstellt wurde, ist der Status im Feld response entweder WAITING_FOR_RESOURCES oder FAILED. Wenn die Ressource in der Warteschlange den Status WAITING_FOR_RESOURCES hat, bedeutet dies, dass die Ressource in der Warteschlange die vorläufige Validierung bestanden hat und auf Kapazität wartet. Sobald Kapazität verfügbar ist, geht die Anfrage auf PROVISIONING über. Der Status WAITING_FOR_RESOURCES bedeutet nicht zwangsläufig, dass Ihnen das Kontingent zugewiesen wird. Es kann einige Zeit dauern, bis der Status von WAITING_FOR_RESOURCES zu ACTIVE wechselt. Wenn die Ressource in der Warteschlange den Status FAILED hat, wird der Fehlergrund in der Ausgabe angegeben. Die Anfrage läuft ab, wenn --valid-until-duration nicht ausgefüllt wird und der Status den Status "FAILED" hat.

Sobald die Ressource in der Warteschlange den Status ACTIVE hat, können Sie mit SSH auf Ihre TPU-VM zugreifen.

Verwenden Sie die Befehle [list](/tpu/docs/managing-tpus-tpu-vm) oder [describe](/tpu/docs/managing-tpus-tpu-vm), um den Status der Ressource in der Warteschlange abzufragen.

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

Der state stellt den Status einer Ressource in der Warteschlange dar. Die Bundesstaaten werden folgendermaßen definiert:

Bundesland Beschreibung
WAITING_FOR_RESOURCES Der Befehl create für die Ressource in der Warteschlange wurde empfangen und beginnt mit der Bereitstellung, sobald Kapazität verfügbar ist.
PROVISIONING Die TPU-Slices werden bereitgestellt.
ACTIVE Alle TPUs sind bereitgestellt und einsatzbereit. Ist ein Startskript angegeben, beginnt es auf allen TPUs mit der Ausführung, sobald der Ressourcenstatus in der Warteschlange in ACTIVE übergeht.
FAILED Die Segmente konnten nicht bereitgestellt werden.
SUSPENDING Mindestens ein Segment wird gelöscht.
SUSPENDED Alle zugrunde liegenden Slices werden gelöscht, aber die Ressource in der Warteschlange bleibt intakt, bis sie explizit gelöscht wird. Derzeit kann eine ausgesetzte Ressource in der Warteschlange nicht fortgesetzt werden und sollte gelöscht werden.
DELETING Die Ressource in der Warteschlange wird gelöscht.

Über SSH eine Verbindung zur TPU-VM herstellen

Im folgenden Abschnitt wird beschrieben, wie Sie Binärprogramme auf jeder TPU-VM in Ihrem TPU-Slice installieren und Code ausführen. Eine TPU-VM wird in diesem Kontext auch als Worker bezeichnet.

Im Abschnitt VM-Typen erfahren Sie, wie viele VMs das Slice enthalten wird.

Stellen Sie mit dem ssh-Befehl tpu-vm eine Verbindung zu Ihrer TPU-VM her, um die Binärdateien zu installieren oder Code auszuführen.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Für den Zugriff auf eine bestimmte TPU-VM mit SSH verwenden Sie das Flag --worker, das einem 0-basierten Index folgt:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Wenn die Kreisformen größer als 8 Chips sind, enthält ein Slice mehrere VMs. Verwenden Sie in diesem Fall das Flag --worker=all, um die Installation auf allen TPU-VMs auszuführen, ohne separat eine Verbindung zu jeder einzelnen herstellen zu müssen. Beispiel:

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

Verwalten

Alle Befehle, die Sie zum Verwalten Ihrer TPU-VMs verwenden können, werden unter TPUs verwalten beschrieben.

Framework-Einrichtung

In diesem Abschnitt wird der allgemeine Einrichtungsprozess für das benutzerdefinierte Modelltraining mit JAX oder PyTorch und TPU v5e beschrieben. TensorFlow-Unterstützung ist in den TPU-Laufzeitversionen tpu-vm-tf-2.15.0-pjrt und tpu-vm-tf-2.15.0-pod-pjrt verfügbar.

Eine Anleitung zum Einrichten von Inferenzen finden Sie unter Einführung in v5e-Inferenzen.

Einrichtung für JAX

Wenn die Kreisformen größer als 8 Chips sind, enthält ein Slice mehrere VMs. In diesem Fall müssen Sie das Flag --worker=all verwenden, um die Installation auf allen TPU-VMs in einem einzigen Schritt auszuführen, ohne sich mit SSH separat bei jeder anzumelden:

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

Mit dem folgenden Befehl können Sie die Anzahl der Geräte prüfen. Die hier gezeigten Ausgaben wurden mit einem v5litepod-16-Slice erstellt. Mit diesem Code wird getestet, ob alles korrekt installiert ist. Dazu wird geprüft, ob JAX die Cloud TPU TensorCores erkennt und grundlegende Vorgänge ausführen kann:

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() gibt die Gesamtzahl der Chips im jeweiligen Segment an. jax.local\_device\_count() gibt die Anzahl der Chips an, auf die eine einzelne VM in diesem Slice zugreifen kann.

# 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 mit dem v5e-Training mit JAX zu beginnen.

Einrichtung für PyTorch

Beachten Sie, dass v5e nur die PJRT-Laufzeit unterstützt und PyTorch 2.1 und höher als Standardlaufzeit für alle TPU-Versionen PJRT verwendet.

In diesem Abschnitt wird beschrieben, wie Sie PJRT unter v5e mit PyTorch/XLA mit Befehlen für alle Worker verwenden.

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 http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

Wenn Sie die Räder für torch/torch_xla/torchvision nicht installieren konnten und ein Fehler wie pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 angezeigt wird, stufen Sie Ihre Version mit diesem Befehl herunter:

pip3 install setuptools==62.1.0

Führen Sie ein Skript mit PJRT aus:

unset LD_PRELOAD

Im folgenden Beispiel wird ein Python-Skript für eine Berechnung auf einer v5e-VM verwendet:

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

Die PyTorch-Anleitungen in diesem Dokument bieten einen Einstieg in das V5e-Training mit PyTorch.

Löschen Sie Ihre TPU und die Ressource in der Warteschlange am Ende der Sitzung. Zum Löschen einer Ressource in der Warteschlange 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

Diese beiden Schritte können auch verwendet werden, um Ressourcenanfragen in der Warteschlange mit dem Status FAILED zu entfernen.

Überwachung und Profil

Cloud TPU v5e unterstützt Monitoring und Profilerstellung mit denselben Methoden wie vorherige Generationen von Cloud TPU. Weitere Informationen zum Monitoring finden Sie unter Profilerstellung für Ihr Modell mit Cloud TPU-Tools und weitere Informationen zur Profilerstellung unter Cloud TPU-VMs überwachen.

JAX/FLAX-Beispiele

ImageNet in v5e trainieren

In dieser Anleitung wird beschrieben, wie ImageNet in v5e mit fiktiven Eingabedaten trainiert wird. Wenn Sie echte Daten verwenden möchten, lesen Sie die 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}

    Sie können eine SSH-Verbindung zu Ihrer 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 QueuedResource den Status ACTIVE hat, sieht die Ausgabe in etwa so aus:

     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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  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 http://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. Zum Generieren falscher Daten benötigt das Modell Informationen zu den Dimensionen des Datasets. Diese finden Sie in den Metadaten des ImageNet-Datasets:

    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 http://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 Ihre TPU und die Ressource in der Warteschlange 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

FLAX-Modelle, die das Gesicht umarmen

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

ViT auf Imagenette trainieren

In dieser Anleitung erfahren Sie, wie Sie das Modell Vision Transformer (ViT) von HuggingFace mithilfe des Datasets Fast AI Imagenette in Cloud TPU v5e trainieren.

Das ViT-Modell war das erste Modell, das einen Transformer-Encoder auf ImageNet erfolgreich trainierte und im Vergleich zu Convolutional Networks hervorragende Ergebnisse lieferte. Weitere Informationen finden Sie in den folgenden Ressourcen:

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}

    Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich Ihre 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Laden Sie das Repository für das Hugging Face herunter und installieren Sie die Anforderungen:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://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.16.1 && 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 http://s3.amaz888onaws.com.fqhub.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

Modell trainieren

Trainieren Sie das Modell mit einem vorab zugeordneten Zwischenspeicher bei 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 Ihre TPU und die in der Warteschlange stehende Ressource 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 folgende Tabelle zeigt die Durchsätze mit verschiedenen Beschleunigertypen.

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

Trainiere Diffusion in Pokémon

In dieser Anleitung erfahren Sie, wie Sie das Stable Diffusion-Modell von HuggingFace mit dem Pokémon-Dataset in Cloud TPU v5e trainieren.

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

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}

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

    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 entsprechende 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Laden Sie das HuggingFace-Repository und die Installationsanforderungen herunter.

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

Modell trainieren

Trainieren Sie das Modell mit einem vorab zugeordneten Zwischenspeicher bei 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 Ressourcen in der Warteschlange löschen

Löschen Sie Ihre TPU und die Ressource in der Warteschlange 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 lief auf v5litepod-4, v5litepod-16 und v5litepod-64. Die folgende Tabelle zeigt die Durchsätze.

Beschleunigertyp v5litepod-4 v5litepod-16 v5litepod-64
Schritt trainieren 1.500 1.500 1.500
Globale Batchgröße 32 64 128
Durchsatz (Beispiele/Sek.) 36,53 43,71 49,36

GPT2 im OSCAR-Dataset trainieren

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

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

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}

    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 entsprechende 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Laden Sie das HuggingFace-Repository und die Installationsanforderungen herunter.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://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. Laden Sie Konfigurationen zum Trainieren des Modells herunter.

    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 bei 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 Ihre TPU und die in der Warteschlange stehende Ressource 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 lief auf v5litepod-4, v5litepod-16 und v5litepod-64. Die folgende Tabelle zeigt die Durchsätze.

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

PyTorch/XLA

ResNet mit der PJRT-Laufzeit trainieren

PyTorch/XLA wird von PyTorch/XLA von PyTorch 2.0 oder höher von XRT zu PjRt migriert. Hier finden Sie die aktualisierten Anweisungen zum Einrichten von v5e für PyTorch/XLA-Trainingsarbeitslasten.

Einrichtung

  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}

    Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die QueuedResource den Status ACTIVE hat:

    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 http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://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 http://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 Ihre TPU und die in der Warteschlange stehende Ressource 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
Benchmark-Ergebnis

Die folgende Tabelle zeigt die Benchmark-Durchsätze.

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 wird beschrieben, wie GPT2 in Version 5e über das HuggingFace-Repository für PyTorch/XLA und das Wikitext-Dataset ausgeführt wird.

Einrichtung

  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}

    Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die QueuedResource den Status ACTIVE hat:

    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 die Torch-/xla-Abhängigkeiten.

    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 http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

  4. Laden Sie das HuggingFace-Repository herunter und installieren Sie die Anforderungen.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone http://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone http://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. Laden Sie Konfigurationen des vortrainierten Modells herunter.

    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 Ihre TPU und die in der Warteschlange stehende Ressource 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

Benchmark-Ergebnis

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

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

ViT in v5e trainieren

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

Einrichtung

  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}

    Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die QueuedResource den 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. Torch-/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 http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

  4. Laden Sie das HuggingFace-Repository herunter und installieren Sie die Anforderungen.

       gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="
      git clone http://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget http://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 Ihre TPU und die in der Warteschlange stehende Ressource 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

Benchmark-Ergebnis

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

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

TensorFlow 2.x

Resnet auf einem einzelnen Host v5e trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet auf v5litepod-4 oder v5litepod-8 mit einem fiktiven Dataset trainieren. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie den Abschnitt 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}

    Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die Ressource in der Warteschlange im Status ACTIVE befindet. Verwenden Sie den folgenden Befehl, um den Status der Ressource in der Warteschlange zu prüfen:

    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 Repository-Verzeichnis des Modells und installieren Sie die Anforderungen.

    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-System mit mehreren Hosts trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet auf v5litepod-16 oder höher mit einem fiktiven Dataset trainieren. Wenn Sie ein anderes Dataset verwenden möchten, lesen Sie 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}

    Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, sobald sich die Ressource in der Warteschlange im Status ACTIVE befindet. Verwenden Sie den folgenden Befehl, um den Status der Ressource in der Warteschlange zu prüfen:

    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 Repository-Verzeichnis des Modells und installieren Sie die Anforderungen.

     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