TPUs zum Trainieren eines Modells verwenden

Tensor Processing Units (TPUs) sind von Google speziell entwickelte ASICs, die dazu dienen, ML-Arbeitslasten zu beschleunigen. Mit Cloud TPU können Sie Ihre Trainingsjobs in AI Platform Training ausführen. AI Platform Training bietet eine Benutzeroberfläche für die Jobverwaltung, sodass Sie die TPU nicht selbst verwalten müssen. Stattdessen können Sie die AI Platform Training jobs API genauso verwenden, wie Sie sie für das Training auf einer CPU oder einer GPU verwenden.

Mit den übergeordneten TensorFlow APIs können Sie Ihre Modelle auf der Cloud-TPU-Hardware ausführen.

Ihre Google Cloud-Umgebung einrichten

Konfigurieren Sie Ihre Google Cloud-Umgebung. Lesen Sie dafür im Startleitfaden den Abschnitt zur Einrichtung.

Cloud TPU für den Zugriff auf das Projekt autorisieren

Führen Sie die folgenden Schritte aus, um den Namen des Cloud TPU-Dienstkontos zu autorisieren, der Ihrem Google Cloud-Projekt zugeordnet ist:

  1. Rufen Sie den Namen Ihres Cloud TPU-Dienstkontos durch einen Aufruf von projects.getConfig ab. Beispiel:

    PROJECT_ID=PROJECT_ID

curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
    https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig

  2. Speichern Sie den Wert der Felder serviceAccountProject und tpuServiceAccount, der von der API zurückgegeben wird.

  3. Initialisieren Sie das Cloud TPU-Dienstkonto:

    curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
  -H "Content-Type: application/json" -d '{}'  \
  https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity

Fügen Sie dem Projekt nun das Cloud TPU-Dienstkonto als Mitglied mit der Rolle Cloud ML-Dienst-Agent hinzu. Führen Sie die folgenden Schritte in der Google Cloud Console oder mit dem gcloud-Befehl aus:

Console

  1. Melden Sie sich in der Google Cloud Console an und wählen Sie das Projekt aus, in dem Sie die TPU verwenden.
  2. Wählen Sie IAM & Verwaltung > IAM aus.
  3. Klicken Sie auf die Schaltfläche Hinzufügen, um ein Mitglied zum Projekt hinzuzufügen.
  4. Geben Sie das TPU-Dienstkonto in das Textfeld Mitglieder ein.
  5. Klicken Sie auf die Drop-down-Liste Rollen.
  6. Aktivieren Sie die Rolle Cloud ML-Dienst-Agent (Dienstverwaltung > Cloud ML-Dienst-Agent).

gcloud

  1. Legen Sie Umgebungsvariablen fest, die die Projekt-ID und das Cloud TPU-Dienstkonto enthalten:

    PROJECT_ID=PROJECT_ID
SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com

  2. Weisen Sie dem Cloud TPU-Dienstkonto die Rolle ml.serviceAgent zu:

    gcloud projects add-iam-policy-binding $PROJECT_ID \
    --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent

Weitere Informationen über das Zuweisen von Rollen zu Dienstkonten finden Sie in der IAM-Dokumentation.

Beispiel: MNIST-Beispielmodell trainieren

In diesem Abschnitt erfahren Sie, wie Sie ein MNIST-Beispielmodell mit einer TPU und Laufzeitversion 2.11 trainieren. In diesem Beispieljob wird als Maschinenkonfiguration die vordefinierte Skalierungsstufe BASIC_TPU verwendet. Weiter unten in dieser Anleitung erfahren Sie, wie Sie eine benutzerdefinierte Konfiguration einrichten.

In diesem Beispiel wird davon ausgegangen, dass Sie eine Bash-Shell mit installierter gcloud CLI verwenden. Führen Sie die folgenden Befehle aus, um den Code abzurufen und den Trainingsjob an AI Platform Training zu senden:

  1. Laden Sie den Code für die Referenzmodelle von TensorFlow herunter und wechseln Sie zum Verzeichnis mit dem Beispielcode:

    git clone https://github.com/tensorflow/models.git \
  --branch=v2.11.0 \
  --depth=1

cd models

  2. Erstellen Sie einesetup.py-Datei im Verzeichnis models. Dadurch berücksichtigt der Befehl gcloud ai-platform jobs submit training beim Erstellen eines Tarball-Archivs Ihres Trainingscodes alle erforderlichen Unterpakete im Verzeichnis models/official. Außerdem wird damit gewährleistet, dass AI Platform Training TensorFlow-Datasets als Abhängigkeit installiert, wenn der Trainingsjob ausgeführt wird. Dieser Trainingscode benötigt TensorFlow-Datasets zum Laden der MNIST-Daten.

    Führen Sie den folgenden Befehl in Ihrer Shell aus, um die Datei setup.py zu erstellen:

    cat << END > setup.py
from setuptools import find_packages
from setuptools import setup

setup(
    name='official',
    install_requires=[
       'tensorflow-datasets~=3.1',
       'tensorflow-model-optimization>=0.4.1'
   ],
    packages=find_packages()
)
END

  3. Senden Sie Ihren Trainingsjob mit dem Befehl gcloud ai-platform jobs submit training:

    gcloud ai-platform jobs submit training tpu_mnist_1 \
  --staging-bucket=gs://BUCKET_NAME \
  --package-path=official \
  --module-name=official.vision.image_classification.mnist_main \
  --runtime-version=2.11 \
  --python-version=3.7 \
  --scale-tier=BASIC_TPU \
  --region=us-central1 \
  -- \
  --distribution_strategy=tpu \
  --data_dir=gs://tfds-data/datasets \
  --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output

    Ersetzen Sie BUCKET_NAME durch den Namen eines Cloud Storage-Buckets in Ihrem Google Cloud-Projekt. Die gcloud CLI lädt den verpackten Trainingscode in diesen Bucket hoch und AI Platform Training speichert die Trainingsausgabe im Bucket.

  4. Überwachen Sie den Trainingsjob. Wenn der Job abgeschlossen ist, können Sie sich seine Ausgabe im Verzeichnis gs://BUCKET_NAME/tpu_mnist_1_output ansehen.

Weitere Möglichkeiten zum Trainieren von Modellen mit Cloud TPU

In diesem Abschnitt erfahren Sie mehr über das Konfigurieren eines Jobs und das Trainieren eines Modells in AI Platform Training mit Cloud TPU.

Region angeben, die TPUs anbietet

Sie müssen Ihren Job in einer Region ausführen, in der TPUs verfügbar sind. Der Zugriff auf TPUs ist derzeit in folgenden Regionen möglich:

  • us-central1
  • europe-west4

Ausführliche Informationen zu den verfügbaren Regionen für AI Platform Training-Dienste, einschließlich Modelltraining und Online-/Batchvorhersagen, finden Sie im Leitfaden für Regionen.

Versionsverwaltung für TensorFlow und AI Platform Training

Die AI Platform Training-Laufzeitversion(en) 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 und 2.11 sind zum Trainieren Ihrer Modelle auf Cloud TPU verfügbar. Hier finden Sie weitere Informationen zu den AI Platform Training-Laufzeitversionen und den entsprechenden TensorFlow-Versionen.

Die Versionsverwaltungsrichtlinie entspricht der Richtlinie für Cloud TPU. Geben Sie in der Trainingsjobanfrage unbedingt eine Laufzeitversion an, die für TPUs verfügbar ist und der Version von TensorFlow im Trainingscode entspricht.

Verbindung zum TPU-gRPC-Server herstellen

Verwenden Sie in Ihrem TensorFlow-Programm TPUClusterResolver, um eine Verbindung zum TPU-gRPC-Server herzustellen, der auf der TPU-VM ausgeführt wird.

In der TensorFlow-Anleitung zur Verwendung von TPUs wird dargestellt, wie Sie TPUClusterResolver mit der Verteilungsstrategie TPUStrategy verwenden.

Wenn Sie TPUClusterResolver für Code verwenden, der in AI Platform Training ausgeführt wird, müssen Sie jedoch eine wichtige Änderung vornehmen: Geben Sie beim Erstellen der TPUClusterResolver-Instanz keine Argumente an. Wenn die Schlüsselwortargumente tpu, zone und project auf den Standardwert None gesetzt sind, liefert AI Platform Training dem Cluster-Resolver automatisch die erforderlichen Verbindungsdetails in Form von Umgebungsvariablen.

Das folgende Beispiel für TensorFlow 2 veranschaulicht, wie ein Cluster-Resolver und eine Verteilungsstrategie für das Training in AI Platform Training initialisiert werden:

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

TPUs in TensorFlow-Code verwenden

Wenn Sie die TPUs auf einem Computer nutzen möchten, verwenden Sie die TPUStrategy API von TensorFlow 2. In der TensorFlow-Anleitung zur Verwendung von TPUs wird dies beschrieben.

Für das Training mit TPUs in TensorFlow 1 können Sie stattdessen die TPUEstimator API verwenden. In der Cloud TPU-Anleitung für die TPUEstimator API wird dies beschrieben.

Außerdem enthält die Cloud TPU-Dokumentation eine Liste der auf Cloud TPU verfügbaren einfachen TensorFlow-Vorgänge.

TPUs im PyTorch-Code verwenden

Für die Nutzung einer TPU bei einem vordefinierten PyTorch-Container nutzen Sie das Paket torch_xla. Mehr über die Verwendung von torch_xla für TPU im Training in der PyTorch-Dokumentation. Weitere Beispiele zur Verwendung von torch_xla finden Sie in den Anleitungen im PyTorch XL-GitHub-Repository.

Beachten Sie, dass Sie beim Trainieren mit einer TPU in AI Platform Training ein einzelnes XLA-Gerät und nicht mehrere XLA-Geräte verwenden.

Weitere Informationen finden Sie im folgenden Abschnitt über die Konfiguration des Trainingsjobs für PyTorch und TPU.

Benutzerdefinierte TPU-Maschine konfigurieren

Ein TPU-Trainingsjob wird in einer Konfiguration mit zwei VMs ausgeführt. Eine VM (der Master) führt Ihren Python-Code aus. Der Master steuert den TensorFlow-Server, der auf einem TPU-Worker ausgeführt wird.

Wenn Sie eine TPU mit AI Platform Training verwenden möchten, haben Sie drei Möglichkeiten, den Trainingsjob für den Zugriff auf eine TPU-fähige Maschine zu konfigurieren:

  • Verwenden Sie die Skalierungsstufe BASIC_TPU. Sie können diese Methode für den Zugriff auf TPU v2-Beschleuniger verwenden.
  • Verwenden Sie für die Master-VM einen cloud_tpu-Worker und einen Legacy-Maschinentyp. Sie können diese Methode für den Zugriff auf TPU v2-Beschleuniger verwenden.
  • Verwenden Sie für die Master-VM einen cloud_tpu-Worker und einen Compute Engine-Maschinentyp. Sie können diese Methode für den Zugriff auf TPU v2- oder TPU v3-Beschleuniger verwenden. TPU v3-Beschleuniger sind in der Betaversion verfügbar.

Einfache TPU-fähige Maschine

Setzen Sie die Skalierungsstufe auf BASIC_TPU, um eine Master-VM und eine TPU-VM mit einer einzigen TPU und acht TPU v2-Kernen zu erhalten, wie Sie es beim Ausführen des vorherigen Beispiels getan haben.

TPU-Worker in einer Konfiguration mit Legacy-Maschinentyp

Wenn Sie auf der Master-VM mehr Rechenressourcen benötigen, können Sie alternativ eine benutzerdefinierte Maschinenkonfiguration einrichten:

  • Setzen Sie die Skalierungsstufe auf CUSTOM.
  • Konfigurieren Sie die Master-VM so, dass ein Legacy-Maschinentyp entsprechend den Jobanforderungen verwendet wird.
  • Setzen Sie workerType auf cloud_tpu, damit Sie eine TPU-VM mit einer Cloud TPU und acht TPU v2-Kernen erhalten.
  • Setzen Sie workerCount auf 1.
  • Geben Sie keinen Parameterserver an, wenn Sie eine Cloud TPU verwenden. Der Dienst weist die Jobanfrage zurück, wenn parameterServerCount größer null ist.

Das folgende Beispiel zeigt eine config.yaml-Datei, die diesen Konfigurationstyp verwendet:

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

TPU-Worker in einer Konfiguration mit Compute Engine-Maschinentyp

Sie können auch eine benutzerdefinierte Maschinenkonfiguration mit einem Compute Engine-Maschinentyp für die Master-VM einrichten und an die TPU-VM eine acceleratorConfig anhängen.

Mit dieser Art der Konfiguration können Sie einen TPU-Worker mit acht TPU v2-Kernen (ähnlich wie bei einer Konfiguration ohne acceleratorConfig) oder einen TPU-Worker mit acht TPU v3-Kernen (Beta) einrichten. Hier erfahren Sie mehr über den Unterschied zwischen TPU v2- und TPU v3-Beschleunigern.

Die Verwendung eines Compute Engine-Maschinentyps bietet auch mehr Flexibilität bei der Konfiguration der Master-VM:

  • Setzen Sie die Skalierungsstufe auf CUSTOM.
  • Konfigurieren Sie die Master-VM für die Verwendung eines Compute Engine-Maschinentyps, der Ihren Jobanforderungen entspricht.
  • Setzen Sie workerType auf cloud_tpu.
  • Fügen Sie eine workerConfig mit dem Feld acceleratorConfig hinzu. In der acceleratorConfig setzen Sie type auf TPU_V2 oder TPU_V3 sowie count auf 8. Sie dürfen keine andere Anzahl von TPU-Kernen anhängen.
  • Setzen Sie workerCount auf 1.
  • Geben Sie keinen Parameterserver an, wenn Sie eine Cloud TPU verwenden. Der Dienst weist die Jobanfrage zurück, wenn parameterServerCount größer null ist.

Das folgende Beispiel zeigt eine config.yaml-Datei, die diesen Konfigurationstyp verwendet:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (Beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

TPU-Pods verwenden

Ein TPU-Pod ist eine Sammlung von TPU-Geräten, die über dedizierte Hochgeschwindigkeits-Netzwerkschnittstellen verbunden sind. Ein TPU-Pod kann bis zu 2.048 TPU-Kerne haben, sodass Sie die Verarbeitungslast auf mehrere TPUs verteilen können.

Wenn Sie TPU-Pods verwenden möchten, müssen Sie zuerst eine Anfrage zur Kontingenterhöhung stellen.

Die folgenden config.yaml-Beispieldateien zeigen die Verwendung von TPU-Pods:

TPU v2-Pods

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2_POD
      count: 128

TPU v3-Pods

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3_POD
      count: 32

Die Anzahl der Pod-Kerne, die für jeden TPU-Typ verwendet werden können, ist begrenzt. Verfügbare Konfigurationen:

TPU-Pod-Typ Anzahl der für die Verwendung verfügbaren Pod-Kerne
TPU_V2_POD 32, 128, 256, 512
TPU_V3_POD 32, 128, 256

Weitere Informationen zur vollständigen Nutzung von TPU-Pod-Kernen finden Sie in der Cloud TPU-Dokumentation zu TPU-Pods.

Vorgefertigten PyTorch-Container auf einem TPU-Worker verwenden

Wenn Sie ein PyTorch-Training mit einer TPU ausführen möchten, müssen Sie das Feld tpuTfVersion in der trainingInput des Trainingsjobs angeben. Legen Sie tpuTfVersion so fest, dass er der Version des vorkonfigurierten PyTorch-Containers entspricht, den Sie für das Training verwenden.

AI Platform Training unterstützt das Training mit TPUs für die folgenden vordefinierten PyTorch-Container:

URI des Container-Images tpuTfVersion
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 pytorch-1.11
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 pytorch-1.10
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 pytorch-1.9
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 pytorch-1.7
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 pytorch-1.6

Wenn Sie beispielsweise mit dem vorkonfigurierten Container PyTorch 1.11 trainieren möchten, können Sie das Training mit der folgenden config.yaml-Datei konfigurieren:

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
    tpuTfVersion: pytorch-1.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

Weitere Informationen finden Sie im vorherigen Abschnitt auf dieser Seite über TPUs in PyTorch-Code verwenden.

Benutzerdefinierte Container für einen TPU-Worker verwenden

Wenn Sie einen benutzerdefinierten Container auf Ihrem TPU-Worker ausführen möchten, anstatt eine der AI Platform Training-Laufzeitversionen zu verwenden, die TPUs unterstützen, müssen Sie beim Senden Ihres Trainingsjobs ein zusätzliches Konfigurationsfeld angeben. Legen Sie die tpuTfVersion auf eine Laufzeitversion fest, die die von Ihrem Container verwendete TensorFlow-Version enthält. Sie müssen eine Laufzeitversion angeben, die derzeit für das Training mit TPUs unterstützt wird.

Da Sie Ihren Job für die Verwendung eines benutzerdefinierten Containers konfigurieren, verwendet AI Platform Training bei der Ausführung Ihres Trainingsjobs nicht die Umgebung dieser Laufzeitversion. Für AI Platform Training ist dieses Feld jedoch erforderlich, damit der TPU-Worker ordnungsgemäß auf die Version von TensorFlow vorbereitet werden kann, die in Ihrem benutzerdefinierten Container verwendet wird.

Das folgende Beispiel zeigt eine config.yaml-Datei mit einer ähnlichen TPU-Konfiguration wie im vorherigen Abschnitt. In diesem Fall führen die Master-VM und der TPU-Worker jeweils unterschiedliche benutzerdefinierte Container aus:

TPU v2

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

TPU v3 (Beta)

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V3
      count: 8

Wenn Sie den Trainingsjob mit dem Befehl gcloud beta ai-platform jobs submit training senden, können Sie das API-Feld tpuTfVersion mit dem Flag --tpu-tf-version statt in einer config.yaml-Datei angeben.

TPUClusterResolver nach der Bereitstellung der TPU verwenden

Wenn Sie einen benutzerdefinierten Container verwenden, müssen Sie warten, bis die TPU bereitgestellt ist, bevor Sie TPUClusterResolver aufrufen können. Der folgende Beispielcode zeigt, wie die TPUClusterResolver-Logik verarbeitet wird:

def wait_for_tpu_cluster_resolver_ready():
  """Waits for `TPUClusterResolver` to be ready and return it.

  Returns:
    A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
    return None.
  Raises:
    RuntimeError: if failed to schedule TPU.
  """
  tpu_config_env = os.environ.get('TPU_CONFIG')
  if not tpu_config_env:
    tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
    return None

  tpu_node = json.loads(tpu_config_env)
  tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)

  num_retries = 40
  for i in range(num_retries):
    try:
      tpu_cluster_resolver = (
          tf.contrib.cluster_resolver.TPUClusterResolver(
              tpu=[tpu_node['tpu_node_name']],
              zone=tpu_node['zone'],
              project=tpu_node['project'],
              job_name='worker'))
      tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
      if 'worker' in tpu_cluster_resolver_dict:
        tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
        return tpu_cluster_resolver
    except Exception as e:
      if i < num_retries - 1:
        tf.logging.info('Still waiting for provisioning of TPU VM instance.')
      else:
        # Preserves the traceback.
        raise RuntimeError('Failed to schedule TPU: {}'.format(e))
    time.sleep(10)

  # Raise error when failed to get TPUClusterResolver after retry.
  raise RuntimeError('Failed to schedule TPU.')

Weitere Informationen zum verteilten Training mit benutzerdefinierten Containern

Nächste Schritte

Zurück
GPUs verwenden