Cloud TPU 自动检查点 [公开预览版]

概览

过去，当 TPU 虚拟机需要维护时，该过程会立即启动，而不给用户留出时间来执行可保持进度的操作（例如保存检查点）。如图 1(a) 所示。

图 1.自动检查点功能图示： (a) 如果没有自动检查点，当即将发生维护事件时，从上一个检查点开始的训练进度会丢失。(b) 借助自动检查点，当即将发生维护事件时，可以保留自上次检查点以来的训练进度。

您可以使用自动检查点（图 1(b)）来保留训练进度，方法是将代码配置为在发生维护事件时保存非计划检查点。发生维护事件时，系统会自动保存自上一个检查点以来的进度。该功能同时适用于单个切片和多切片。

自动检查点功能可与可捕获 SIGTERM 并随后保存检查点的框架配合使用。支持的框架包括 MaxText、Pax 和采用 Orbax 的 JavaX。对其他框架的支持将在推出时公布。

目前，只有通过 Cloud TPU API 创建的 TPU（v2-v4 和 v5e）可以使用此功能。我们将在 GKE 中宣布对 TPU 的支持。

使用自动检查点

自动检查点功能默认处于停用状态。创建 TPU 或加入队列的资源时，您可以在预配 TPU 时添加 --autocheckpoint-enabled 标志来启用它。启用此功能后，Cloud TPU 会在收到维护事件的通知后执行以下步骤：

1. 使用 TPU 设备捕获发送到进程的 SIGTERM，
2. 等待进程退出或已经过 5 分钟（以先发生者为准），并对受影响的切片执行维护。

请注意，Autocheckpoint 使用的基础架构独立于机器学习框架。任何机器学习框架都可以支持 Autocheckpoint，前提是它可以捕获 SIGTERM 信号并启动检查点流程。

在应用代码中，您需要启用机器学习框架提供的 Autocheckpoint 功能。例如，在 Pax 中，这意味着在启动训练时启用命令行标志（请参阅 Pax 的自动检查点快速入门）。在后台，框架会在收到 SIGTERM 后保存非计划检查点，并且在 TPU 不再使用时，受影响的 TPU 虚拟机进行维护。

快速入门：使用 MaxText 实现自动检查点

MaxText 是一种“经过良好测试的高性能、可任意伸缩的开源 LLM，以纯 Python/JAX 为目标平台、针对 Cloud TPU 编写的 LLM”。MaxText 包含使用自动检查点功能所需的所有设置。

MaxText 自述文件介绍了大规模运行 MaxText 的两种方法：

• 使用 multihost_runner.py（建议用于实验）
• 使用 multihost_job.job（建议用于生产环境）

使用 multihost_runner.py 时，唯一需要进行的更改是在预配排队的资源时设置 autocheckpoint-enabled 标志。使用 multihost_job.py 时，唯一需要进行的更改是在启动作业时指定 ENABLE_AUTOCHECKPOINT=true 命令行 flag。

快速入门：在单个 Slice 上使用 Pax 实现自动检查点

在本部分中，我们将举例说明如何在单个 Slice 上设置和使用带有 Pax 的 Autocheckpoint。通过适当的设置：

• 发生维护事件时，系统会保存一个检查点。
• 保存检查点后，Cloud TPU 将在受影响的 TPU 虚拟机上执行维护。
• 在 Cloud TPU 完成维护后，您可以照常使用 TPU 虚拟机。

1. 创建 TPU 虚拟机或已加入队列的资源时，请使用 autocheckpoint-enabled 标志。

   例如：

   PROJECT=your-gcp-project-name
   ZONE=zone-you-want-to-use
   NODE_ID=your-node-id
   ACCELERATOR_TYPE=your-accelerator-type
   gcloud config set project $PROJECT
   gcloud config set compute/zone $ZONE
   

   gcloud alpha compute tpus tpu-vm create $NODE_ID \
   --accelerator-type $ACCELERATOR_TYPE \
   --version tpu-ubuntu2204-base \
   --autocheckpoint-enabled
   

   注意： 支持 Autocheckpoint 的 Pax 版本需要 Python 3.10 附带的运行时版本 tpu-ubuntu2204-base。

2. 在单个切片上安装 Pax

   自动检查点功能适用于 Pax 1.1.0 及更高版本。在 TPU 虚拟机上，安装 jax[tpu] 和最新的 paxml：

   pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

3. 使用适当的配置启动训练

   以下示例展示了如何配置 LmCloudSpmd2B 模型以将自动检查点触发的检查点保存到 Google Cloud Storage 存储桶：

   JOB_LOG_DIR=gs://your-storage-bucket
   
   { python3 .local/lib/python3.10/site-packages/paxml/main.py
   --jax_fully_async_checkpoint=1 \
   --exit_after_ondemand_checkpoint=1 \
   --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2B \
   --job_log_dir=$JOB_LOG_DIR; } 2>&1 | tee pax_logs.txt
   

   请注意传递给此命令的两个标志：

   • jax_fully_async_checkpoint：启用此标志后，将使用 orbax.checkpoint.AsyncCheckpointer。 当训练脚本收到 SIGTERM 信号时，AsyncCheckpointer 类会自动保存检查点。
   • exit_after_ondemand_checkpoint：启用此标志后，TPU 进程会在自动检查点保存成功后退出，这会触发立即执行维护。如果您不使用此标志，训练将在检查点保存后继续，Cloud TPU 将在执行必要的维护之前等待超时（5 分钟）。

快速入门：在多切片中使用 Pax 实现自动检查点

自动检查点不仅适用于单个切片，也适用于多切片。本部分详细介绍了结合使用自动检查点与多切片所需的步骤。

1. 在创建排队的资源期间指定自动检查点。

   多切片环境只能通过已加入队列的资源请求进行预配。与单切片的情况类似，在调用中使用 autocheckpoint-enabled 标志来创建加入队列的资源。

   QR_ID=your-qr-id
   NODE_COUNT=your-node-count
   ACCELERATOR_TYPE=your-accelerator-type
   
   gcloud compute tpus queued-resources create $QR_ID \
   --node-count $NODE_COUNT \
   --accelerator-type $ACCELERATOR_TYPE \
   --runtime-version tpu-ubuntu2204-base \
   --autocheckpoint-enabled
   

   如需详细了解所有可用选项，请参阅多切片用户指南。创建排队的资源请求并使其处于 ACTIVE 状态后，请按照后续步骤运行带有 Autocheckpoint 的 Pax。

2. 在多切片环境中的所有虚拟机上安装 Pax。

   在 TPU 虚拟机上，在多切片环境中的所有 TPU 虚拟机上安装 jax[tpu] 和最新的 paxml：

   pip install paxml && pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

3. 使用适当的配置启动训练

   以下示例展示了如何在多切片环境中训练时为 Autocheckpoint 配置模型 LmCloudSpmd2B。在运行训练脚本之前，请将 DCN_MESH_SHAPE 设置为 [2, 1, 1]，如以下代码所示：

   @experiment_registry.register
   class LmCloudSpmd2B(LmCloudSpmd):
   
   """SPMD model with 2B params.
   
   Global batch size = 2 * 2 * 1 * 32 = 128
   """
   PERCORE_BATCH_SIZE = 8
   
   NUM_LAYERS = 18
   MODEL_DIMS = 3072
   HIDDEN_DIMS = MODEL_DIMS * 4
   
   CHECKPOINT_POLICY = layers.AutodiffCheckpointType.SAVE_NOTHING
   ICI_MESH_SHAPE = [1, 4, 1]
   DCN_MESH_SHAPE = [2, 1, 1]
   

   启动训练时，除了单切片案例中讨论的命令行标志之外，还需要三个标志：

   • num_hosts：主机的总数。在本示例中，它是 2。
   • host_index：启动训练的主机的索引。范围在 0 到 N-1 之间，其中 N 是主机的总数。
   • server_addr：节点 0 的工作器 0 的 IP 地址，具有未使用的端口（例如 8476）。如需查找此信息，请在节点 0 的工作器 0 上使用 hostname -i。

使用 Orbax 的自动检查点

自动检查点功能不限于 MaxText 或 Pax。任何能够捕获 SIGTERM 信号并启动检查点流程的框架都适用于 Autocheckpoint 提供的基础架构。Orbax 是一个为 JAX 用户提供通用实用程序库的命名空间，可提供这些功能。

如 Orbax 文档中所述，默认情况下，系统会为 orbax.checkpoint.CheckpointManager 的用户启用这些功能。在每个步骤后调用的 save 方法会自动检查维护事件是否即将发生，如果是，则保存检查点，即使步骤序号不是 save_interval_steps 的倍数。GitHub 文档还说明了如何在保存自动检查点后通过修改用户代码退出训练。