Auf TPU-Pods trainieren

Übersicht

TPUs wurden für die Hochskalierung auf einen 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, sodass Sie die Verarbeitungslast auf mehrere TPUs verteilen können. Jedes TPU-Board ist mit einer leistungsstarken CPU-basierten Hostmaschine für das Laden und Vorverarbeiten von Daten verbunden. Sie müssen mehrere Parameter für Trainingsaufgaben optimieren, um die TPUs in vollem Umfang nutzen zu können. In diesem Dokument werden einige häufige Probleme, Änderungen, die Sie an Ihren Modellen vornehmen müssen, und Best Practices erläutert, um Pod-Fehler zu vermeiden oder zu vermeiden.

Batchgröße und Trainingsschritte skalieren

Für eine lineare Skalierung auf größeren TPU-Typen 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 das Flag train_steps, wobei ein Schritt der Verarbeitung eines einzelnen Datenbatches entspricht. Wenn Sie die Batchgröße erhöhen, müssen Sie die Anzahl der Trainingsschritte herunterskalieren, damit 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 nun vier Worker und eine effektive Batchgröße von 4.000 haben, müssen Sie die Anzahl der Schritte auf 25 für die Verarbeitung der gleichen 100.000 Beispiele einstellen. Wenn das Modell das Flag epochs verwendet, müssen Sie die Anzahl der Schritte nicht skalieren.

Größere Batchgrößen können das Konvergenzverhalten des Modells ändern. Daher können Sie auch einige Hyperparameter wie die Lernrate 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, wenn TPU-Pods verwendet werden, 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.

NFS zum Speichern von Daten 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 versuchen, von mehreren Knoten aus auf dasselbe NFS-Verzeichnis zuzugreifen. Das gleiche Verzeichnis hat dann unterschiedliche Eigentümer auf verschiedenen Knoten und die Standard-Linux-Berechtigungen gelten nicht knotenübergreifend. Ein Prozess in einem Knoten kann beispielsweise nicht in das Logverzeichnis schreiben, das von einem anderen Knoten erstellt wurde.

Zur Umgehung dieses Problems aktivieren Sie OS Login. 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 Workflows für die Entwicklung auf TPU-Pods

Beim Entwickeln einer neuen TPU-Arbeitslast ist es oft optimal, mit der Entwicklung auf den kleinsten TPUs zu beginnen und schrittweise zu größeren TPU-Größen zu iterieren. Verwenden Sie zuerst die kleinste TPU (z. B. v2-8 oder v3-8).

  • Arbeitslast auf Funktionalität testen
  • Leistung mithilfe der Leistungstools testen und validieren

Wenn Ihre Arbeitslast funktioniert und Ihre Leistungsziele erreicht ist, können Sie eine Hochskalierung auf v2-32 oder v3-32 vornehmen. Erhöhen Sie die TPU-Größe schrittweise und schrittweise, während Sie die Skalierbarkeit (Funktionalität und Leistung) prüfen, bis Sie die gewünschte TPU-Größe erreichen.