Pathways を使用した復元力のあるトレーニング

このページでは、Google Kubernetes Engine（GKE）で Pathways ワークロードの復元機能を構成し、大規模なトレーニング ジョブが計画的な中断や計画外のハードウェア障害から復旧できるようにする方法について説明します。

Pathways は、次の方法で復元力のメリットを提供します。

  • 一時停止と再開: プリエンプション通知などの計画的な中断をワークロードで許容できるようにします。カスタムのプリエンプション処理コードは必要ありません。
  • エラスティック トレーニング: トレーニング コードの復元ロジックを使用して、クライアントをクラッシュさせることなく、計画外のハードウェア障害にワークロードが耐えられるようにします。
  • 多層チェックポイント処理: ワークロードは、同じ場所に配置された Python を使用して、TPU ワーカー間で高頻度のローカル チェックポイントとピア レプリケート チェックポイントを保存できます。また、非同期の Cloud Storage バックアップにより、復元時間を短縮できます。
  • ホットスワップ: ハードウェア障害が発生したときに、優先度の高い GKE JobSet ワークロードが、スペアまたは優先度の低い TPU スライスを迅速に引き継ぐことができます。

始める前に

始める前に、次のタスクを完了します。

エラスティック トレーニングを使用する場合は、ワークロード コンテナ イメージに次のパッケージをインストールします。

  • pathwaysutils
  • Orbax（orbax-checkpoint）: >= 0.12.6（ocp.v1.training.pathways.Snapshotter を使用したインメモリ スナップショットに必要）
  • グレイン（grain）: grain.experimental.ElasticIterator を使用する場合は必須

多層チェックポイント処理を使用する場合は、次の前提条件も満たしてください。

  • 多層チェックポイント処理を使用して GKE で大規模な ML モデルをトレーニングするの手順に沿って、多層チェックポイント処理用に GKE クラスタを構成します。
    • クラスタで、Workload Identity Federation for GKE とともに、多層チェックポイント処理と Cloud Storage FUSE（GcsFuseCsiDriver）アドオンを有効にします。
    • 階層型名前空間が有効になっている Cloud Storage バケットを作成します。
    • 多層チェックポイント処理の Kubernetes サービス アカウントに roles/storage.objectUser と roles/storage.insightsCollectorService の Identity and Access Management（IAM）ロールを付与します。
    • Cloud Storage バケット、TPU ノードセレクタ、インメモリ ボリューム サイズ（inMemoryVolumeSize）を指定して、CheckpointConfiguration カスタム リソースをデプロイします。
  • サンプルのサイドカー Dockerfile などを使用して、コロケーションされた Python サイドカー コンテナ イメージをビルドして push します。ワークロード イメージとサイドカー イメージで、Pathways リリースに一致し、次の最小バージョン要件を満たす jax バージョンと jaxlib バージョンを使用します。
    • Orbax のバージョン: >= 0.12.4
    • JAX バージョン（コロケーションされた Python を含む）: >= 0.10.1

一時停止と再開

GKE は通常、Pod をプリエンプトする前にプリエンプション通知をアクセラレータ Pod に送信します。Pathways では、すべての Cloud Deploy でプリエンプション耐性がデフォルトで有効になっており、Pathways アクセラレータ ジョブはこれらの通知をリッスンします。

プリエンプション通知が届くと、Pathways はまず、ワークロードを透過的に保存して復元できるかどうかを判断します。可能な場合、Pathways は、GKE がアクセラレータ ジョブを強制排除する前に、現在の状態を Cloud Storage などの永続ストレージに書き出すことで、ML ワークロードを透過的に一時停止しようとします。後で GKE がジョブを再スケジュールすると、Pathways は永続化された状態を読み戻して ML ワークロードを再開します。

ワークロードを復元できない場合、Pathways はアクセラレータ ジョブをシャットダウンし、エラスティック トレーニングが構成されている場合は、ワークロードに障害を転送します。エラスティック トレーニングが構成されていない場合、GKE は JobSet の再起動ポリシーに基づいてワークロード全体を再起動します。

JAX を使用して定義された一般的な ML ワークロードは、ステートレスな Pathways XLA コンポーネントに依存しています。これらのコンポーネントは、高帯域幅メモリ（HBM）スナップショットを使用して復元できます。JAX コロケーション Python API を使用して定義された ML ワークロードなど、一部の ML ワークロードは、ステートフルな Pathways コンポーネントに依存しており、復元できません。

弾力性のあるトレーニング

エラスティック トレーニングを使用すると、ハードウェア障害やスライスのプリエンプションが発生した場合でも、ワークロード全体を再起動することなく、トレーニング ジョブを続行できます。マルチスライス トレーニング ジョブでは、個々のスライスが使用できなくなっても、Pathways は単一コントローラのクライアント プロセスと正常な TPU スライスを維持します。

エラスティック トレーニングでは、Pathways システムの機能とトレーニング コードの復元ロジックを組み合わせます。

  • 障害検出: TPU スライスでハードウェア障害またはプリエンプションが発生すると、Pathways は影響を受けるスライスの割り当てを解除し、ワークロードがそのスライスのデータに次回アクセスしたときに jax.errors.JaxRuntimeError を発生させます。クライアント プロセスはアクティブなままになるため、コードでこのエラーをキャッチし、リソースを再構成して、トレーニングを続行できます。
  • リカバリモード（Recovery mode）: pathwaysutils.elastic.manager.Manager クラスとその elastic_retry デコレータは、次の 2 つの復元戦略をサポートしています。
    • 一時停止と再開: minimum_slice_count を None またはスライスの合計数に設定すると、スライスが失敗したときにトレーニングが一時停止し、すべてのスライスが正常な状態に戻ると再開されます。スライス数は試行間で変更されないため、ワークロードは既存のデバイスメッシュとコンパイル済みの実行可能ファイルを再利用できます。
    • レプリカのサイズ変更: minimum_slice_count を合計スライス数より少なく設定すると、スライスが失敗した場合に、デバイスメッシュを再構築してより小さなトポロジ用に再コンパイルすることで、残りのアクティブなスライスでトレーニングが続行されます。バックグラウンド モニター スレッドは、非アクティブなスライスを定期的にチェックし、新しく正常になったスライスを elastic_manager.available_inactive_slices に記録します。トレーニング ループは、交換用スライスを使用してデバイスメッシュを再構築し、拡張されたトポロジ用に再初期化して、現在の状態を新しいメッシュに配置できます。
  • 状態の復元: エラスティック イベント後にトレーニングを再開するために、ワークロードはインメモリ スナップショットを使用してモデル パラメータ、オプティマイザーの状態、データ イテレータの状態を復元します。orbax.checkpoint の Snapshotter クラスは、複製された状態を TPU HBM から各スライスの pinned_host CPU メモリに非同期でコピーします。復元中、Snapshotter は、残存しているホストレプリカから状態を再構築し、リモート ストレージから読み取ることなく、アクティブな TPU メッシュに直接再シャーディングします。

Pathways プロキシ サーバーでエラスティシティを有効にする

GKE の Pathways でエラスティック デバイスセットを有効にするには、JobSet マニフェストの pathways-proxy コンテナの args フィールドに --num_elastic_instances フラグを追加し、ワークロード内のエラスティック TPU スライスの数を指定します。

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

次のように置き換えます。

  • BUCKET_NAME: Pathways の一時ファイルに使用される Cloud Storage バケット。
  • NUM_SLICES: ワークロード内の TPU スライスの数。

アクティブなデバイスをフィルタしてデバイス メッシュを構築する

Pathways ワークロードでは、jax.devices() はクライアント プロセスの存続期間の初期割り当てのすべてのデバイスを返します。エラスティック トレーニングを使用する場合は、デバイス メッシュとバッチサイズが正常でアクティブなスライスのみを使用するように elastic_manager.active_devices を使用します。

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

レプリカのサイズ変更モードでアクティブなスライスの数が変更された場合は、grain.experimental.ElasticIterator などのエラスティック データローダーを使用します。標準のホストごとのデータ イテレータは、ホストごとに状態を保存し、異なる数のホストに復元することはできません。一方、ElasticIterator は、任意のアクティブなスライスで復元できる単一のグローバル データセット位置を追跡します。

弾力性のあるトレーニング ループ

エラスティック トレーニング ループは、次のコード例に示すように、スライス モニタリング、障害復旧、インメモリ スナップショットを調整します。

  • train 関数は、再試行ラッパーの外部に単一の manager.Manager() インスタンスと Orbax Snapshotter を作成します。これにより、スライスのトラッキングとホストメモリのスナップショットが試行間で保持されます。
  • @elastic_manager.elastic_retry を run_training_loop に適用すると、関数は正常なスライスを待機し、jax.default_device をアクティブなスライスに設定し、バックグラウンド モニター スレッドをアクティブに保ち、jax.errors.JaxRuntimeError で再試行するように構成されます。
  • 復元を試みる際、コードは elastic_manager.active_devices からデバイス メッシュを再構築し、snapshotter.load を呼び出して、更新されたメッシュ シャーディングに最新のスナップショットを復元する必要があります。失敗した試行のデバイス バッファをスコープ外にして、JAX がメモリを解放できるようにします。
  • ループの各ステップで、コードは定期的に snapshotter.save を呼び出してトレーニング状態をバックアップし、elastic_manager.available_inactive_slices をチェックして正常な交換スライスを検出します。置換スライスが利用可能になると、コードはそれらを elastic_manager.active_slice_indices に追加し、メッシュを再構築して、トレーニング ループを再初期化し、状態を拡張メッシュに配置します。

次の例は、manager.Manager、elastic_retry、Orbax のインメモリ Snapshotter を組み合わせて、エラスティック トレーニング ループを実装する方法を示しています。

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

エラスティック マネージャーと再試行動作を構成する

manager.Manager(devices=None) はプロセスごとに 1 回インスタンス化します。devices を省略すると、マネージャーは jax.devices() を使用し、slice_index 属性でデバイスをグループ化します。

Manager インスタンスは、トレーニング中にクラスタ トポロジを検査するために、次のプロパティを公開します。

  • slice_to_devices: 各 slice_index から、そのスライスに属する jax.Device オブジェクトのシーケンスへのマッピング。
  • all_slice_indices: 割り当て内のすべてのスライス インデックスの frozenset。
  • active_slice_indices: アクティブなコンピューティング メッシュに参加しているスライス インデックスの frozenset。
  • active_devices: アクティブなスライスに属する jax.Device オブジェクトのリスト。slice_index で並べ替えられます。
  • inactive_slice_indices: 割り当て内のスライス インデックスの frozenset。コンピューティング メッシュに参加していません。
  • available_inactive_slices: バックグラウンド モニター スレッドが正常で再参加の準備ができていると検出した、非アクティブなスライス インデックスの frozenset。
  • total_slice_count と active_slice_count: 割り当てられたスライスの合計数とアクティブなスライスの数。
  • default_device: アクティブなスライスのデバイス 0。elastic_retry は、各試行のデフォルトの JAX デバイスとして自動的に設定されます。

@elastic_manager.elastic_retry に次の引数を渡すことで、再試行とスケーリングの動作をカスタマイズできます。

  • minimum_slice_count: 最初の実行または再試行を開始する前に必要な正常なスライスの最小数。None（デフォルト）に設定すると、elastic_retry は割り当て（total_slice_count）内のすべてのスライスを待機し、一時停止と再開モードで動作します。より小さい整数（1 など）に設定すると、elastic_retry はレプリカのサイズ変更モードで動作します。
  • poll_interval: トレーニング中にエラスティック トレーニング マネージャー（Manager）が実行するスライス ヘルスチェックの間隔（秒単位）。デフォルトは 10 です。
  • timeout: TimeoutError が発生する前に minimum_slice_count スライスがアクティブになるまで待機する最大秒数。デフォルトは None（タイムアウトなし）です。
  • pre_callback: マネージャーが active_slice_indices を更新した後、ラップされたトレーニング ループ関数が実行される前に、各試行の開始時に実行されるオプションのコールバック。
  • on_elastic_event_callback: elastic_retry がスライスダウン jax.errors.JaxRuntimeError をキャッチした後、次の再試行が開始される前に実行されるオプションのコールバック。このコールバックを使用して、不完全なチェックポイントをクリーンアップしたり、データ イテレータをリセットしたりできます。
  • retry_policy: エラスティック イベント後に再試行を続行するかどうかを決定する呼び出し可能な (attempt: int, error: Exception) -> bool。manager.ElasticRetryLimit(max_attempts) を使用して、試行の合計数の上限を設定します。省略すると、elastic_retry は無期限に再試行します。

エラスティック イベントでは、スタック フレームがアンワインドされ、それらの配列を参照する変数がない場合、JAX は失敗した試行からデバイス バッファを解放します。失敗した試行のデバイス配列への参照をグローバル変数や存続期間の長いコールバック内に保持しないようにして、次の試行が開始される前にメモリが解放されるようにします。

Orbax を使用したメモリ内スナップショット

ocp.v1.training.pathways.Snapshotter クラス（orbax-checkpoint バージョン 0.12.6 以降で使用可能）は、スライス間で pinned_host メモリにトレーニング状態のメモリ内バックアップを保持します。

  • Snapshotter(replica_axis_index=0): スナップショットを初期化し、データ並列スライス レプリカを表すメッシュ軸を指定します。保存された状態のすべての JAX 配列は、このメッシュ軸に沿って複製する必要があります。これにより、スライスが失敗した場合でも、残りのスライスにモデルとオプティマイザーの状態の完全なコピーが保持されます。状態に単一デバイス配列が含まれている場合は、save を呼び出す前にメッシュ全体に複製します。
  • save(step, state): state を pinned_host メモリに非同期で転送し、replica_axis_index に沿ってレプリカごとのツリーに分割します。以前のスナップショット転送がまだ進行中の場合、save は新しいリクエストをスキップして、トレーニング ループが停止しないようにします。
  • load(abstract_state, reset_snapshot_state=True): 保存されたホストレプリカをプローブして、存続しているスライスを特定し、正常なレプリカを連結して、abstract_state で指定されたターゲット シャーディングに状態を再シャーディングします。
  • join(): キューに格納されたスナップショット オペレーションが完了するまでブロックします。
  • latest: 最後に完了したスナップショットのメタデータ（step を含む）を返します。完了したスナップショットがない場合は None を返します。

スライスの失敗時にシャーディングを減らす

実行中に TPU スライスが失敗するか、プリエンプトされると、Pathways は jax.errors.JaxRuntimeError（DATA_LOSS、DEADLINE_EXCEEDED、NOT_FOUND、INTERNAL など）を発生させます。@elastic_manager.elastic_retry で装飾された関数では、elastic_retry がこの例外をキャッチし、elastic.is_error_due_to_slice_down(error) を確認します。

  • エラーがスライスダウン イベントによって発生したものでない場合、elastic_retry はすぐに例外を再発生させます。
  • エラーがスライスダウン イベントによって発生した場合、elastic_retry はアクティブな JAX プロファイラ トレースを停止し、on_elastic_event_callback を実行して retry_policy(attempt, error) を評価します。
  • 次の試行を開始する前に、elastic_retry は elastic.wait_for_slices を 1 回呼び出して、少なくとも minimum_slice_count スライスがアクティブになるまでスライスの健全性をポーリングします。次に、elastic_manager.active_slice_indices を更新し、残りのスライスでトレーニング ループ関数を実行します。

スライスが再結合したときにシャーディングを増やす

ワークロードがレプリカ サイズ変更モードで削減された数のスライスで実行されている場合、バックグラウンド モニター スレッドは elastic_manager.inactive_slice_indices で定期的にヘルスチェックを実行します。修復または再スケジュールされたスライスがヘルスチェックに合格すると、モニタースレッドはスライス インデックスを elastic_manager.available_inactive_slices に追加します。

トレーニング ステップの終了時に elastic_manager.available_inactive_slices が空でない場合、ワークロードはトレーニングの進行状況を失うことなく、その場でスケールアップします。

  • elastic_manager.available_inactive_slices のスライス インデックスを elastic_manager.active_slice_indices に追加すると、再結合されたスライスが elastic_manager.active_devices に組み込まれ、available_inactive_slices からクリアされます。
  • elastic_manager.active_devices からデバイスメッシュを再構築し、トレーニング ステップとデータ イテレータを再初期化することで、拡張されたトポロジに計算を適応させます。
  • jax.device_put を呼び出すと、スナップショットから復元せずに、現在のトレーニング状態が新しいメッシュ シャーディングに配置されます。

多層チェックポイント処理

大規模なトレーニング ジョブでは中断が頻繁に発生し、Cloud Storage との間でチェックポイントを直接書き込みまたは復元すると、トレーニングの進行がボトルネックになる可能性があります。多層チェックポイント処理では、GKE の 3 つのストレージ階層を調整することで、チェックポイント処理のオーバーヘッドと復旧時間を短縮します。

  • ローカル RAM ディスク: TPU ワーカーは、各ワーカーノードにマウントされたインメモリ CSI ボリューム（multitier-checkpoint.csi.storage.gke.io）に高頻度のチェックポイントを書き込みます。このボリュームは、同じノードでの Pod の再起動後も保持されます。
  • クラスタ内ピア レプリケーション: GKE レプリケーター サービスは、クラスタ内のピア TPU ワーカー間でローカル チェックポイント シャードを複製します。
  • Cloud Storage バックアップ: レプリケーター サービスは、TPU トレーニング ステップをブロックすることなく、階層型名前空間が有効になっている Cloud Storage バケットにローカル チェックポイントを非同期でバックアップします。

各 TPU ホストが独自のトレーニング スクリプトを実行するマルチコントローラ JAX とは異なり、Pathways は pathways-head Pod で単一のコントローラを実行し、pathways-worker Pod に作業をディスパッチします。Pathways は、単一のコントローラを介してチェックポイント データをルーティングすることなく、各 TPU ホストのローカル RAM ディスクでチェックポイント シャードを直接読み書きするために、コロケーション Python を使用します。各 worker Pod は、pathways-worker コンテナとともに colocated-python-sidecar コンテナを実行します。両方のコンテナが、多層チェックポイント処理 CSI ボリュームと共有メモリ ディレクトリ（/tmp/sidecar）をマウントします。初期化とトレーニング中に、Orbax はチェックポイントの設定、保存、復元を Pathways コントローラから各 TPU ホストのワーカー サイドカーにディスパッチします。Pathways のマルチティア チェックポイントは、アクティブな TPU スライスの数をスケールダウンまたはスケールアップする際のエラスティック トレーニングもサポートしています。

Pathways は、MaxText での多層チェックポイント処理をサポートしています。設定と構成手順については、MaxText マルチティア チェックポイント ガイドをご覧ください。

Pathways ワークロードを構成する

マルチティア チェックポイント用に Pathways ワークロードを構成するには、Cluster Toolkit（gcluster job submit）を使用するか、カスタム Pathways JobSet マニフェストを構成します。

Cluster Toolkit を使用する

gcluster job submit --pathways を使用してジョブを送信する場合は、--gke-mtc-enabled、--gke-mtc-ramdisk-dir、--pathways-colocated-python-sidecar-image フラグを渡します。

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

SIDECAR_IMAGE は、コロケーションされた Python サイドカー コンテナ イメージの URI に置き換えます。

これらのフラグを設定すると、gcluster は次のものを自動的に構成します。

  • --sidecar_name=external を pathways-proxy コンテナに渡します。
  • 各 worker Pod で colocated-python-sidecar コンテナを実行し、pathways-worker と colocated-python-sidecar の両方に共有 /tmp/sidecar ディレクトリ（--cloud_pathways_sidecar_shm_directory=/tmp/sidecar）をマウントします。
  • worker Pod の pathways-worker と colocated-python-sidecar の両方で --gke-mtc-ramdisk-dir で指定されたパスに multitier-checkpoint.csi.storage.gke.io CSI ボリュームをマウントします。pathways-head Pod にはマウントしません。

カスタムの Pathways JobSet マニフェストを使用する

gcluster を使用せずに独自の Pathways JobSet マニフェストを管理する場合は、--sidecar_name=external を pathways-proxy に渡します。worker Pod 仕様で、共有サイドカー メモリ ボリュームと multitier-checkpoint.csi.storage.gke.io CSI ボリュームの両方を pathways-worker コンテナと colocated-python-sidecar コンテナにマウントします。

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

カスタム JAX ワークロードで Orbax 多層チェックポイント処理を使用する

Pathways で実行されるカスタム JAX トレーニング スクリプトで、orbax.checkpoint.experimental.emergency.multi_tier_checkpointing パッケージを使用し、initialize_multi_tier_checkpointing と ReplicatorCheckpointManagerOptions の両方で use_colocated_python=True を設定します。

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

ホットスワップ

ホットスワップは、優先度の高いジョブが優先度の低いジョブからリソースを迅速に引き継ぎ、ダウンタイムを最小限に抑えて復旧を高速化できる GKE JobSet API の機能です。

JobSet を作成すると、GKE は JobSet 構成で指定された複数のスライスにワークロードをスケジュールします。1 つ以上のスライスでハードウェア障害が発生すると、GKE は影響を受ける Pod を失敗としてマークします。GKE クラスタに優先度の低いジョブを実行している予備のスライスがある場合、JobSet コントローラは、優先度の高いジョブから失敗したスライスのワークロードをその予備のスライスに再マッピングします。この再マッピングは通常、1 分以内に完了します。

JobSet が再起動すると、次の状況でホットスワップが発生する可能性があります。

  • デフォルト モード: 同じクラスタ内で予備のアイドル状態の TPU スライスが使用可能な場合、Kubernetes スケジューラは、障害が発生したスライスの修復を待つのではなく、再起動されたジョブをこれらのスライスにスケジュールすることを優先します。このスケジューリング動作により、復旧が迅速化されます。
  • 異種ワークロード: Kubernetes PriorityClass リソースが構成された複数のワークロードを実行しているクラスタでは、再起動された JobSet がホットスワップをトリガーできます。再起動されたジョブのアフィニティが優先度の低いジョブのリソースと一致する場合、Kubernetes は優先度の低いジョブをプリエンプトして、優先度の高いジョブがスライスの修復を待たずに開始できるようにします。たとえば、PriorityClass リソースを使用して、Pathways ワーカー Pod を異なる優先度で構成できます。

クラスタで優先度を使用するには、PriorityClass マニフェストを定義して high-prior-job.yaml として保存します。

apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

この YAML を GKE クラスタに適用します。

kubectl apply -f high-prior-job.yaml

次に、新しい優先度クラスを Pathways ワークロードに接続します。Cluster Toolkit を使用してワークロードを送信する場合は、--priority=high-prior-job を gcluster job submit コマンドに渡します。これにより、pathways-head Pod と worker Pod の両方に priorityClassName が設定されます。独自の JobSet マニフェストを作成する場合は、pathways-worker コンテナを実行する worker レプリケートされたジョブの Pod テンプレートで priorityClassName を設定します。

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

Pathways を有効にして Cluster Toolkit TPU ブループリントを使用してクラスタを作成した場合、クラスタにはすでに very-low、low、medium、high、very-high の優先度クラスが含まれているため、独自の優先度クラスを作成する代わりに、そのうちの 1 つを --priority に渡すことができます。

次のステップ