Auf TPU-Pods trainieren

Übersicht

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. Ein TPU-Pod kann bis zu 2.048 TPU-Kerne haben. Damit 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. In diesem Dokument 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.

Viele 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.

NFS zur Datenspeicherung verwenden

Wenn Sie einen TPU-Pod erstellen, wird für jeden TPU-Knoten eine separate VM erstellt. Standardmäßig wird jeder TPU-VM eine andere Nutzer-ID (UID) zugewiesen. Dies führt zu Problemen, wenn Sie von mehreren Knoten aus auf dasselbe NFS-Verzeichnis zugreifen möchten. Ein Verzeichnis hat dann möglicherweise unterschiedliche Inhaber auf verschiedenen Knoten und Standard-Linux-Berechtigungen gelten nicht für mehrere Knoten. Beispielsweise kann ein Prozess in einem Knoten nicht in das Logverzeichnis schreiben, das von einem anderen Knoten erstellt wurde.

Zur Umgehung dieses Problems können Sie OS Login aktivieren. OS Login kann für eine bestimmte VM-Instanz oder für ein Projekt konfiguriert werden. Weitere Informationen finden Sie unter OS Login einrichten.

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 die kleinste TPU (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 bis auf 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.