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.x、PyTorch、または JAX/flax をご覧ください。

メンテナンスイベントの検出

TPU でメンテナンスイベントが発生するかどうかとタイミングは、次の gcloud describe コマンドを使用して検出できます。

TPU VM

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

TPU ノード

$ gcloud compute tpus describe tpu-name --zone=zone | grep 'health'

このコマンドの出力には、TPU の現在の状態と最新のメンテナンスイベントの説明が表示されます。出力は次のようになります。

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

メンテナンスイベントログ

TPU のメンテナンスイベントの履歴ログは、システムイベント監査ログで確認できます。

Google Cloud コンソールのナビゲーションメニューで、[Compute Engine] > [VM インスタンス] をクリックして検索します。次に例を示します。

"tpu.nodes.terminate" OR "tpu.nodes.restart"

検索期間内に、TPU ワーカーの中断と修復が表示されます。ログには、イベントの日時、イベントの種類、「終了」イベントについては、protoPayload.metadata.terminateReason で終了した理由が表示されます。

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

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

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

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

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

本番環境に対応したトレーニングパイプラインでは、Google Kubernetes Engine（GKE）などのリソース管理システムを使用する必要があります。TPU VM アーキテクチャで Google Kubernetes Engine を使用する方法については、TPU ワークロードをデプロイするをご覧ください。TPU ノードアーキテクチャとともに Google Kubernetes Engine を使用する方法については、Google Kubernetes Engine で TPU アプリケーションを実行するをご覧ください。それ以外の場合は、bash スクリプトを実装すると、完了するまでトレーニングスクリプトの再試行を継続できます。例:

TPU ノードの場合:

（VM から） bash 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 からではなく、Google Cloud Shell またはターミナルから実行する必要があります）。
TPU の障害を検出して回復する

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

TPU VM を作成または再作成する場合は、--metadata startup-script パラメータで起動スクリプトを指定できます。起動スクリプトは、TPU VM が作成されるたびに実行されます。詳細については、標準のインストールスクリプトを実行するをご覧ください。

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

概要

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

メンテナンス イベントの検出