Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

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

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

一時停止 / 再開: カスタムのプリエンプション処理コードを必要とせずに、プリエンプション通知などの計画的な中断をワークロードで許容できます。
エラスティックトレーニング: トレーニングコードの復元ロジックを使用して、クライアントをクラッシュさせることなく、計画外のハードウェア障害にワークロードが耐えられるようにします。

始める前に

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

一時停止と再開

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

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

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

JAX を使用して定義された一般的な ML ワークロードは、高帯域幅メモリ（HBM）スナップショットを使用して復元可能なステートレス Pathways XLA コンポーネントに依存しています。JAX コロケーション Python API を使用して定義されたワークロードなど、特定の 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/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 を作成します。これにより、スライスのトラッキングとホストメモリのスナップショットが試行間で保持されます。
内部の run_training_loop 関数を @elastic_manager.elastic_retry でデコレートすると、正常なスライスを待機し、jax.default_device をアクティブなスライスに設定し、ステップ間でバックグラウンドモニタースレッドをアクティブに保ち、jax.errors.JaxRuntimeError が発生したときにループを再試行します。
試行の開始時にスライス障害から復旧する場合、コードは elastic_manager.active_devices からデバイスメッシュを再構築し、snapshotter.load を呼び出して、更新されたメッシュシャーディングに最新のスナップショットを復元する必要があります。また、失敗した試行のデバイスバッファをスコープ外にする必要があります。
ループの各ステップで、コードは定期的に 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: スライスダウン 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_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 を呼び出すと、スナップショットから復元せずに、現在のトレーニング状態が新しいメッシュシャーディングに配置されます。

ホットスワップ

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

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

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

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

クラスタで優先度を使用するには、PriorityClass マニフェストを定義します。

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-worker Pod の podspec に次のテキストを追加して、新しい優先度クラスを Pathways ワーカージョブにアタッチします。

priorityClassName: high-prior-job

Pathways を使用した復元力のあるトレーニング コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。