Fehlerbehebung für JAX – TPU

Dieser Leitfaden enthält Verweise auf Informationen zur Fehlerbehebung für JAX, damit Sie Probleme erkennen und beheben können, die beim Trainieren von JAX-Modellen auf Cloud TPU auftreten können.

Eine allgemeinere Anleitung zum Einstieg in Cloud TPU finden Sie in der JAX-Kurzanleitung.

Hinweis: Wenn Sie Ihr Problem mit diesem Leitfaden nicht lösen können, finden Sie unter Support erhalten weitere Unterstützung.

Allgemeine JAX-Probleme

Wenn bei der Entwicklung Ihres Trainingsmodells oder beim Training mit JAX Probleme auftreten, lesen Sie die FAQs zu JAX.

Allgemeinere Programmierfehler, die beim Schreiben einer Trainingsanwendung mit JAX auftreten können, finden Sie unter JAX-Fehler.

Profilerstellung für JAX-Leistung

Mithilfe der unter Profilerstellung für JAX-Leistung beschriebenen Tools können Sie nachvollziehen, wie Ihre TPU-Ressourcen genutzt werden.

Arbeitsspeicherprobleme beheben

Hinweis: Der TPU-Arbeitsspeicher (High Bandwidth Memory) wird von den TensorCores auf jedem TPU-Chip gemeinsam genutzt. Dies ermöglicht eine effiziente Koordination zwischen den TensorCores auf den Chips.

Mit dem JAX Device Memory Profiler können Sie überwachen, wie der Arbeitsspeicher verwendet wird. Sie können dessen Verwendung jedoch nicht direkt verwalten.

Der JAX Device Memory Profiler kann verwendet werden, um:

• herauszufinden, welche Arrays und ausführbare Dateien zu einem bestimmten Zeitpunkt im TPU-Arbeitsspeicher sind oder
• um Arbeitsspeicherlecks zu finden.

Sie können nicht vorgeben, wie der TPU-Arbeitsspeicher für bestimmte Vorgänge zugewiesen wird. Weitere Informationen zu JAX-spezifischen TPU-Leistungsproblemen finden Sie unter Leistungshinweise zur Verwendung von TPUs mit JAX.

TPU-Probleme beheben

In den folgenden Abschnitten wird beschrieben, wie Sie einige häufige Probleme beheben können, die beim Ausführen eines JAX-Programms auf einer TPU auftreten können.

Wie kann ich prüfen, ob die TPU ausgeführt wird?

Alle Vorgänge werden auf der TPU ausgeführt, solange JAX nicht „No GPU/TPU found, falling back to CPU“ ausgibt.

Sie können prüfen, ob die TPU aktiv ist. Dazu rufen Sie jax.devices() auf und es sollten dann mehrere TPU-Geräte angezeigt werden. Oder Sie überprüfen dies programmatisch mithilfe von assert jax.devices()[0].platform == 'tpu'.

RuntimeError: Unable to initialize backend 'tpu': UNAVAILABLE: No TPU Platform available.

Es wird diese Laufzeitfehlermeldung oder die folgende Meldung auf der TPU-VM in /tmp/tpu_logs/tpu_driver.WARNING angezeigt: W1118 17:40:20.985243 23901 tpu_version_flag.cc:57] No hardware is found. Using default TPU version:xxxxxx. Beides kann darauf hinweisen, dass Sie die falsche TPU-VM-Version verwenden.

Prüfen Sie in diesem Fall, ob Sie die aktuelle JAX-Laufzeitversion verwenden, und versuchen Sie es noch einmal.

Fehlerbehebung bei TPU- und GKE-Problemen

Zur Fehlerbehebung können Sie im GKE-Arbeitslastmanifest das ausführliche Logging aktivieren und die Logs dann an den GKE-Support senden.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

In den folgenden Abschnitten werden Fehlermeldungen im Zusammenhang mit TPU- und GKE-Einrichtungen sowie deren Lösungen beschrieben.

No endpoints available for service 'jobset-webhook-service'

Dieser Fehler bedeutet, dass das JobSet nicht richtig installiert wurde. Prüfen Sie, ob die Kubernetes-Pods für die jobset-controller-manager-Bereitstellung ausgeführt werden. Weitere Informationen finden Sie in der Dokumentation zur Fehlerbehebung für JobSet.

TPU initialization failed: Failed to connect

Die GKE-Knotenversion muss 1.30.4-gke.1348000 oder höher sein (GKE 1.31 wird nicht unterstützt).