多層チェックポイント処理を使用して GKE で大規模な ML モデルをトレーニングする

Standard

このページでは、多層チェックポイント処理を使用して、GKE での ML モデルのトレーニング中にチェックポイントを確実に保存して管理する方法について説明します。チェックポイントの保存と管理は、数千を超えるノードを使用する大規模なトレーニングジョブにとって非常に重要です。このような大規模なジョブは頻繁に中断され（1 時間に 1 回の可能性あり）、復元に時間がかかることがあります。

利点

多層チェックポイント処理を使用すると、次のメリットがあります。

次のワークロードのバックアップ、レプリケーション、自動復元など、完全にオーケストレートされたチェックポイントデータの管理:
- TPU で実行される JAX トレーニング実行で状態管理に Orbax を使用。
- GPU 上の PyTorch ワークロード。
ローカルノードに保存されたチェックポイントからトレーニングジョブを迅速に復元。トレーニングクラスタ内の別のノードに保存されたチェックポイントを使用して復元することもできます。
最悪のシナリオ（クラスタ内チェックポイントがない場合）で、Cloud Storage バックアップに保存されたチェックポイントからトレーニングジョブを迅速に復元。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

Google Kubernetes Engine API を有効にする。

Google Kubernetes Engine API の有効化

このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。すでに gcloud CLI をインストールしている場合は、gcloud components update を実行して最新のバージョンを取得する。
注: gcloud CLI がすでにインストールされている場合には、必ず compute/region プロパティを設定してください。主にゾーンクラスタを使用する場合は、代わりに compute/zone を設定します。デフォルトのロケーションを設定することで、gcloud CLI のエラー（One of [--zone, --region] must be supplied: Please specify location など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

プロジェクトで使用する Cloud Storage バケットを作成します（まだ作成していない場合）。階層型名前空間を有効にしてください。有効にしないと、バックアップが失敗します。

要件

多層チェックポイント処理には、GKE クラスタバージョン 1.32.4-gke.1415000 以降が必要です。

制限事項

Autopilot クラスタはサポートされていません。

多層チェックポイント処理を使用するように GKE ノードを構成する

このセクションでは、新規と既存のクラスタで GKE ノードを構成する方法について説明します。

新しいクラスタでノードを構成する

多層チェックポイント処理、Cloud Storage FUSE CSI ドライバ、Workload Identity Federation for GKE を有効にしてクラスタを作成します。ML ワークロードに TPU スライスを使用する場合は、TPU スライスノードプールの構成を含めるようにクラスタ作成コマンドを調整する必要があります。
```
gcloud container clusters create CLUSTER_NAME \
    --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
    --node-locations=NODE_LOCATION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --cluster-version=CLUSTER_VERSION
    --location=CLUSTER_LOCATION \
    --machine-type=MACHINE_TYPE \
    --num-nodes=NUM_NODES
```
次の値を置き換えます。
- CLUSTER_NAME: クラスタの名前。
- NODE_LOCATION: クラスタノードのゾーン。ここの TPU 容量を使用します。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- CLUSTER_VERSION: クラスタのバージョン。サポートされている最小バージョンは 1.32.4-gke.1415000 です。
- CLUSTER_LOCATION: クラスタを作成するリージョン。
- MACHINE_TYPE: JobSet コントローラや多層チェックポイント処理コントローラなどのコンポーネントを実行するノードに使用されるマシンタイプ。大規模なトレーニングには、少なくとも e2-standard-4 マシンを使用することをおすすめします。このマシンタイプはモデルのトレーニングには使用しません。その目的のため、個別のノードプールを作成します。多くの場合、アクセラレータ最適化 VM ファミリーを使用します。
- NUM_NODES: クラスタの各ゾーンに作成するノードの数。

既存のクラスタでノードを構成する

既存のクラスタで多層チェックポイント処理を使用するには、次のコマンドを使用して、Cloud Storage FUSE CSI ドライバと Workload Identity Federation for GKE とともに有効にします。既存のクラスタバージョンは 1.32.3-gke.1170000 以降である必要があります。

Workload Identity Federation for GKE を有効にします。
```
gcloud container clusters update CLUSTER_NAME \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=CLUSTER_LOCATION
```
次の値を置き換えます。
- CLUSTER_NAME: クラスタの名前。
- PROJECT_ID: 実際の Google Cloud プロジェクト ID。
- CLUSTER_LOCATION: クラスタのリージョン。

多層チェックポイント処理と Cloud Storage FUSE CSI ドライバを有効にします。

gcloud container clusters update CLUSTER_NAME \
    --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
    --location=CLUSTER_LOCATION

多層チェックポイント処理を使用するための権限を構成する

このセクションでは、多層チェックポイント処理を使用するための権限を構成する方法について説明します。

Cloud Storage バケットへのアクセス権を付与する

多層チェックポイント処理の CSI ドライバで使用されるエフェメラルボリュームは、既存の Cloud Storage バケットを使用する必要があります。

チェックポイントを Cloud Storage バケットに保存するには、多層チェックポイント処理がバケットにアクセスする必要があります。バケットに対する Storage オブジェクトユーザー（roles/storage.objectUser）IAM ロールを、多層チェックポイント処理の Kubernetes サービスアカウントに付与します。

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

次の値を置き換えます。

GCS_BUCKET: データの転送元となる Cloud Storage バケットの名前。
PROJECT_ID: 実際の Google Cloud プロジェクト ID。
PROJECT_NUMBER: プロジェクトに対して自動的に生成される固有識別子。この値を確認するには、プロジェクトの作成と管理をご覧ください。

（省略可）Compute Engine のデフォルトサービスアカウントにアクセス権を付与する

Compute Engine インスタンスに Cloud Storage バケットへの読み取りアクセス権が必要な場合は、Compute Engine のデフォルトサービスアカウントにストレージオブジェクト閲覧者（roles/storage.objectViewer）IAM ロールを付与します。

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --role roles/storage.objectViewer

JobSet コントローラをデプロイする

JobSet コントローラは、GKE でモデルトレーニングを実行するバッチジョブの管理を担当します。そのリソース割り当ては、ワークロードを効率的に処理するように調整されます。トレーニングジョブランチャーが JobSet をデプロイして使用していることを確認します。

JobSet デプロイのマネージャーコンテナのメモリリクエストを 1 Gi、メモリ上限を 2 Gi、CPU リクエストを 1 に増やすには、次のパッチコマンドを実行します。

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

多層チェックポイント処理の CSI ドライバを初期化する

このセクションでは、ワークロードが実行されるノードで多層チェックポイント処理の CSI ドライバを初期化する方法について説明します。CSI ドライバは、モデルのトレーニングプロセス中のチェックポイントの保存と管理を担当します。

CheckpointConfiguration を作成する

CheckpointConfiguration は、多層チェックポイント処理の CSI ドライバのデプロイのプロパティを指定する Kubernetes カスタムリソースです。このリソースはクラスタスコープです。

次の checkpoint.yaml マニフェストを作成します。
```
apiVersion: checkpointing.gke.io/v1
kind: CheckpointConfiguration
metadata:
  name: MTC_CONFIG_NAME-configuration
spec:
    cloudStorageBucketName: GCS_BUCKET
    nodeSelector:
        node.kubernetes.io/instance-type: MACHINE_TYPE
    tolerations:
    - key: TOLERATION_KEY
        operator: Exists
        effect: NoSchedule
    inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
    gcsFuseMountOptions:
    - implicit-dirs
    - metadata-cache:negative-ttl-secs:0
    - metadata-cache:ttl-secs:-1
    - metadata-cache:stat-cache-max-size-mb:-1
    - metadata-cache:type-cache-max-size-mb:-1
    - file-cache:max-size-mb:-1
    - file-cache:cache-file-for-range-read:true
    - file-system:kernel-list-cache-ttl-secs:0
    - file-cache:enable-parallel-downloads:true
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true
```
次のように置き換えます。
- MTC_CONFIG_NAME: CheckpointConfiguration の名前。この名前はクラスタに対してグローバルであり、ジョブ固有ではありません。
- GCS_BUCKET: チェックポイントデータを保存する Cloud Storage バケットの名前。権限を使用して Cloud Storage バケットを設定するの手順で設定したバケットを使用します。
- MACHINE_TYPE: 対応するアクセラレータのマシンタイプ。値は次のいずれかになります。
  - TPU v5p: ct5p-hightpu-4t
  - TPU v5e: ct5e-hightpu-4t
  - TPU v6e: ct6e-standard-4t
  - NVIDIA H100 80 GB GPU（A3 シリーズ）:
    - A3 High: a3-highgpu-8g
    - A3 Mega: a3-megagpu-8g
    - A3 Ultra: a3-ultragpu-8g
  GKE で、分散ワークロードを GPU 上で実行する方法については、マルチインスタンス GPU の実行をご覧ください。TPU の場合は、TPU スライスノードプールを作成するをご覧ください。
- TOLERATION_KEY: このフィールドを使用すると、一致する taint を持つノードに CSI ドライバをスケジュールできます。さまざまなアクセラレータタイプでの taint の仕組みについては、次のページをご覧ください。
  - TPU
  - GPU
- IN_MEMORY_VOLUME_SIZE: インメモリチェックポイントキャッシュのサイズ。数量と単位（200 Gi など）を指定します。この値は次のようになります。
  - TPU のローカルチェックポイントサイズに 2.2 を掛けた値
  - 単一ピアを持つ GPU のローカルチェックポイントサイズに 6.6 を掛けた値。
次のようにマニフェストを適用します。
```
kubectl apply -f checkpoint.yaml
```

CSI ドライバが実行されていることを確認します。

kubectl get pod -n gke-managed-checkpointing

出力例を以下に示します。エントリは複数あり、高速化されたノードごとに 1 つのエントリがあります。

NAME                                                          READY   STATUS    RESTARTS   AGE
multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s

多層チェックポイント処理の CSI ドライバをアンインストールする

多層チェックポイント処理の CSI ドライバのデプロイを解除する場合は、CheckpointConfiguration リソースを削除します。多層チェックポイント処理コントローラは、ノードから CSI ドライバを削除します。これにより、RAM ディスクが削除され、他のワークロード用にメモリが解放されます。例:

kubectl delete -f checkpoint.yaml

Cloud Storage バックアップのデータ保持とガベージコレクションを管理する

チェックポイントのバックアップを Cloud Storage に保持するためのポリシーは、お客様の責任で実装していただく必要があります。多層チェックポイント処理は、チェックポイントのバックアップを Cloud Storage に書き込むだけで、変更や削除は行いません。

多くのオープンソースツールで、保持とガベージコレクションを処理できます。たとえば、次のようなツールがあります。

次の例では、Cloud Storage FUSE を使用するバックアップの場所に backup ディレクトリがマウントされている backup-warden を使用します。

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

ワークロード JobSet マニフェストを更新する

大規模なチェックポイントボリュームを含めるように、ジョブの JobSet マニフェストを更新します。詳細はワークロードによって異なります。

たとえば、GKE に TPU マルチスライスをデプロイするのサンプル JobSet を拡張するには、次の操作を行います。

jax-tpu コンテナに次の行を追加します。
```
volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR
```
CHECKPOINT_DIR は、チェックポイントディレクトリのパスに置き換えます。これは、replicator.yaml が生成され、多層チェックポイント処理がチェックポイント保存のオペレーションを実行する場所です。詳細については、アプリケーションに多層チェックポイント処理を統合するをご覧ください。

Job 仕様の spec.template.spec フィールドに次の行を追加します。

volumes:
- name: checkpoint
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

アプリケーションに多層チェックポイント処理を統合する

チェックポイントの場所とレプリケーションの準備状況に関する情報を共有するには、次のプロトコルを使用して多層チェックポイント処理と通信するようにアプリケーションを変更します。

スタートアップ

このセクションでは、アプリケーションが多層チェックポイント処理とやり取りするために必要な初期手順について説明します。

レプリケーターは、多層チェックポイント処理のコアコンポーネントであり、CSI ドライバの一部としてすべてのノードで実行されます。レプリケーターは、ローカル RAM ディスクからピアノード、Cloud Storage などの外部ストレージまで、ストレージ階層間のチェックポイントレプリケーションを管理します。

replicator.yaml ファイルは、ML トレーニングジョブ（フレームワークコード）とレプリケーターコンポーネントの間の動的コントロールプレーンとして機能します。ML アプリケーションは、トレーニングジョブとレプリケーターサービスの両方からアクセス可能なローカルボリューム（RAMDisk）にこのファイルをプログラムで生成します。このマニフェストにより、ML フレームワークは、バックエンド設定時に定義された静的インフラストラクチャパラメータ（Cloud Storage のアップロード頻度など）とは別に、ランタイム構成とライフサイクル管理の手順をレプリケーターに提供できます。

このインタラクションの具体的な例については、以下をご覧ください。

Cloud TPU の JAX にこのアーキテクチャを利用する MaxText プロジェクト。
このアーキテクチャを GPU 上で PyTorch と NVIDIA NeMO とともに利用する PyTorch リファレンスの例。

アプリケーションは、起動時に次の処理を実行する必要があります。

replicator.yaml ファイルがなくなるまで待ちます。これは、レプリケーターがアプリケーションによって構成される準備ができたことを示します。replicator.yaml ファイルは、ワークロード JobSet マニフェストを更新するセクションで構成した CHECKPOINT_DIR の場所に生成されます。

モデルトレーニングジョブが最初に作成されたとき、replicator.yaml ファイルは存在しないため、アプリケーションはすぐに続行できます。ただし、ジョブが再開された場合（障害や手動介入など）、システムが以前のジョブインスタンスをまだ処理している可能性があり、そのインスタンスの replicator.yaml がローカルボリュームにまだ存在している可能性があります。
アプリケーションまたは ML ジョブは、次のような構成で replicator.yaml ファイルを作成します。
Orbax
```
job-name: orbax
framework: orbax
assume-data-parallelism: 3
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30
```
PyTorch
```
job-name: nemo
framework: pytorch.distributed
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30
```
この構成例には次のフィールドがあります。
- name: トレーニングジョブの名前。
- framework: トレーニングジョブで使用されている ML フレームワーク。
- node-rank: 分散トレーニングジョブの現在のノードの固有識別子。これは、このファイルを作成するノードのノードランクを表します。実行に参加する各ノードには独自のランクが割り当てられます。
- nodes: 分散トレーニングジョブに参加しているノードの合計数。この値は、Pod のメタデータから取得されます。ML トレーニングジョブでこの値を確認することもできます。
- peer-ranks または peers-per-node: レプリケーショントポロジを指定する 2 つの方法。これらの 2 つのパラメータのうち、1 つのみが存在する必要があります。
  - peer-ranks: 現在のノードのチェックポイントデータを複製するピアノードの明示的なランク。これにより、どのノードをレプリケーションパートナーとして使用するかをきめ細かく制御できます。
  - peers-per-node: レプリケーターがレプリケーション用に自動的に選択するノードあたりのピアノード数。
- backup-interval-minutes: チェックポイントが Cloud Storage にバックアップされる頻度（分単位）。この値は 30 分以上に設定することをおすすめします。
新しい replicator.yaml ファイルがシステムによって削除されるまで待ちます。これは、レプリケーターが再起動してクリーンアップを実行したことを意味します。この手順により、アプリケーションが次のセクションの手順を実行するときに、ローカルボリュームに古いファイルや一時ファイルが残るのを防ぐことができます。

最後に確認された正常な（LKG）チェックポイントから復元する

レプリケーターが初期化されると、多層チェックポイント処理は TPU または GPU ワーカーごとに 1 つのシンボリックリンクを作成します。これらのシンボリックリンクは、replicator.yaml ファイルと同じマウントされたローカルボリュームに作成されます。このボリュームにジョブはチェックポイントを保存します。

シンボリックリンクの形式は <job-name>-s{step}-n<node-rank>-w<worker-index>.restore です。
対応する .restore ファイルから各ワーカーを復元します。例については、次のセクションの Orbax 複製チェックポイントマネージャーの例をご覧ください。

チェックポイントを保存する

トレーニングジョブの進行中に、アプリケーションはこれらの処理を複数回実行します。保存オペレーションは、ワークロード JobSet マニフェストを更新するで構成した CHECKPOINT_DIR ロケーションで行われます。

Orbax

Orbax チェックポイントを作成します。ディレクトリにはステップ番号が付けられます。レプリケーターは、新しく作成されたチェックポイントディレクトリを検出して、必要に応じてレプリケーションまたはバックアップを実行し、自動的にクリーンアップします。

Orbax レプリケーターチェックポイントマネージャーの使用方法の詳細については、MaxtTest checkpointing ファイルをご覧ください。レプリケーターサービスのインタラクションの例については、MaxText max_utils ファイルをご覧ください。

PyTorch

InClusterLocalCheckpointIO をカスタム pytorch_lightning.CheckpointIO として使用して、ローカルストレージで正しい分散チェックポイントを有効にします。次のコマンド例では、NVIDIA NeMo フレームワーク上に構築されたリファレンス実装を使用して、多層チェックポイントを有効にします。

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

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

CHECKPOINT_DIR: チェックポイントディレクトリのパス。
JOB_NAME: トレーニングジョブワークロードの名前。

クラスタのアップグレード

クラスタのアップグレードでは、アップグレードの前または後に CheckpointConfiguration オブジェクトを削除して再作成できます。このオブジェクトによって動的にデプロイされるノードチェックポイントドライバの DaemonSet は自動的にアップグレードされないため、この操作が必要になります。

DaemonSet 仕様を同じに保つことを特に希望する場合は、何もする必要はありません。

トラブルシューティング

このセクションでは、多層チェックポイント処理に関する問題のトラブルシューティングについて説明します。一般的なストレージのトラブルシューティングについては、GKE での Cloud Storage のトラブルシューティングをご覧ください。

多層チェックポイント処理が有効になっていない

次のエラーは、クラスタで多層チェックポイント処理が有効になっていないことを示します。

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

このエラーは、CheckpointConfiguration を作成するの手順で kubectl apply -f checkpoint.yaml を実行した後に発生することがあります。

この問題を解決するには、次のコマンドを使用して、クラスタで多層チェックポイント処理が有効になっているかどうかを確認します。

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

多層チェックポイント処理が有効になっている場合、出力は次のようになります。

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

多層チェックポイント処理が無効になっている場合は、クラスタを更新し多層チェックポイント処理を有効にします。

多層チェックポイント処理の CSI ドライバがボリュームをマウントできない

この問題は、CSI ドライバが Cloud Storage ボリュームをマウントできない場合に発生することがあります。このような行が複数あることがあります。

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

この問題を解決するには、次の例に示すように、CSI ドライバ Pod イベントを確認します。

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

例に示すように、Cloud Storage バケットの PermissionDenied エラーが原因で問題が発生した場合は、権限を正しく設定することで問題を解決できます。

多層チェックポイント処理を使用して GKE で大規模な ML モデルをトレーニングする コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

利点

始める前に

要件

制限事項

多層チェックポイント処理を使用するように GKE ノードを構成する

新しいクラスタでノードを構成する

既存のクラスタでノードを構成する

多層チェックポイント処理を使用するための権限を構成する

Cloud Storage バケットへのアクセス権を付与する

（省略可）Compute Engine のデフォルト サービス アカウントにアクセス権を付与する

JobSet コントローラをデプロイする

多層チェックポイント処理の CSI ドライバを初期化する

CheckpointConfiguration を作成する

多層チェックポイント処理の CSI ドライバをアンインストールする

Cloud Storage バックアップのデータ保持とガベージ コレクションを管理する

ワークロード JobSet マニフェストを更新する

アプリケーションに多層チェックポイント処理を統合する

スタートアップ

Orbax

PyTorch

最後に確認された正常な（LKG）チェックポイントから復元する

チェックポイントを保存する

Orbax

PyTorch

クラスタのアップグレード

トラブルシューティング

多層チェックポイント処理が有効になっていない

多層チェックポイント処理の CSI ドライバがボリュームをマウントできない

次のステップ

多層チェックポイント処理を使用して GKE で大規模な ML モデルをトレーニングする

（省略可）Compute Engine のデフォルトサービスアカウントにアクセス権を付与する

Cloud Storage バックアップのデータ保持とガベージコレクションを管理する