Cloud TPU Pod を使用したメンテナンス イベントの管理

概要

TPU ノードと TPU VM は、TPU ハードウェアが接続された Compute Engine VM のインスタンスです。Compute Engine VM は Compute Engine VM のメンテナンス イベントの対象となります。各 TPU は Compute Engine VM に接続されているため、使用する TPU(TPU Pod など)を増やすと、いずれかの VM でメンテナンス イベントが発生する可能性があります。

このドキュメントでは、Cloud TPU で長時間実行トレーニング ジョブのメンテナンス イベントを処理するさまざまな方法について説明します。

チェックポイントを使用したメンテナンス イベントからの迅速な復元

チェックポイントは、メンテナンス イベントから短時間の復元を行う場合に重要であるため、頻繁に保存する必要があります。経験則として、チェックポイントは約 1 時間ごとに保存することをおすすめします。チェックポイントが不十分であると、メンテナンス イベントやその他のトレーニングの中断が原因となり、多くのトレーニングの進行状況が失われるリスクが発生します。

一般に、チェックポイントはトレーニングで使用されるすべての保存済みパラメータ(モデルの重みなど)を参照します。チェックポイントの保存に要する時間は、数秒から数分程度になることがあります。

メンテナンス イベントのほとんどは自動的に復元され、トレーニング ジョブは手動で操作しなくても継続されます。ただし、ジョブが再起動せず自動的に続行するケースが稀にあります。この場合、TPU リソースを削除して再作成し、保存したチェックポイントからトレーニング ジョブを再起動する必要があります。自動復元の失敗の検出と復旧方法については、以下の TPU の障害を検出して回復するをご覧ください。

チェックポイントの保存と読み込みに使用されるメカニズムは、ML フレームワークごとに異なります。通常、サポートされている Cloud TPU モデルにはチェックポインが組み込まれています。チェックポインティングの詳細については、TensorFlow 2.xPyTorch、または JAX/flax をご覧ください。

メンテナンス イベントを処理する

メンテナンス イベントの中断を緩和する方法はいくつかあります。

  1. チェックポイントを定期的に保存する

    理想的なシナリオでは、「割り込みイベント」が発生した場合、トレーニングは最新のチェックポイントから再開します。

  2. トレーニング スクリプトの再試行

    トレーニング スクリプトは、「中断イベント」の結果として停止する場合があります。bash スクリプトを使用すると、トレーニングが完了するまでトレーニング スクリプトを継続的に再試行できます。再試行は最新のチェックポイントから続行するため、再試行スクリプトは常にチェックポイントと組み合わせて使用する必要があります。

    本番環境に対応したトレーニング パイプラインでは、GKE またはその他のリソース管理システムを使用する必要があります。それ以外の場合は、bash スクリプトを実装すると、完了するまでトレーニング スクリプトの再試行を継続できます。次に例を示します。

    TPU ノードの場合:

    (VM から) while ! python3 [training command]; do sleep 1; done

    TPU VM の場合: while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done

    (TPU VM コマンドは、TPU VM からではなく、GCP またはターミナルから実行する必要があります)。

  3. TPU の障害を検出して回復する

    TPU がメンテナンス イベントから回復しない場合、復元スクリプトを使用して TPU の状態を検出し、TPU を削除して再作成できます。このスクリプトの例は、こちらをご覧ください。 TPU の手動削除と再作成の詳細については、TPU の管理をご覧ください。