Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

使用多層檢查點功能，在 GKE 上訓練大型機器學習模型

標準

本頁面說明如何在 GKE 訓練機器學習模型時，使用多層級檢查點功能，可靠地儲存及管理檢查點。檢查點的儲存和管理對於大規模訓練作業至關重要，這類作業是指使用超過數千個節點的作業。這類大規模工作經常中斷 (可能每小時一次)，且復原速度可能較慢。

優點

使用多層級檢查點功能可享以下優點：

全面調度檢查點資料管理，包括備份、複製和自動還原下列工作負載：
- 使用 Orbax 進行狀態管理的 JAX 訓練執行作業，在 TPU 上執行。
- 在 GPU 上執行 PyTorch 工作負載。
從儲存在本機節點的查核點快速復原訓練工作。您也可以使用訓練叢集中另一個節點儲存的檢查點進行復原。
在最壞的情況下，如果叢集內沒有查核點，可以從儲存在 Cloud Storage 備份的查核點快速還原訓練工作。

事前準備

開始之前，請確認您已執行下列工作：

啟用 Google Kubernetes Engine API。

啟用 Google Kubernetes Engine API

如要使用 Google Cloud CLI 執行這項工作，請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI，請執行 gcloud components update 指令，取得最新版本。較舊的 gcloud CLI 版本可能不支援執行本文中的指令。
注意：如果是現有的 gcloud CLI 安裝項目，請務必設定 compute/region 屬性。如果您主要使用可用區叢集，請改為設定 compute/zone。設定預設位置後，即可避免 gcloud CLI 發生錯誤，例如：One of [--zone, --region] must be supplied: Please specify location。如果叢集位置與您設定的預設位置不同，您可能需要在特定指令中指定位置。

如果專案沒有 Cloud Storage bucket，請建立一個。請務必啟用階層命名空間，否則備份會失敗。

需求條件

多層檢查點需要 GKE 叢集 1.32.4-gke.1415000 以上版本。

限制

不支援 Autopilot 叢集。

設定 GKE 節點，以便使用多層級檢查點

本節說明如何在新叢集和現有叢集上設定 GKE 節點。

設定新叢集的節點

建立叢集時，請啟用多層檢查點、Cloud Storage FUSE CSI 驅動程式和 Workload Identity Federation for GKE。如果使用 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 bucket 存取權

多層級檢查點 CSI 驅動程式使用的臨時磁碟區，必須使用現有的 Cloud Storage bucket。

如要將檢查點儲存在 Cloud Storage bucket，多層級檢查點功能必須具備 bucket 存取權。將 bucket 的「Storage 物件使用者」(roles/storage.objectUser) IAM 角色授予 Multi-Tier Checkpointing 的 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 bucket 名稱。
PROJECT_ID：您的 Google Cloud 專案 ID。
PROJECT_NUMBER：系統自動為專案產生的專屬 ID。如要找出這個值，請參閱「建立及管理專案」。

(選用) 授予 Compute Engine 預設服務帳戶存取權

如果 Compute Engine 執行個體需要 Cloud Storage bucket 的讀取權限，請將「Storage 物件檢視者」(roles/storage.objectViewer) IAM 角色授予 Compute Engine 預設服務帳戶。

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 是 Kubernetes 自訂資源，用於指定部署多層檢查點 CSI 驅動程式的屬性。這項資源的範圍為叢集。

建立下列 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-system:kernel-list-cache-ttl-secs:0
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true
```
更改下列內容：
- MTC_CONFIG_NAME：CheckpointConfiguration 的名稱。這個名稱適用於整個叢集，並非工作專屬。
- GCS_BUCKET：用於儲存檢查點資料的 Cloud Storage bucket 名稱。使用您在「設定具有權限的 Cloud Storage bucket」步驟中設定的 bucket。
- MACHINE_TYPE：對應加速器的機型。可能的值如下：
  - TPU v5p：ct5p-hightpu-4t
  - TPU v5e：ct5e-hightpu-4t
  - TPU v6e：ct6e-standard-4t
  - Ironwood (TPU7x)：tpu7x-standard-4t
  - NVIDIA H100 80GB GPU (A3 系列)：
    - A3 High：a3-highgpu-8g
    - A3 Mega：a3-megagpu-8g
    - A3 Ultra：a3-ultragpu-8g
  如要進一步瞭解如何使用 GKE 在 GPU 上執行分散式工作負載，請參閱「執行多實體 GPU」。如為 TPU，請參閱「建立 TPU 配量節點集區」。
- TOLERATION_KEY：這個欄位可讓 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

輸出內容大致如下。系統會顯示多個項目，每個加速節點各一個。

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，不會修改或刪除。

許多開放原始碼工具都能處理保留和垃圾回收作業，包括：

以下範例使用 backup-warden，其中 backup 目錄掛接至使用 Cloud Storage FUSE 的備份位置：

# 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 的位置，也是多層查核點執行查核點儲存作業的位置。詳情請參閱「在應用程式中整合多層檢查點」。

在工作規格的 spec.template.spec 欄位中，新增下列幾行。

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

在應用程式中整合多層級檢查點

如要分享檢查點位置和複製準備狀態的相關資訊，請修改應用程式，使用下列通訊協定與多層檢查點通訊。

下列架構和程式庫已導入這項整合功能：

MaxText (使用 Orbax)：MaxText 內建支援多層級檢查點 (MTC)。詳情請參閱 MaxText MTC 說明文件。
任何使用 Orbax 的模型：整合功能包含兩個主要元件：用於初始設定的 MTC 初始化模組 (範例請見 MaxText)，以及用於從本機 RAM 磁碟儲存及還原檢查點的 Replicator Checkpoint Manager (範例請見 MaxText)。
PyTorch 參考範例：這個範例適用於使用 NeMo 搭配 PyTorch 的工作負載。

啟動

本節說明應用程式與多層檢查點機制互動所需的初始步驟。

複製器是多層級檢查點的核心元件，在每個節點上執行，是 CSI 驅動程式的一部分。複製器會管理各儲存層的檢查點複製作業，從本機 RAM 磁碟到對等節點，以及 Cloud Storage 等外部儲存空間。

replicator.yaml 檔案是機器學習訓練工作 (框架程式碼) 和 Replicator 元件之間的動態控制層。機器學習應用程式會在本機磁碟區 (RAMDisk) 上，以程式輔助方式產生這個檔案，訓練工作和 Replicator 服務都能存取這個檔案。這個資訊清單可讓機器學習架構向 Replicator 提供執行階段設定和生命週期管理操作說明，與後端設定期間定義的靜態基礎架構參數 (例如 Cloud Storage 上傳頻率) 不同。

如需這項互動的具體範例，請參閱：

MaxText 專案，這個專案會將此架構用於 JAX on Cloud TPU。
PyTorch 參考範例：在 GPU 上搭配 PyTorch 和 NVIDIA NeMO 使用這個架構。

應用程式應在啟動期間執行下列步驟：

請等到 replicator.yaml 檔案消失，這表示 Replicator 已準備好由應用程式設定。replicator.yaml 檔案會在您於「更新工作負載 JobSet 資訊清單」一節中設定的 CHECKPOINT_DIR 位置生成。

首次建立模型訓練工作時，系統不會有 replicator.yaml 檔案，應用程式可以立即繼續。不過，如果工作重新啟動 (例如因失敗或手動介入)，系統可能仍在處理先前的例項，而該例項的 replicator.yaml 可能仍存在於本機磁碟區。
應用程式或機器學習工作會建立 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：訓練工作使用的機器學習架構。
- node-rank：分散式訓練工作中的目前節點專屬 ID。這代表建立此檔案的節點節點排名。參與執行的每個節點都會有自己的排名。
- nodes：參與分散式訓練工作的節點總數。這個值來自 Pod 的中繼資料。機器學習訓練工作也可以查看這個值。
- peer-ranks 或 peers-per-node：指定複製拓撲的兩種替代方式。這兩個參數只能提供其中一個。
  - peer-ranks：要複製目前節點檢查點資料的同層級節點明確排名。這樣一來，您就能精細控管哪些特定節點做為複製合作夥伴。
  - peers-per-node：複製器應自動選取用於複製的每個節點的對等節點數量。
- backup-interval-minutes：將檢查點備份至 Cloud Storage 的頻率 (以分鐘為單位)。建議您將這個值設為 30 分鐘以上。
等待系統刪除新的 replicator.yaml 檔案。這表示複製器已重新啟動並完成清理作業。應用程式執行下一節中的步驟時，這個步驟可避免本機磁碟區出現任何過時或暫存檔案。

從最後已知良好 (LKG) 查核點還原

初始化複製器後，多層級檢查點會為每個 TPU 或 GPU worker 建立一個符號連結。這些符號連結會建立在與 replicator.yaml 檔案相同的掛接本機磁碟區中，也就是工作儲存查核點的位置。

符號連結的形式為 <job-name>-s{step}-n<node-rank>-w<worker-index>.restore。
從對應的 .restore 檔案還原每個 worker。如需範例，請參閱下一節的 Orbax 複製檢查點管理員範例。

根據預設，LKG 查核點是多層級查核點功能可找到的最新完整查核點。系統會優先使用叢集內檢查點 (位於本機 RAM 磁碟和對等互連節點)，只有在沒有完整的叢集內檢查點時，才會使用 Cloud Storage 備份。如要覆寫這項選取項目，請參閱「從特定備份檔案還原」。

從特定備份還原

您可以從特定 Cloud Storage 備份還原工作，而不是從 LKG 檢查點還原，例如將訓練作業還原至較早的狀態。工作開始前，請先完成下列步驟：

找出要用來還原的備份。多層檢查點會將備份檔寫入您在 CheckpointConfiguration 中設定的 Bucket，並以工作名稱命名目錄。每個備份都是一個子目錄，以備份時的世界標準時間命名，格式為 YYYY-MM-DD_HH-MM：
```
gcloud storage ls gs://GCS_PATH/JOB_NAME/
```
更改下列內容：
- GCS_PATH：多層級檢查點寫入備份的 Cloud Storage 路徑。這個值是您在 CheckpointConfiguration 的 cloudStorageBucketName 欄位中指定的 bucket，包括 bucket 內的任何目錄。
- JOB_NAME：訓練工作的名稱，如 replicator.yaml 檔案的 job-name 欄位所指定。
在 bucket 的工作目錄中，建立名為 restore.version 的檔案。檔案只能包含備份目錄的名稱，例如：
```
echo "BACKUP_NAME" | gcloud storage cp - gs://GCS_PATH/JOB_NAME/restore.version
```
將 BACKUP_NAME 替換為備份目錄的名稱，例如 2025-06-30_14-30。
啟動工作。

如果存在 restore.version 檔案，多層檢查點的行為如下：

所有節點都會從指定的備份檔還原。即使叢集內查核點比備份新，也會遭到忽略。
工作在還原後儲存第一個檢查點時，系統會自動刪除 restore.version 檔案。覆寫作業只會套用一次：如果工作稍後重新啟動，則會照常從 LKG 檢查點還原。
如果指定的備份不存在或不完整，工作就會在沒有檢查點的情況下啟動。多層查核點不會回溯至叢集內查核點或其他備份。

儲存檢查點

訓練工作進行期間，應用程式會多次執行這些步驟。儲存作業會在您於「更新工作負載 JobSet 資訊清單」中設定的 CHECKPOINT_DIR 位置執行。

Orbax

建立 Orbax 檢查點。目錄名稱為步驟編號。複製器會偵測新建立的檢查點目錄、視需要執行複製或備份作業，並自動清理。

如要進一步瞭解如何使用 Orbax 複製器檢查點管理工具，請參閱 MaxText 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 bucket PermissionDenied 錯誤所致 (如範例所示)，請正確設定權限，即可解決問題。

使用多層檢查點功能，在 GKE 上訓練大型機器學習模型 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

優點

事前準備

需求條件

限制

設定 GKE 節點，以便使用多層級檢查點

設定新叢集的節點

設定現有叢集的節點

設定使用多層級檢查點的權限

授予 Cloud Storage bucket 存取權

(選用) 授予 Compute Engine 預設服務帳戶存取權

部署 JobSet 控制器

初始化多層級檢查點 CSI 驅動程式

建立 CheckpointConfiguration

解除安裝多層級檢查點 CSI 驅動程式

管理 Cloud Storage 備份的資料保留和垃圾回收

更新工作負載 JobSet 資訊清單

在應用程式中整合多層級檢查點

啟動

Orbax

PyTorch

從最後已知良好 (LKG) 查核點還原

從特定備份還原

儲存檢查點

Orbax

PyTorch

叢集升級

疑難排解

未啟用多層檢查點

多層檢查點 CSI 驅動程式無法掛接磁碟區

後續步驟

使用多層檢查點功能，在 GKE 上訓練大型機器學習模型