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 v5e das umsatzstärkste Produkt für Transformer-, Text-to-Image- und CNN-Training, -Feinabstimmung und -Bereitstellung.

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. Segmente, 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 auf der Seite TPU-Konfigurationen beschrieben.

Inferenz

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

Slices

Ein Segment stellt eine Sammlung von Chips dar, die sich alle im selben Pod befinden und über Hochgeschwindigkeits-Inter-Chip-Interconnect-Verbindungen (ICI) verbunden sind. v5e hat 2D-Slice-Formen. Unterstützte Kreisformen sind in der Tabelle im Abschnitt Beschleunigertypen aufgeführt.

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 dem es Vorhersagen oder Entscheidungen treffen kann. Latenz und Verfügbarkeit auf Dienstebene sind für die Bereitstellung wichtig.

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 acht 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 Segmenten und Chips finden Sie auf der Seite TPU-Systemarchitektur.

v5e unterstützt das Training mit mehreren Hosts und die Inferenz mit mehreren 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.

Mehr erfahren

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

Kapazität sichern

Wenden Sie sich an den Cloud-Vertrieb, um Cloud TPU v5e für Ihre KI-Arbeitslasten zu verwenden.

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 Cloud-Projekt auswählen oder erstellen.
  3. Die Abrechnungseinrichtung ist für die gesamte Nutzung von Google Cloud erforderlich. Achten Sie also darauf, dass die Abrechnung für Ihr Projekt aktiviert ist.

  4. Installieren Sie gcloud alpha Components.

  5. Wenn Sie als TPU-Nutzer vorhandene gcloud alpha-Komponenten wiederverwenden, aktualisieren Sie diese, damit relevante 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 es auch über die Google Cloud Console aktivieren.)

    gcloud services enable tpu.googleapis.com
    
  7. Aktivieren Sie das TPU-Dienstkonto.

    Dienstkonten ermöglichen es dem Cloud TPU-Dienst, auf andere Google Cloud-Dienste zuzugreifen. Ein nutzerverwaltetes Dienstkonto wird für Google Cloud empfohlen. Folgen Sie diesen Anleitungen zum Erstellen und Zuweisen von Rollen. Die folgenden Rollen sind erforderlich:

    • TPU-Administrator
    • Storage-Administrator: erforderlich für den Zugriff auf Cloud Storage
    • Logautor: erforderlich zum Schreiben von Logs mit der Logging API
    • Monitoring-Messwert-Autor: erforderlich zum Schreiben von Messwerten in Cloud Monitoring
  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 alpha 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

Cloud TPU v5es werden mit dem Befehl queued-resource create als Ressourcen in der Warteschlange bereitgestellt. Sie können aber 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 rot markierten Variablen 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 befindet sich in us-west4-a.
RUNTIME_VERSION Verwenden Sie v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT Dies ist die Adresse Ihres Dienstkontos, die Sie in der Google Cloud Console -> IAM -> Dienstkonten finden. Beispiel: tpu-service-account@myprojectID.iam.gserviceaccount.com
TPU_NAME Die vom Nutzer zugewiesene Text-ID der TPU, die bei der Zuweisung der in die Warteschlange gestellten Ressourcenanfrage 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 Das Flag kann entweder reserved oder best-effort sein. Informationen zu den verschiedenen von Cloud TPU unterstützten Kontingenttypen finden Sie unter TPU-Kontingenttypen. Das Best-Effort-Kontingent ist die Standardeinstellung, daher müssen Sie diesen Wert nicht festlegen.
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 sie die vorläufige Validierung bestanden hat und auf Kapazität wartet. Sobald Kapazität verfügbar ist, geht die Anfrage in PROVISIONING über. Der Status WAITING_FOR_RESOURCES bedeutet nicht, dass Ihnen das zugewiesene Kontingent zugewiesen wird. Es kann einige Zeit dauern, bis der Status vom Status WAITING_FOR_RESOURCES zu ACTIVE wechselt. Wenn die Ressource in der Warteschlange den Status FAILED hat, wird der Grund für den Fehler in der Ausgabe angegeben. Die Anfrage läuft ab, wenn das Feld --valid-until-duration nicht ausgefüllt wird, und der Status lautet „FAILED“.

Sie können mit SSH auf Ihre TPU-VM zugreifen, sobald Ihre QueuedResource den Status ACTIVE hat.

Verwenden Sie die Befehle list oder describe, um den Status der Ressource in der Warteschlange abzufragen.

    gcloud alpha 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 Status sind wie folgt definiert:

Bundesland Beschreibung
WAITING_FOR_RESOURCES Der Befehl create der Ressource in der Warteschlange ist eingegangen und wird bereitgestellt, sobald Kapazität verfügbar ist.
PROVISIONING Die TPU-Slices werden bereitgestellt.
ACTIVE Alle TPUs sind bereitgestellt und einsatzbereit. Wenn ein Startup-Skript angegeben wird, beginnt es auf allen TPUs mit seiner Ausführung, sobald der Ressourcenstatus der Warteschlange in ACTIVE übergeht.
FAILED Die Segmente konnten nicht bereitgestellt werden.
SUSPENDING Mindestens ein Segment wird gerade gelöscht.
SUSPENDED Alle zugrunde liegenden Segmente werden gelöscht. Die Ressource in der Warteschlange bleibt jedoch 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 können. In diesem Kontext wird eine TPU-VM auch als Worker bezeichnet.

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

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

gcloud alpha 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 alpha compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Wenn Sie Kreisformen mit mehr als 8 Chips haben, 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 alpha 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'

Verwalten

Alle Befehle, mit denen Sie Ihre TPU-VMs verwalten können, werden unter TPUs verwalten beschrieben.

Framework einrichten

In diesem Abschnitt wird der allgemeine Einrichtungsprozess für das benutzerdefinierte Modelltraining mit JaX oder PyTorch mit 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 zur Einrichtung von Inferenzen finden Sie unter Einführung in die V5e-Inferenz.

Einrichtung für JAX

Wenn Sie Kreisformen mit mehr als 8 Chips haben, 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 über SSH bei jeder einzelnen VM anzumelden:

gcloud alpha 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 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 alpha 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 Slice an. jax.local_device_count() gibt die Anzahl der Chips an, auf die eine einzelne VM in diesem Segment 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 alpha 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+ PJRT als Standardlaufzeit für alle TPU-Versionen 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 alpha 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
      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 Sie keine Räder für torch/torch_xla/torchvision 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, führen Sie mit dem folgenden Befehl ein Downgrade Ihrer Version aus:

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 alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker all --command='
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/$USER/.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/$USER/.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 mit dem V5e-Training mit PyTorch zu beginnen.

Löschen Sie Ihre TPU und die Ressource in der Warteschlange am Ende der Sitzung. Um eine Ressource in der Warteschlange zu löschen, löschen Sie das Segment und dann die Ressource in der Warteschlange in zwei Schritten:

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

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

Mit diesen beiden Schritten lassen sich auch Ressourcenanfragen in der Warteschlange entfernen, die den Status FAILED haben.

Überwachung und Profil

Cloud TPU v5e unterstützt Monitoring und Profilerstellung mit denselben Methoden wie vorherige Cloud TPU-Generationen. Unter Profil für Modell mit Cloud TPU-Tools erstellen finden Sie weitere Informationen zur Profilerstellung und unter Cloud TPU-VMs überwachen, um mehr über das Monitoring zu erfahren.

JAX/FLAX-Beispiele

ImageNet in Version 5e trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet in v5e mit fiktiven Eingabedaten trainieren. 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the QueuedResource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installieren Sie die neueste Version von JAX und jaxlib:

    gcloud alpha 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. Klonen Sie das ImageNet-Modell und installieren Sie die entsprechenden Anforderungen:

    gcloud alpha 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. Zum Generieren von fiktiven Daten benötigt das Modell Informationen zu den Dimensionen des Datasets. Diese Informationen können Sie den Metadaten des ImageNet-Datasets entnehmen:

    gcloud alpha 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

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

gcloud alpha 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 alpha compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet
gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --quiet

FLAX-Modelle, die sich umarmen

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

ViT auf Imagenette trainieren

In dieser Anleitung erfahren Sie, wie Sie das Vision Transformer(ViT)-Modell von HuggingFace mit dem Fast AI-Imagenette-Dataset in Cloud TPU v5e trainieren.

Das ViT-Modell war das erste Modell, das erfolgreich einen Transformer-Encoder auf ImageNet trainiert hat, der 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 alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installieren Sie JAX und die zugehörige Bibliothek:

    gcloud alpha 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. Laden Sie das Hugging Face-Repository herunter und installieren Sie die Anforderungen:

    gcloud alpha 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.15.0 && pip install -r examples/flax/vision/requirements.txt'
    
  5. Laden Sie das Imagenette-Dataset herunter:

    gcloud alpha 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 Puffer von 4 GB.

   gcloud alpha 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 Ressource in der Warteschlange am Ende der Sitzung.

   gcloud alpha compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --quiet
   
   gcloud alpha 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 in Pokémon Diffusion

In dieser Anleitung erfahren Sie, wie Sie das stabile 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installieren Sie JAX und die zugehörige Bibliothek.

    gcloud alpha 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. Laden Sie das HuggingFace-Repository herunter und installieren Sie die Anforderungen.

    gcloud alpha 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.15.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
    

Modell trainieren

Trainieren Sie das Modell mit einem vorab zugeordneten Puffer von 4 GB.

   gcloud alpha 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 Ihre TPU und die Ressource in der Warteschlange am Ende der Sitzung.

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

Benchmarking-Ergebnisse für die Diffusion

Das Trainingsskript wurde auf v5litepod-4, v5litepod-16 und v5litepod-64 ausgeführt. Die folgende Tabelle zeigt den Durchsatz.

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 mit dem 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 Labeling durch Menschen 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output
    will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installieren Sie JAX und die zugehörige Bibliothek.

    gcloud alpha 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. Laden Sie das HuggingFace-Repository herunter und installieren Sie die Anforderungen.

    gcloud alpha 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. Laden Sie Konfigurationen zum Trainieren des Modells herunter.

    gcloud alpha 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 Puffer von 4 GB.

   gcloud alpha 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 Ressource in der Warteschlange am Ende der Sitzung.

   gcloud alpha compute tpus tpu-vm delete ${TPU_NAME} 
--project=${PROJECT_ID}
--zone=${ZONE}
--quiet
   gcloud alpha 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. Die folgende Tabelle zeigt den Durchsatz.

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+ von XRT zu PjRt migriert. Hier finden Sie die aktualisierte Anleitung 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Torch-/XLA-spezifische Abhängigkeiten installieren

    gcloud alpha 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 alpha 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/$USER/.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 Ihre TPU und die Ressource in der Warteschlange am Ende der Sitzung.

gcloud alpha compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet
gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --quiet
Benchmarkergebnis

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 mit dem 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state,
    the output will be similar to the following:
    
    state: ACTIVE
    
  3. Installieren Sie Torch/xla-Abhängigkeiten.

    gcloud alpha 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. Laden Sie das HuggingFace-Repository herunter und installieren Sie die Anforderungen.

    gcloud alpha 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 alpha 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 alpha 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 Ressource in der Warteschlange am Ende der Sitzung.

   gcloud alpha compute tpus tpu-vm delete ${TPU_NAME} 
--project=${PROJECT_ID}
--zone=${ZONE}
--quiet
   gcloud alpha 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 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 auf v5e trainieren

In dieser Anleitung wird beschrieben, wie Sie VIT in v5e mit dem HuggingFace-Repository unter 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Torch-/xla-Abhängigkeiten installieren

    gcloud alpha 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. Laden Sie das HuggingFace-Repository herunter und installieren Sie die Anforderungen.

    gcloud alpha 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 alpha 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/$USER/.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 Ressource in der Warteschlange am Ende der Sitzung.

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

   gcloud alpha 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
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 trainieren (v5e)

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, finden Sie weitere Informationen unter 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
  3. Über SSH eine Verbindung zu Ihrer TPU herstellen

    gcloud alpha 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 alpha compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
    
  2. Ressourcenanfrage in der Warteschlange löschen

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

Resnet auf einem Multi-Host v5e trainieren

In dieser Anleitung wird beschrieben, wie Sie ImageNet mit einem fiktiven Dataset auf v5litepod-16 oder höher trainieren. 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}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
  3. Über SSH eine Verbindung zu Ihrer TPU (Worker Null) herstellen

    gcloud alpha 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 alpha compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
    
  2. Ressourcenanfrage in der Warteschlange löschen

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