Cloud TPU v5e-Training [Öffentliche Vorschau]

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

Konzepte von Cloud TPU v5e

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

In diesem Abschnitt werden Konzepte definiert, die sich speziell auf v5e beziehen. Komponenten, die in anderen TPU-Konfigurationen verwendet werden (z. B. Segmente, Hosts, Chips und TensorCores), werden auf der Seite TPU-Systemarchitektur beschrieben.

Inferenz

Inferenz ist die Verwendung eines trainierten Modells, um Vorhersagen für neue Daten zu treffen. Sie wird für den Bereitstellungsprozess verwendet.

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 für die Bereitstellung wichtig.

Slices

Ein Segment 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-Segmentformen. Die unterstützten Kreisformen sind in der Tabelle im Abschnitt Beschleunigertypen aufgeführt.

Chipform und Chiptopologie beziehen sich auch auf Kreisformen.

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.

Segmente mit weniger als 8 Chips verwenden höchstens einen Host. Segmente mit mehr als acht Chips haben Zugriff auf mehr als einen Host und können ein 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 einem einzelnen Host. In dieser Vorschau wird keine Inferenz oder Bereitstellung aus mehreren Hosts unterstützt. In einem einzelnen Bereitstellungsjob sind maximal 8 Chips zulässig.

TPU-VM

Eine virtuelle Maschine mit Linux, die Zugriff auf die zugrunde liegenden TPUs hat. Bei v5e-TPUs hat jede TPU-VM je nach benutzerdefiniertem Beschleunigertyp direkten Zugriff auf 1, 4 oder 8 Chips. Eine TPU-VM wird auch als Worker bezeichnet.

Worker

Siehe TPU-VM.

v5e-Systemarchitektur

Jeder v5e-Chip enthält einen TensorCore. Jeder TensorCore hat vier Matrix Multiply Units (MXU), eine Vektoreinheit und eine Skalareinheit. Die folgende Tabelle enthält die wichtigsten Chipspezifikationen und ihre Werte für eine v5e. Pod-Spezifikationen sind in der Tabelle aufgeführt, die den Chipspezifikationen entspricht.

Wichtige Chipspezifikationen v5e-Werte
Spitzenberechnung pro Chip (bf16) 197 TFLOPs
Spitzenberechnung pro Chip (Int8) 393 TFLOPs
HBM2-Kapazität und Bandbreite 16 GB, 819 GB/s
Interchip Interconnect-BW 1.600 Gbit/s
Wichtige Pod-Spezifikationen v5e-Werte
TPU-Pod-Größe 256 Chips
Interconnect-Topologie 2D Torus
Spitzen-Computing pro Pod 100 PetaOps(Int8)
Gesamte Bandbreite pro Pod reduzieren 51,2 TB/s
Bandbreite pro Pod aufteilen 1,6 TB/s
Netzwerkbandbreite für Rechenzentren pro Pod 6,4 Tbit/s

TPU v5e-Chip

Das folgende Diagramm veranschaulicht einen TPU v5e-Chip.

v5e-Pod-Chip

Beschleunigertypen

Cloud TPU v5e ist ein kombiniertes Trainings- und Inferenzprodukt (Bereitstellung). Das Flag AcceleratorType wird verwendet, um zwischen der für das Training bereitgestellten TPU-Umgebung und der für die Bereitstellung bereitgestellten Umgebung zu unterscheiden.

Trainingsjobs sind hinsichtlich Durchsatz und Verfügbarkeit optimiert, während Bereitstellungsjobs für Latenz optimiert sind. Ein Trainingsjob auf TPUs, die für die Bereitstellung bereitgestellt werden, kann also zu einer schlechten Nutzererfahrung führen (geringere Verfügbarkeit). Ebenso könnte ein Bereitstellungsjob, der auf TPUs ausgeführt wird, die für das Training bereitgestellt wurden, zu einer schlechten Leistung führen (höhere Latenz).

AcceleratorType verwendet die Anzahl der TensorCores im Knoten, um die Größe des Segments zu beschreiben. AcceleratorType ist ein String im Format v$VERSION_NUMBER-$CORES_COUNT.

Weitere Informationen zu Beschleunigertypen finden Sie unter Unterstützte TPU-Konfigurationen.

Eine v5e hat einen Kern pro Chip. Die folgenden 2D-Segmentformen werden für V5es unterstützt:

Topologie Anzahl der TPU-Chips Anzahl der Hosts
1x1 1 1/8
2x2 4 02.01.
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

Die Anzahl der VMs pro Topologie finden Sie unter VM-Typen.

Cloud TPU v5e-Typen

Training wird für bis zu 256 Chips unterstützt.

Verwenden Sie zum Bereitstellen von TPUs für einen v5e-Trainingsjob einen der folgenden Beschleunigertypen in Ihrer Anfrage zur Erstellung einer Befehlszeilen- oder API-TPU:

  • V5e-1
  • V5e-4
  • V5e-8
  • V5e-16
  • V5e-32
  • V5e–64
  • V5e-128
  • V5e-256

Die Bereitstellung für Inferenzen wird in der Einführung in v5e-Inferenzen beschrieben.

TPU-VM-Typen

Jede TPU-VM in einem TPU-Segment enthält 1, 4 oder 8 Chips. 4-Chips und kleinere Segmente haben denselben NUMA-Knoten (Non Uniform Memory Access). Weitere Informationen zu NUMA-Knoten finden Sie im Dokument zur Systemarchitektur. Bei 8-Chip-TPU-VMs ist die CPU-TPU-Kommunikation innerhalb der NUMA-Partitionen effizienter. In der folgenden Abbildung ist die CPU0-Chip0-Kommunikation beispielsweise schneller als die CPU0-Chip4-Kommunikation.

Kommunikation mit NUMA-Knoten

Mehr erfahren

Kapazität sichern

Cloud TPU v5e ist jetzt in der öffentlichen Vorschau verfügbar. Wenden Sie sich an den Cloud-Vertrieb, wenn Sie Cloud TPU v5e für Ihre KI-Arbeitslasten verwenden möchten.

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 Google Cloud-Nutzung erforderlich. Achten Sie also darauf, dass die Abrechnung für Ihr Projekt aktiviert ist.

    Die Abrechnung für die öffentliche Vorschau richtet sich nach den regionalen Standardpreisen, die auf der Seite mit den Cloud TPU-Preisen aufgeführt sind.

  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 ist eine empfohlene Vorgehensweise für Google Cloud. Folgen Sie diesen Anleitungen, um Rollen zu erstellen und Rollen zuzuweisen. 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. Projekt und Zone konfigurieren

    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 als Ressourcen in der Warteschlange verwaltet. Mit dem Befehl queued-resource create können Sie Kapazität bereitstellen.

Umgebungsvariablen erstellen

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

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-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

Variable Beschreibung
PROJECT_ID Name des Google Cloud-Projekts
ACCELERATOR_TYPE Die unterstützten Segmentformen finden Sie im Abschnitt Beschleunigertypen.
Zone Die gesamte öffentliche Vorschau-Kapazität befindet sich in us-west4-a.
RUNTIME_VERSION v2-alpha-tpuv5-lite für die öffentliche Vorschau verwenden.
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 erstellt wird, wenn die Ressourcenanfrage in der Warteschlange zugewiesen wird.
RESSOURCEN-ID_IN DER WARTESCHLANGE Die vom Nutzer zugewiesene Text-ID der Ressourcenanfrage in der Warteschlange. Informationen zu Ressourcen in der Warteschlange finden Sie im Dokument Ressourcen in der Warteschlange.

Eine 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} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved

Wenn die Ressource in der Warteschlange erfolgreich erstellt wurde, hat das Feld response den Status ACCEPTED oder FAILED. Wenn die Ressource in der Warteschlange den Status ACCEPTED 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, wird die Anfrage auf PROVISIONING umgestellt. Der Status ACCEPTED bedeutet nicht zwangsläufig, dass Ihnen das Kontingent zugewiesen wird. Es kann einige Zeit dauern, bis der Status von ACCEPTED zu ACTIVE wechselt. Wenn die Ressource in der Warteschlange den Status FAILED hat, wird der Fehlergrund in der Ausgabe angegeben.

Sie können eine SSH-Verbindung zu Ihrer TPU-VM herstellen, 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 Bundesstaaten sind definiert als:

Bundesland Beschreibung
Angenommen Der Befehl create der Ressource ist eingegangen und wird bereitgestellt, sobald Kapazität verfügbar ist.
PROVISIONING Die TPU-Segmente werden bereitgestellt.
ACTIVE Alle TPUs sind bereitgestellt und einsatzbereit. Wenn ein Startskript angegeben wird, wird es auf allen TPUs ausgeführt, wenn der Status der in der Warteschlange befindlichen Ressourcen in ACTIVE übergeht.
FAILED Die Segmente konnten nicht bereitgestellt werden.
SUSPENDING Mindestens ein Segment wird 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 angehaltene Ressource in der Warteschlange nicht fortgesetzt werden und sollte gelöscht werden.
DELETING Die Ressource in der Warteschlange wird gelöscht.

SSH-Verbindung zu TPU-VMs herstellen

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

Im Abschnitt VM-Typen können Sie berechnen, wie viele VMs Ihr Segment haben wird.

Zum Installieren der Binärdateien oder Ausführen von Code stellen Sie eine SSH-Verbindung zur TPU-VM mit tpu-vm ssh command her.

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}

Wenn Sie eine SSH-Verbindung zu einer bestimmten TPU-VM herstellen möchten, 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, haben Sie mehrere VMs in einem Segment. Verwenden Sie in diesem Fall das Flag --worker=all, um die Installation auf allen TPU-VMs in einem einzigen Schritt auszuführen, ohne jede separate SSH-Verbindung 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.6" -f https://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 mit TPU v5e beschrieben. Die TensorFlow-Unterstützung wird bald aktiviert.

Eine Anleitung zur Einrichtung von Inferenzen finden Sie in der Einführung in die Inferenz von V5e.

Einrichtung für JAX

Wenn Sie Kreisformen mit mehr als 8 Chips haben, haben Sie mehrere VMs in einem Segment. In diesem Fall müssen Sie das Flag --worker=all verwenden, um die Installation auf allen TPU-VMs in einem Schritt auszuführen, ohne sich mit SSH in 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.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Sie können den folgenden Befehl ausführen, um die Anzahl der Geräte zu prüfen. Die hier gezeigten Ausgaben wurden mit einem v5litepod-16-Segment erzeugt. 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 Segment 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 loszulegen.

PyTorch einrichten

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 in 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 numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
'

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, führen Sie mit dem folgenden Befehl ein Downgrade Ihrer Version aus:

pip3 install setuptools==62.1.0

Führen Sie ein einfaches Skript mit PJRT aus:

Bei Modellen mit großen, häufigen Zuweisungen kann die Verwendung von tcmalloc die Trainingszeit im Vergleich zur Standardimplementierung malloc erheblich verbessern. Daher ist der Standard-malloc, der auf einer TPU-VM verwendet wird, tcmalloc. Je nach Arbeitslast (z. B. bei DLRM, das sehr große Zuweisungen für die eingebetteten Tabellen hat) kann tcmalloc jedoch zu einer Verlangsamung führen. In diesem Fall können Sie versuchen, die folgende Variable mit dem Standardwert malloc aufzuheben:

unset LD_PRELOAD

Im folgenden Beispiel wird ein Python-Skript verwendet, um eine einfache Berechnung auf einer v5e-VM auszuführen:

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

Die PyTorch-Anleitungen in diesem Dokument helfen Ihnen dabei, mit dem V5e-Training mit PyTorch loszulegen.

Löschen Sie die 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 können auch Ressourcenanfragen in der Warteschlange entfernt werden, die den Status FAILED haben.

Monitoring und Profil

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

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=v5e-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-test-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

  2. Erstellen Sie eine TPU-Ressource:

    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} \
--service-account ${SERVICE_ACCOUNT} \
--reserved

    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.6" -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 -b throughput_testing https://github.com/gkroiz/flax.git && cd flax && pip install -e . && cd examples/imagenet && pip install -r requirements.txt && pip install --upgrade git+https://github.com/google/flax.git'

  5. Zum Generieren fiktiver Daten benötigt das Modell Informationen zu den Dimensionen des Datasets. Dies 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

Sobald alle vorherigen Schritte abgeschlossen sind, können Sie das Modell trainieren. Im GitHub-Repository gibt es mehrere Konfigurationsdateien, die verwendet werden können. Alternativ können Sie Konfigurationen in der Befehlszeile überschreiben.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='cd flax/examples/imagenet && JAX_USE_PJRT_C_API_ON_TPU=1 JAX_PLATFORMS=tpu TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 TPU_MIN_LOG_LEVEL=0 python3 imagenet_fake_data_benchmark.py --config=configs/fake_data_16_chips.py'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die 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 das Gesicht umarmen

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

ViT auf Imagenette trainieren

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

Das ViT-Modell war das erste, das erfolgreich einen Transformer-Encoder auf ImageNet trainiert hat, der im Vergleich zu Faltungsnetzwerken 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=v5e-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-test-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

  2. Eine 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} \
--service-account ${SERVICE_ACCOUNT}

    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.6" -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 && 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 && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die 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 in v5e-4, v5e-16 und v5e-64 ausgeführt. Die folgende Tabelle zeigt den Durchsatz mit verschiedenen Beschleunigertypen.

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

Trainiere die Diffusion in Pokémon

In dieser Anleitung erfahren Sie, wie Sie das Modell für die stabile Diffusion von HuggingFace mit dem Dataset Pokémon 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=v5e-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-test-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

  2. Erstellen Sie eine TPU-Ressource:

    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} \
--service-account ${SERVICE_ACCOUNT} \
--reserved

    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. JAX und die zugehörige Bibliothek installieren

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -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 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 && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --fake_data=1 --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die 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 Streuung

Das Trainingsskript lief auf v5e-4, v5e-16 und v5e-64. In der folgenden Tabelle sind die Durchsatze aufgeführt.

Beschleunigertyp V5e-4 V5e-16 V5e–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.

Das GPT2 ist ein Transformer-Modell, das mit Rohtexten ohne menschliches Labeling vortrainiert wurde. Es wurde 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=v5e-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-test-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

  2. Eine 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} \
--service-account ${SERVICE_ACCOUNT} \
--reserved

    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. JAX und die zugehörige Bibliothek installieren

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -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-v2-images-artifacts/config/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 die 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 lief auf v5e-4, v5e-16 und v5e-64. In der folgenden Tabelle sind die Durchsatze aufgeführt.

V5e-4 V5e-16 V5e–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 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=v5e-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=your_queued_resource_id

  2. Erstellen Sie eine TPU-Ressource:

    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} \
--service-account ${SERVICE_ACCOUNT} \
--reserved

    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. PyTorch/XLA klonen

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command='cd /usr/share/
sudo git clone --recursive https://github.com/pytorch/pytorch
cd pytorch/
sudo git clone --recursive https://github.com/pytorch/xla.git
cd xla/
sudo git checkout 3b10663
sudo pip3 install numpy'

  4. Torch, Torch/xla, Torchvision installieren

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} --project ${PROJECT_ID} --worker=all --command='
pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
sudo pip3 install torch_xla[tpuvm]
'

  5. 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 numactl libopenblas-dev -y
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
pip3 install torch_xla[tpuvm]
'

  6. libtpu herunterladen

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/libtpu.so /home/$USER/.local/lib/python3.8/site-packages/libtpu/libtpu.so
"

  7. tpu.py mit der V5e-Konfiguration aktualisieren

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/tpu.py /home/$USER/.local/lib/python3.8/site-packages/torch_xla/experimental/tpu.py
"
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 die 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
Benchmark-Ergebnis

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

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

GPT2 in V5e trainieren

In dieser Anleitung wird beschrieben, wie GPT2 in v5e mit dem HuggingFace-Repository auf PyTorch/XLA unter Verwendung des Wikitext-Datasets ausgeführt wird.

Einrichtung

  1. Erstellen Sie Umgebungsvariablen:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5e-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

  2. Eine 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} \
--service-account ${SERVICE_ACCOUNT}

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

    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 numactl libopenblas-dev -y
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
pip3 install torchvision-0.16.0a0+fc838ad-cp310-cp310-linux_x86_64.whl
pip3 install torch_xla[tpuvm]
'

  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="
rm -rf transformers/
git clone https://github.com/huggingface/transformers.git; \
cd transformers; \
sudo pip3 install -e .; \
git checkout -q ebdb185befaa821304d461ed6aa20a17e4dc3aa2; \
pip3 install datasets; \
pip3 install evaluate; \
pip3 install scikit-learn;"

  5. Laden Sie Konfigurationen des vortrainierten Modells herunter.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
cd transformers &&
gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/my_config_2.json examples/pytorch/language-modeling/
gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/fsdp_config.json 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
   export PT_XLA_DEBUG=0
   export USE_TORCH=ON
   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 transformers
   numactl --cpunodebind=0 python3 -u 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
    --output_dir /tmp/test-clm
    --overwrite_output_dir
    --config_name examples/pytorch/language-modeling/my_config_2.json --cache_dir /tmp
    --tokenizer_name gpt2
    --block_size 1024
    --optim adafactor
    --adafactor true
    --save_strategy no
    --logging_strategy no
    --fsdp "full_shard"
    --fsdp_config examples/pytorch/language-modeling/fsdp_config.json'

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die 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

Benchmark-Ergebnis

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

V5e-4 V5e-16 V5e–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 auf 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=v5e-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

  2. Eine 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} \
--service-account ${SERVICE_ACCOUNT} \
--reserved

    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 numactl libopenblas-dev -y
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly+20230726-cp310-cp310-linux_x86_64.whl
pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20230726-cp310-cp310-linux_x86_64.whl
gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
pip3 install torch_xla[tpuvm]
'

  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;
      numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10
      --remove_unused_columns=False
      --label_names=pixel_values
      --mask_ratio=0.75
      --norm_pix_loss=True
      --do_train=true
      --do_eval=true
      --base_learning_rate=1.5e-4
      --lr_scheduler_type=cosine
      --weight_decay=0.05
      --num_train_epochs=3
      --warmup_ratio=0.05
      --per_device_train_batch_size=8
      --per_device_eval_batch_size=8
      --logging_strategy=steps
      --logging_steps=30
      --evaluation_strategy=epoch
      --save_strategy=epoch
      --load_best_model_at_end=True
      --save_total_limit=3
      --seed=1337
      --output_dir=MAE
      --overwrite_output_dir=true
      --logging_dir=./tensorboard-metrics
      --tpu_metrics_debug=true"

TPU und Ressource in der Warteschlange löschen

Löschen Sie die TPU und die 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

Benchmark-Ergebnis

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

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

Support und Feedback

Wir freuen uns über jedes Feedback! Wenn Sie Feedback geben oder Support anfordern möchten, kontaktieren Sie uns oder senden Sie eine E-Mail an cloudtpu-preview-support@google.com.

Nutzungsbedingungen

Alle Informationen, die Google Ihnen im Rahmen dieser öffentlichen Vorschau zur Verfügung gestellt hat, sind vertrauliche Informationen von Google und unterliegen den Vertraulichkeitsbestimmungen der Nutzungsbedingungen der Google Cloud Platform (oder anderer Vereinbarungen, die Ihre Nutzung der Google Cloud Platform regeln).