We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText
Google Developers Blogが、MaxText、JAX、Pathways、Cloud TPUを使ったelastic trainingの解説を公開しました。分散学習中にTPUワーカーを意図的に停止しても、メインプロセスを再起動せず約2分以内に復旧する例を示しています。
Googleは、JAX AIスタックのMaxText、Pathways、Orbax、Cloud TPU、GKEを使い、分散LLM学習の途中でワーカー障害が起きても学習プロセスを生かしたまま復旧するelastic trainingを紹介しました。対象は大規模モデルをTPU上で分散学習するMLエンジニアや研究開発チームです。
通常の分散学習では、1台のマシンが消えるとall-reduceが停止し、全ワーカーが終了してジョブ全体を再起動する必要があります。記事の例では、Pathwaysが障害をPython例外として扱い、壊れたワーカーだけを置き換え、OrbaxがCloud Storage上のチェックポイントから復旧します。
実験では、TPUワーカーを意図的に停止した後、次の学習ステップまでの総ダウンタイムは2分未満で、その大半はKubernetesが代替Podをスケジュールする時間だったと説明されています。コントローラのPythonプロセスは同じPIDのまま再起動されません。
提供条件として、記事は再現方法を含む開発者向け解説です。対象API/ランタイムはMaxText、JAX、Pathways、Orbax、Cloud TPU、GKEで、モデル提供やGemini/Gemmaの新リリースではありません。破壊的変更や移行期限は示されていませんが、実運用ではチェックポイント、Pod置換、TPU/GKE構成の設計が重要になります。