通过 Pathways 进行弹性训练

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

  • 暂停-恢复：在面对计划内中断（例如 抢占通知）时具有容错能力，而无需用户编写任何自定义抢占 处理代码。
  • 弹性训练：在面对计划外硬件故障时具有容错能力 不会导致客户端崩溃，但需要用户编写特定于模型 的恢复代码。

准备工作

请确保您已备妥：

暂停-恢复

通常，GKE 会在抢占加速器 Pod 之前向其发送抢占通知。在所有云部署中，Pathways 抢占容错功能默认处于启用状态，Pathways 加速器作业会监听这些通知。

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

如果工作负载不可恢复，Pathways 会关闭加速器 作业，并在配置了弹性训练 的情况下将故障转发给作业。如果未配置弹性训练，GKE 会根据 JobSet 重启政策重启整个工作负载。

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

弹性训练

借助弹性训练，即使发生硬件故障，训练作业也能继续进行。这是通过 Pathways 系统功能和用户定义的模型恢复逻辑相结合来实现的：

  • 故障检测：当发生硬件故障（例如 TPU 工作器崩溃）时，Pathways 系统会检测到此故障，并在下次访问该 硬件上的数据时通过异常通知用户的 训练作业。此通知不会导致工作负载崩溃；它允许您的代码处理通知并重新配置资源，以继续处理或正常退出。
  • 用户定义的弹性处理程序：用户的模型代码需要能够 处理此异常。这就是“特定于模型的恢复”的含义。
    • 快照：最常见的方法是定期保存模型状态的快照 。发生故障时，您可以从最新的快照加载以恢复训练。
    • 重新配置：您可能需要重新配置训练作业，以 适应可用切片的数量。例如，如果一个切片停止工作，您可以将活跃切片的数量减少一个，直到有替换切片可用为止。如需了解详情，请参阅弹性处理程序。
    • 数据/计算图更新：您的代码需要通过根据需要重新创建 计算图来处理可用于计算的设备数量的任何变化。这可能涉及重新分区数据或重新编译模型。
  • Pathways 在恢复中的作用：Pathways 提供支持 用户定义重新配置的原语：
    • 切片替换：如果替换了发生故障的切片，则可以在新切片可用后 通知客户端。然后，您的代码可以重新配置以使用此新切片。
    • 透明恢复：Pathways 处理 恢复的较低级别详细信息，例如重新建立与 集群健康部分的连接。
  • pathwaysutils 中的实用程序：在 pathways-utils中定义的一组 Pathways 实用程序。

实现弹性处理程序

您必须编写的大部分代码都将位于用户定义的弹性处理程序中。此处理程序通过重新创建网格并重新初始化训练循环来响应弹性事件（例如 TPU 切片变得不可用）。

每个工作负载都是独一无二的。弹性处理程序的复杂性可能会随着工作负载的复杂性而变化。处理程序的输入和输出应是重新初始化训练循环所需的最少参数和返回值。

def elastic_handler(elastic_utils, *args, **kwargs):
  mesh = initialize_mesh(**kwargs["mesh_kwargs"])
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **kwargs["initialize_training_loop_kwargs"])

  step, snapshot = elastic_utils.get_next_snapshot()
  state = initial_state.replace(**snapshot)

  return state, step, mesh, jitted_train_step, other_variables

更新训练循环

您需要对训练循环进行以下更改：

  1. 创建弹性管理器
  2. 将训练循环封装在处理 jax.errors.JaxRuntimeError 的 try-except 块内
  3. 在 jax.errors.JaxRuntimeError 处理程序中，调用 maybe_reshard_down。 如果错误与弹性事件相关，弹性管理器将重新分片，否则将重新引发该错误。
  4. 在训练循环结束时调用 maybe_snapshot 和 maybe_reshard_up
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()

def initialize_mesh(**kwargs):
  ...


def initialize_training_loop(**kwargs):
  ...


def train_loop(
    final_step,
    elastic_manager,
    mesh_kwargs,
    initialize_training_loop_kwargs,
):
  mesh = initialize_mesh(**mesh_kwargs)
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **initialize_training_loop_kwargs)

  step = initial_step
  while step < final_step:
    try:
      state = jitted_train_step(state)

      elastic_manager.maybe_snapshot(step=step, snapshot=state)
      handler_returns = elastic_manager.maybe_reshard_up(
          step=step,
          snapshot=state,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns
      step += 1
    except jax.errors.JaxRuntimeError as error:
      handler_returns = elastic_manager.maybe_reshard_down(
          error=error,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns

  return state


def main():
  elastic_manager = manager.Manager(
      devices=jax.devices(),
      snapshot_period=10,
      snapshot_buffer_size=1,
      reshard_check_period=5,
      max_elastic_down_event_count=10,
      max_reshard_retry_count=3,
  )

  train_loop(100, elastic_manager, {}, {})

配置弹性管理器

弹性管理器可以通过几种不同的方式进行配置。快照的频率由快照周期决定。快照周期会影响因弹性事件而丢失的平均步数。重新分片检查周期决定了训练循环轮询切片可用性的频率。 max_elastic_down_event_count 可让您设置训练循环将支持的因切片丢失而发生的弹性事件数。max_reshard_retry_count 指定了弹性管理器应重试重新分片的次数。管理器是单例对象，应仅创建一次。

快照

根据弹性管理器配置，该函数可能会将数据快照到主机内存中，弹性处理程序在弹性事件期间可以使用这些数据。

减少分片

捕获 jax.errors.JaxRuntimeError 后，Pathways 将检查错误是否是由于切片丢失而发生的弹性事件所致。如果是，它将循环调用弹性处理程序，直到成功或达到最大重试次数。如果错误不是由弹性事件引起的，则会再次引发该错误。弹性处理程序的返回值将传递给调用方。

增加分片

根据弹性管理器配置以及是否存在不可用的切片，Pathways 将检查是否有其他切片可用。如果是，它将立即保存快照（如果尚未针对当前步拍摄现有快照），并循环调用弹性处理程序，直到成功或达到最大重试次数。如果发生重新分片，弹性处理程序的返回值将传递给调用方。否则，将返回 None。

热插拔

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

创建 JobSet 后，GKE 会根据 JobSet 配置中的规定，在多个切片上调度工作负载。如果一个或多个切片发生硬件故障，受影响的 Pod 将被标记为失败。在重新调度此 JobSet 时，如果您选择在 GKE 集群中保留一个备用切片，该切片可用于优先级较低的作业，则 JobSet 系统会将优先级较高的作业的失败切片的工作负载重新映射到同一 GKE 集群中优先级较低的作业正在使用的备用切片上。此重新映射通常需要不到一分钟的时间。

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

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

如需在集群中使用优先级，请定义优先级类，例如：

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

后续步骤