通过 Pathways 进行弹性训练

Pathways 通过以下方式提供弹性优势：

• 暂停-恢复：让您的工作负载能够容忍计划内中断（例如抢占通知），而无需自定义抢占处理代码。
• 弹性训练：通过在训练代码中使用恢复逻辑，让您的工作负载能够容忍意外硬件故障，而不会导致客户端崩溃。

准备工作

在开始之前，请完成以下任务：

• 设置 Cluster Toolkit
• 安装 Kubernetes 工具
• 安装 Google Cloud CLI
• 启用 Google Kubernetes Engine API

暂停-恢复

Google Kubernetes Engine (GKE) 通常会在抢占加速器 Pod 之前向其发送抢占通知。默认情况下，所有 Cloud Deploy 都会启用 Pathways 抢占容忍度，并且 Pathways 加速器作业会监听这些通知。

收到抢占通知后，Pathways 会首先确定是否可以透明地保存和恢复工作负载。如果是，则在 GKE 逐出加速器作业之前，它会尝试通过将当前状态写入永久性存储（例如 Cloud Storage）来透明地暂停机器学习工作负载。当 GKE 稍后重新安排作业时，Pathways 会通过读回其持久化状态来恢复机器学习工作负载。

如果工作负载无法恢复，Pathways 会关闭加速器作业，并将故障转发到您的作业（如果配置了弹性训练）。如果未配置弹性训练，GKE 会根据 JobSet 重启政策重启整个工作负载。

使用 JAX 定义的典型机器学习工作负载依赖于无状态 Pathways XLA 组件，这些组件可以使用高带宽内存 (HBM) 快照进行恢复。某些机器学习工作负载（例如使用 JAX 共置 Python API 定义的工作负载）依赖于有状态 Pathways 组件；这些工作负载无法恢复。

弹性训练

弹性训练可让您的训练作业在发生硬件故障或切片抢占时继续运行，而无需重启整个工作负载。在多切片训练作业中，当单个切片不可用时，Pathways 会使单控制器客户端进程和健康的 TPU 切片保持活跃状态。

弹性训练将 Pathways 系统功能与训练代码中的恢复逻辑相结合：

• 故障检测：当 TPU 切片发生硬件故障或抢占时，Pathways 会取消分配受影响的切片，并在您的工作负载下次访问该切片上的数据时引发 jax.errors.JaxRuntimeError。由于客户端进程保持活跃状态，因此您的代码可以捕获此错误、重新配置资源并继续训练。
• 恢复模式：pathwaysutils.elastic.manager.Manager 类及其 elastic_retry 修饰器支持两种恢复策略：
  • 暂停和恢复：当 minimum_slice_count 为 None 或等于作业中的切片总数时，训练会在切片丢失时暂停，并等待所有切片再次恢复正常后恢复。由于切片数量在多次尝试中不会发生变化，因此您的工作负载可以重复使用其现有设备网格和已编译的可执行文件。
  • 副本大小调整：当 minimum_slice_count 小于切片总数时，系统会通过重建设备网格并针对较小的拓扑重新编译，立即在剩余的活跃切片上继续训练。后台监控线程会定期检查非活跃切片，并在 elastic_manager.available_inactive_slices 中记录新恢复正常的切片。然后，训练循环可以使用替换切片重建设备网格，针对扩展的拓扑重新初始化，并将当前状态放置到新网格上。
• 状态恢复：如需在弹性事件后恢复训练，您的工作负载会使用内存中快照恢复模型参数、优化器状态和数据迭代器状态。orbax.checkpoint 中的 Snapshotter 类会将复制的状态从 TPU HBM 异步复制到每个切片上的 pinned_host CPU 内存。在恢复期间，Snapshotter 会根据幸存的主机副本重构状态，并将其直接重新分片到活跃的 TPU 网格上，而无需从远程存储空间读取。

在 Pathways 代理服务器上启用弹性

如需在 GKE 中的 Pathways 上启用弹性设备集，请将 --num_elastic_instances 标志添加到 JobSet 清单中 pathways-proxy 容器的 args 字段，并指定工作负载中的弹性 TPU 切片数量：

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

替换以下内容：

• BUCKET_NAME：用于 Pathways 临时文件的 Cloud Storage 存储桶。
• NUM_SLICES：工作负载中的 TPU 切片数量。

过滤活跃设备并构建设备网格

在 Pathways 工作负载中，jax.devices() 会在客户端进程的生命周期内返回初始分配中的所有设备。使用弹性训练时，请使用 elastic_manager.active_devices，以便设备网格和批次大小仅使用正常运行的活跃切片：

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

在副本大小调整模式下，当活跃切片数量发生变化时，请使用弹性数据加载器，例如 grain.experimental.ElasticIterator。标准单主机数据迭代器会保存每个主机的状态，并且无法在不同数量的主机上恢复，而 ElasticIterator 会跟踪单个全局数据集位置，该位置可在任意数量的活跃切片中恢复。

弹性训练循环

弹性训练循环可协调切片监控、故障恢复和内存中快照，如以下代码示例所示：

• train 函数会在重试封装容器之外创建一个 manager.Manager() 实例和一个 Orbax Snapshotter，以便在多次尝试之间保留切片跟踪和主机内存快照。
• 使用 @elastic_manager.elastic_retry 修饰内部 run_training_loop 函数会等待健康切片，将 jax.default_device 设置为活跃切片，在各个步骤中保持后台监控器线程处于活跃状态，并在发生 jax.errors.JaxRuntimeError 时重试循环。
• 从尝试开始时的切片故障中恢复时，您的代码必须从 elastic_manager.active_devices 重新构建设备网格，并调用 snapshotter.load 将最新快照恢复到更新后的网格分片上，同时让失败尝试中的设备缓冲区超出范围。
• 在循环的每个步骤中，您的代码会定期调用 snapshotter.save 来备份训练状态，并检查 elastic_manager.available_inactive_slices 以检测运行状况良好的替换切片。当替换切片可用时，您的代码会将其添加到 elastic_manager.active_slice_indices，重新构建网格，重新初始化训练循环，并将状态放置在扩展网格上。

以下示例展示了如何通过结合使用 manager.Manager、elastic_retry 和 Orbax 的内存中 Snapshotter 来实现弹性训练循环：

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

配置弹性管理器和重试行为

您需要为每个进程实例化一次 manager.Manager(devices=None)。如果您省略 devices，管理器会使用 jax.devices() 并按设备的 slice_index 属性对设备进行分组。

Manager 实例公开了以下属性，以便在训练期间检查集群拓扑：

• slice_to_devices：从每个 slice_index 到属于该切片的 jax.Device 对象序列的映射。
• all_slice_indices：分配中所有切片索引的 frozenset。
• active_slice_indices：参与活跃计算网格的切片索引的 frozenset。
• active_devices：属于活跃切片的 jax.Device 对象列表，按 slice_index 排序。
• inactive_slice_indices：分配中未参与计算网格的切片索引的 frozenset。
• available_inactive_slices：后台监控线程检测到健康且准备好重新加入的非活跃切片索引的 frozenset。
• total_slice_count 和 active_slice_count：分配的切片总数和活跃切片数。
• default_device：来自活跃切片的设备 0，elastic_retry 会自动将其设置为每次尝试的默认 JAX 设备。

您可以通过将以下参数传递给 @elastic_manager.elastic_retry 来自定义重试和伸缩行为：

• minimum_slice_count：在开始初始执行或重试尝试之前，所需的健康切片数量下限。如果设置为 None（默认值），elastic_retry 会等待分配中的所有切片 (total_slice_count)，并在暂停和恢复模式下运行。如果设置为较小的整数（例如 1），elastic_retry 会在副本大小调整模式下运行。
• poll_interval：弹性训练管理器 (Manager) 在训练期间执行切片健康检查的时间间隔（以秒为单位）。默认值为 10。
• timeout：在引发 TimeoutError 之前，等待 minimum_slice_count 切片变为活跃状态的最长秒数。默认值为 None（无超时）。
• pre_callback：一个可选回调，在每次尝试开始时，在 active_slice_indices 更新后且在封装的训练循环函数运行之前运行。
• on_elastic_event_callback：一个可选回调，在捕获到切片 jax.errors.JaxRuntimeError 后、下一次重试开始前运行。您可以使用此回调来清理不完整的检查点或重置数据迭代器。
• retry_policy：一个可调用的 (attempt: int, error: Exception) -> bool，用于确定是否在弹性事件后继续重试。使用 manager.ElasticRetryLimit(max_attempts) 限制总尝试次数。如果省略，elastic_retry 将无限期重试。

在弹性事件中，一旦堆栈帧展开且没有变量引用这些数组，JAX 就会释放失败尝试中的设备缓冲区。避免在全局变量或长期回调中保留对失败尝试中的设备数组的引用，以便在下次尝试开始之前释放内存。

将内存中快照与 Orbax 结合使用

ocp.v1.training.pathways.Snapshotter 类（在 orbax-checkpoint 0.12.6 版及更高版本中提供）在切片之间的 pinned_host 内存中维护训练状态的内存中备份：

• Snapshotter(replica_axis_index=0)：初始化快照工具，并指定哪个网格轴表示数据并行切片副本。保存状态中的所有 JAX 数组都必须沿此网格轴进行复制，以便在切片发生故障时，幸存的切片仍保留模型和优化器状态的完整副本。如果您的状态包含单设备数组，请在调用 save 之前在网格中复制这些数组。
• save(step, state)：将 state 异步传输到 pinned_host 内存，并沿 replica_axis_index 将其拆分为每个副本树。如果之前的快照传输仍在进行中，save 会跳过新请求，以避免训练循环停滞。
• load(abstract_state, reset_snapshot_state=True)：探测已保存的主机副本以识别幸存切片，串联正常副本，并将状态重新分片到 abstract_state 指定的目标分片。
• join()：在任何已加入队列的快照操作完成之前，阻止其他操作。
• latest：返回最近完成的快照的元数据（包括 step），如果没有任何快照完成，则返回 None。

减少切片失败时的分片

当 TPU 切片在执行期间失败或被抢占时，Pathways 会引发 jax.errors.JaxRuntimeError（例如 DATA_LOSS、DEADLINE_EXCEEDED、NOT_FOUND 或 INTERNAL）。elastic_retry 修饰器会捕获此异常并检查 elastic.is_error_due_to_slice_down(error)：

• 如果错误不是由切片事件引起的，elastic_retry 会立即重新引发异常。
• 如果错误是由切片缩减事件引起的，elastic_retry 会停止所有活跃的 JAX 性能分析器跟踪记录，运行 on_elastic_event_callback，并评估 retry_policy(attempt, error)。
• 在开始下一次尝试之前，elastic_retry 会调用 elastic.wait_for_slices 一次，该函数会轮询切片运行状况，直到至少有 minimum_slice_count 个切片处于活跃状态，然后更新 elastic_manager.active_slice_indices，并在幸存的切片上运行训练循环函数。

在切片重新加入时增加分片

当工作负载在副本大小调整模式下以较少的切片数量运行时，后台监控线程会定期对 elastic_manager.inactive_slice_indices 运行健康检查。当修复或重新安排的切片通过健康检查时，监控线程会将其切片索引添加到 elastic_manager.available_inactive_slices。

如果 elastic_manager.available_inactive_slices 在训练步骤结束时非空，则工作负载会就地扩容，而不会丢失训练进度：

• 将 elastic_manager.available_inactive_slices 中的切片索引添加到 elastic_manager.active_slice_indices 会将重新联接的切片纳入 elastic_manager.active_devices，并将其从 available_inactive_slices 中清除。
• 通过 elastic_manager.active_devices 重新构建设备网格，并重新初始化训练步骤和数据迭代器，使计算适应扩展后的拓扑。
• 调用 jax.device_put 会将当前训练状态放置到新的网格分片上，而无需从快照恢复。

热插拔

热插拔是指 GKE JobSet API 的一项功能，其中优先级较高的作业可以快速接管优先级较低的作业的资源，从而最大限度地减少停机时间并确保更快恢复。

创建 JobSet 后，GKE 会按照 JobSet 配置中的指定，在多个切片之间调度工作负载。如果一个或多个切片发生硬件故障，受影响的 Pod 会被标记为失败。如果您的 GKE 集群有一个运行较低优先级作业的备用切片，JobSet 系统会将故障切片的工作负载从较高优先级作业重新映射到该备用切片。此重新映射过程通常用时不超过 1 分钟。

JobSet 重启时，在以下情况下可能会发生热插拔：

• 默认模式：如果同一集群中有空闲的备用 TPU 切片，Kubernetes 调度器会优先将重启的作业调度到这些切片上，而不是等待故障切片修复。这种调度行为可加快恢复速度。
• 异构工作负载：在运行多个工作负载且配置了 Kubernetes PriorityClass 资源的集群中，重启的 JobSet 可能会触发热插拔。如果重启的作业的亲和性与优先级较低的作业的资源匹配，则 Kubernetes 会抢占优先级较低的作业，以便优先级较高的作业可以启动，而无需等待切片修复。例如，您可以使用 PriorityClass 资源为 Pathways 工作器 Pod 配置不同的优先级。

如需在集群中使用优先级，请定义 PriorityClass 清单：

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

将此 YAML 应用到您的 GKE 集群：

kubectl apply -f high-prior-job.yaml

接下来，通过将以下文本添加到 pathways-worker Pod 的 podspec，将新的优先级类附加到 Pathways 工作器作业。

priorityClassName: high-prior-job

后续步骤

• 创建具有 Pathways 的 GKE 集群
• 使用 Pathways 运行批量工作负载
• Pathways 互动模式
• 将 JAX 工作负载移植到 Pathways
• 排查 Pathways on Cloud 问题