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

概要

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 コマンドを使用して検出できます。

$ gcloud compute tpus tpu-vm 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 の終了の理由（「終了」イベントの場合）が表示されます。

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

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

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

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

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

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

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

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

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

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

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