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 ist für jedes Framework unterschiedlich. Über die folgenden Links erhalten Sie detaillierte Informationen zum Training auf Pods mit jedem Framework:

• TensorFlow

• PyTorch

• JAX

In den folgenden Abschnitten werden einige häufige Probleme, Änderungen, die Sie an Ihren Modellen vornehmen müssen, sowie Best Practices zur Reduzierung oder Vermeidung von Pod-Fehlern erläutert.

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 das Flag train_steps, wobei ein Schritt der Verarbeitung eines einzelnen Batches von Daten 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 Ressourcenregion ist besonders wichtig bei der Verwendung von TPU-Pods, da die Datenübertragungsraten höher sind, wenn sich der Google Cloud Storage-Bucket und die 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

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