本頁面由 Cloud Translation API 翻譯而成。

在 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 安裝項目，請務必設定 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 驅動程式和 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 驅動程式和 GKE 適用的工作負載身分聯盟，方法是執行下列指令。現有叢集版本必須高於 1.32.3-gke.1170000。

為 GKE 啟用 Workload Identity Federation：
```
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 值區的存取權

Multi-Tier Checkpointing CSI 驅動程式使用的臨時磁碟區必須使用現有的 Cloud Storage bucket。

如要將檢查點儲存在 Cloud Storage 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 驅動程式

本節說明如何在工作負載執行的節點上，初始化 Multi-Tier Checkpointing 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-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 bucket 名稱。使用您在「設定具備權限的 Cloud Storage 值區」步驟中設定的值區。
- MACHINE_TYPE：對應加速器的機器類型。可能的值如下：
  - TPU v5p：ct5p-hightpu-4t
  - TPU v5e：ct5e-hightpu-4t
  - TPU v6e：ct6e-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 驅動程式

如要取消部署 Multi-Tier Checkpointing CSI 驅動程式，請刪除 CheckpointConfiguration 資源。Multi-Tier Checkpointing 控制器會從節點中移除 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 的位置，且多層檢查點會執行檢查點儲存作業。詳情請參閱「在應用程式中整合多層檢查點」。

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

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

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

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

啟動

本節說明應用程式與多層檢查點互動時需要執行的初始步驟。

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

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

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

MaxText 專案，這個專案會使用此架構在 Cloud TPU 上執行 JAX。
PyTorch 參考範例，在 GPU 上使用 PyTorch 和 NVIDIA NeMO，採用這個架構。

應用程式啟動時應執行下列步驟：

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

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

初始化複製器後，多層檢查點會為每個 TPU 或 GPU 工作站建立一個符號連結。這些符號連結會建立在與 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

如果未啟用多層檢查點，請更新叢集以啟用多層檢查點。

Multi-Tier Checkpointing 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 上使用多層檢查點訓練大型機器學習模型 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

優點

事前準備

需求條件

限制

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

在新叢集上設定節點

設定現有叢集中的節點

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

授予 Cloud Storage 值區的存取權

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

部署 JobSet 控制器

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

建立 CheckpointConfiguration

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

管理 Cloud Storage 備份的資料保留時間和垃圾收集

更新工作負載 JobSet 資訊清單

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

啟動

Orbax

PyTorch

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

儲存檢查點

Orbax

PyTorch

叢集升級

疑難排解

未啟用多層檢查點

Multi-Tier Checkpointing CSI 驅動程式無法掛接磁碟區

後續步驟

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