Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

使用多层级检查点在 GKE 上训练大规模机器学习模型

Standard

本页面介绍了如何在 GKE 上训练机器学习模型期间使用多层级检查点设置来可靠地存储和管理检查点。检查点存储和管理对于大规模训练作业至关重要，大规模训练作业是指利用数千个节点的作业。这些大规模作业经常中断（可能每小时中断一次），并且恢复速度可能很慢。

优势

使用多层级检查点设置具有以下优势：

完全编排的检查点数据管理，包括以下工作负载的备份、复制和自动恢复：
- 使用 Orbax 进行状态管理并在 TPU 上运行的 JAX 训练运行。
- 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 存储桶，请创建一个。请确保启用分层命名空间，否则备份将失败。

要求

多层级检查点设置需要 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 机器。您不会将此机器类型用于模型训练；而是会为此目的创建单独的节点池，通常会利用加速器优化的虚拟机系列。
- 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 存储桶中，多层级检查点需要访问该存储桶。向 Kubernetes 服务账号授予存储桶的 Storage Object User (roles/storage.objectUser) IAM 角色，以实现多层级检查点。

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 默认服务账号授予 Storage Object Viewer (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 是一种 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 存储桶的名称。使用您在设置具有权限的 Cloud Storage 存储桶步骤中设置的存储桶。
- 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 驱动程序安排到具有匹配污点的节点上。如需详细了解污点在不同加速器类型上的工作方式，请参阅以下页面：
  - 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 清单，以包含大规模检查点卷。详细信息取决于您的工作负载。

例如，如需扩展示例 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 检查点管理器 （如需查看示例，请参阅 MaxText）。
PyTorch 参考示例：此示例适用于将 NeMo 与 PyTorch 搭配使用的工作负载。

启动

本部分介绍了应用与多层级检查点设置交互所需的初始步骤。

复制器是多层级检查点设置的核心组件，作为 CSI 驱动程序的一部分在每个节点上运行。复制器管理跨存储层级的检查点复制，从本地 RAM 磁盘到对等节点，再到 Cloud Storage 等外部存储。

replicator.yaml 文件充当机器学习训练作业（框架代码）和复制器组件之间的动态控制平面。机器学习应用以编程方式在本地卷 (RAMDisk) 上生成此文件，训练作业和复制器服务都可以访问该文件。此清单允许机器学习框架向复制器提供运行时配置和生命周期管理说明，这与后端设置期间定义的静态基础架构参数（例如 Cloud Storage 上传频率）不同。

如需查看此交互的具体示例，请参阅：

MaxText 项目，该项目将此架构用于 Cloud TPU 上的 JAX。
PyTorch 参考示例，该示例将此架构与 GPU 上的 PyTorch 和 NVIDIA NeMO 搭配使用。

您的应用应在启动期间执行以下步骤：

等待直到 replicator.yaml 文件不存在，这表示复制器已准备好由您的应用进行配置。replicator.yaml 文件是在您在CHECKPOINT_DIR位置配置的更新工作负载 JobSet 清单部分生成的。

首次创建模型训练作业时，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：分布式训练作业中当前节点的唯一标识符。这表示创建此文件的节点的节点排名。参与运行的每个节点都有自己的排名。
- nodes：参与分布式训练作业的节点总数。此值来自 Pod 的元数据。机器学习训练作业也可以查看此值。
- peer-ranks 或 peers-per-node：指定复制拓扑的两种替代方法。这两个参数中只能存在一个。
  - peer-ranks：当前节点的检查点数据应复制到的对等节点的显式排名。这样可以精确控制哪些特定节点充当复制合作伙伴。
  - peers-per-node：复制器应自动选择用于复制的每个节点的对等节点数。
- 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 复制检查点管理器示例。

默认情况下，LKG 检查点是多层级检查点设置可以找到的最新完整检查点。首选集群内检查点（在本地 RAM 磁盘上和对等节点上），并且仅当不存在完整的集群内检查点时才使用 Cloud Storage 备份。如需替换此选择，请参阅从特定备份恢复。

从特定备份恢复

您可以从特定的 Cloud Storage 备份（而不是 LKG 检查点）恢复作业，例如，将训练回滚到较早的状态。在作业开始之前，请执行以下步骤：

找到要从中恢复的备份。多层级检查点设置会将备份写入您在 CheckpointConfiguration 中配置的存储桶，并写入以作业命名的目录中。每个备份都是一个以备份拍摄时的 UTC 时间命名的子目录，格式为 YYYY-MM-DD_HH-MM：
```
gcloud storage ls gs://GCS_PATH/JOB_NAME/
```
替换以下内容：
- GCS_PATH：多层级检查点设置写入备份的 Cloud Storage 路径。此值是您在 CheckpointConfiguration 的 cloudStorageBucketName 字段中指定的存储桶，包括存储桶中的任何目录。
- JOB_NAME：训练作业的名称，如 job-name 文件的 replicator.yaml 字段中所指定。
在存储桶中的作业目录中，创建一个名为 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 检查点恢复。
如果指定的备份不存在或不完整，作业将启动，但没有检查点。多层级检查点设置不会回退到集群内检查点或其他备份。

保存检查点

在训练作业进行期间，您的应用会多次执行这些步骤。保存操作发生在您在CHECKPOINT_DIR位置您在更新工作负载 JobSet 清单中配置的。

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 对象。此操作是必需的，因为此对象动态部署的节点检查点设置驱动程序守护进程集不会自动升级。

如果您特别希望守护进程集规范保持不变，则无需执行任何操作。

问题排查

本部分提供了解决多层级检查点设置问题的问题排查指南。如需了解常规存储问题排查，请参阅排查 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 上训练大规模机器学习模型 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

优势

准备工作

要求

限制

配置 GKE 节点以使用多层级检查点设置

在新集群上配置节点

在现有集群上配置节点

配置使用多层级检查点设置的权限

授予对 Cloud Storage 存储桶的访问权限

（可选）授予 Compute Engine 默认服务帐号访问权限

部署 JobSet 控制器

初始化多层级检查点设置 CSI 驱动程序

创建 CheckpointConfiguration

卸载多层级检查点设置 CSI 驱动程序

管理 Cloud Storage 备份的数据保留和垃圾回收

更新工作负载 JobSet 清单

在应用中集成多层级检查点设置

启动

Orbax

PyTorch

从上次已知良好 (LKG) 检查点恢复

从特定备份恢复

保存检查点

Orbax

PyTorch

集群升级

问题排查

未启用多层级检查点设置

多层级检查点设置 CSI 驱动程序无法装载卷

后续步骤

使用多层级检查点在 GKE 上训练大规模机器学习模型