Auf TPU-Pods trainieren

Überblick

TPUs wurden für die horizontale Skalierung auf einem TPU-Pod entwickelt. Ein TPU-Pod ist eine Sammlung von TPU-Geräten, die über dedizierte Hochgeschwindigkeits-Netzwerkschnittstellen verbunden sind. Mit einem TPU-Pod können Sie die Verarbeitungslast auf mehrere TPUs verteilen. Jedes TPU-Board ist zum Beispiel für das Laden und Vorverarbeiten von Daten mit einem hochleistungsfähigen CPU-basierten Hostcomputer verbunden. Um eine größere Anzahl von TPUs bestmöglich nutzen zu können, müssen Sie mehrere Trainingsaufgabenparameter optimieren.

Die Einrichtung für das Training mit TPU-Pods unterscheidet sich je nach Framework. Unter den folgenden Links finden Sie ausführliche Informationen zum Training mit Pods für die einzelnen Frameworks:

In den folgenden Abschnitten werden einige häufig auftretende Probleme, Änderungen, die Sie an Ihren Modellen vornehmen müssen, sowie Best Practices erläutert, um Pod-Fehler zu reduzieren oder zu vermeiden.

Batchgröße und Trainingsschritte skalieren

Wenn Sie bei größeren TPU-Typen eine lineare Skalierung erzielen möchten, sollten Sie die Batchgröße pro Kern beibehalten.

Wenn Sie beispielsweise für einen v2-8-TPU-Typ eine Batchgröße von 1.024 verwenden, sollten Sie für einen v2-32-TPU-Typ eine Batchgröße von 4.096 (4 * 1.024) verwenden. Dadurch wird die TPU-Hardware bestmöglich genutzt. Sie können kleinere Batchgrößen verwenden, aber das Training wird dann nicht linear skaliert.

Einige Modelle enthalten ein Flag train_steps, wobei ein Schritt einem einzelnen Datenbatch entspricht. Wenn Sie die Batchgröße erhöhen, müssen Sie die Anzahl der Trainingsschritte herunterskalieren, sodass die Gesamtzahl der Trainingsbeispiele gleich bleibt.

Wenn Sie beispielsweise eine Batchgröße von 1.000 für 100 Schritte haben, werden während des Trainings 100.000 Beispiele verarbeitet. Wenn Sie jetzt vier Worker und eine effektive Batchgröße von 4.000 haben, müssen Sie die Anzahl der Schritte auf 25 anpassen, um dieselben 100.000 Beispiele zu verarbeiten. Wenn das Modell ein Flag epochs verwendet, müssen Sie die Anzahl der Schritte nicht skalieren.

Größere Batches können zu einem geänderten Konvergenzverhalten des Modells führen. Es empfiehlt sich daher ggf., auch gewisse Hyperparameter wie die Lernrate zu optimieren.

Regionale Google Cloud Storage-Buckets in der Region des TPU-Pods verwenden

Die Best Practice für das TPU-Training besteht im Allgemeinen darin, immer nur Ressourcen aus einer Region zu verwenden. Die Region der Ressourcen ist insbesondere dann wichtig, wenn Sie TPU-Pods verwenden. Die Datenübertragungsraten sind höher, wenn sich Ihr Google Cloud Storage-Bucket und Ihre TPU in derselben Region befinden.

Verwenden Sie für Training-Datasets und Prüfpunkte unbedingt einen regionalen Google Cloud Storage-Bucket in derselben Region wie die TPU.

Best Practices für den Workflow bei der Entwicklung auf TPU-Pods

Wenn Sie eine neue TPU-Arbeitslast erstellen, ist es oft sinnvoll, mit der Entwicklung auf den kleinsten TPUs zu beginnen und schrittweise auf größere TPUs zu iterieren. Verwenden Sie zuerst eine kleine TPU-Version (z. B. v2-8 oder v3-8).

  • Arbeitslast auf Funktionalität testen
  • Leistung mithilfe der Performance-Tools testen und validieren

Wenn Ihre Arbeitslast funktionsfähig ist und Ihre Leistungsziele erreicht hat, skalieren Sie auf einen größeren TPU-Typ wie v2-32 oder v3-32 hoch. Erhöhen Sie die TPU-Größe graduell und iterativ und prüfen Sie dabei die Skalierbarkeit (Funktionalität und Leistung), bis Sie die gewünschte TPU-Größe erreichen.