容错分布式 JAX

容错分布式 JAX#

回想一下，多控制器 JAX 允许您在多台机器上运行分布式 JAX 程序。默认情况下，如果其中任何一台机器发生故障，所有机器都会失败。也就是说，多控制器 JAX 默认不是容错的。

本文分为三个部分。在第一部分中，我们将解释如何编写容错多控制器 JAX 程序的基础知识。在第二部分中，我们将展示一些容错多控制器 JAX 程序的示例。在第三部分中，我们将深入了解多控制器 JAX 是如何实现容错的。

警告

JAX 对容错的支持尚处于实验阶段。目前它仅在 GPU 上能完全工作。它还存在一些粗糙之处，可能包含 Bug，并且可能会发生变化。使用风险自负。

注意

如果您正在寻找 TPU 上多控制器 JAX 容错的替代方案，请考虑 Pathways。

第 1 部分：容错基础#

默认不具备容错性#

默认情况下，多控制器 JAX 程序不具备容错能力。如果任何进程崩溃，则所有其他进程也将故意崩溃。为了具体说明这一点，请考虑以下微不足道的脚本 example.py，它通过调用 jax.distributed.initialize 初始化多控制器 JAX，然后进入无限循环。

example.py#
 1from absl import app
 2from absl import flags
 3from collections.abc import Sequence
 4import jax
 5import time
 6
 7_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
 8_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
 9
10
11def main(_: Sequence[str]) -> None:
12  jax.distributed.initialize(
13      coordinator_address="localhost:9000",
14      num_processes=_NUM_PROCESSES.value,
15      process_id=_PROCESS_ID.value,
16      local_device_ids=[_PROCESS_ID.value],
17      heartbeat_timeout_seconds=10,
18  )
19  print(f'{jax.devices()=}')
20  print(f'{jax.local_devices()=}')
21  while True:
22    print(time.time())
23    time.sleep(1)
24
25
26if __name__ == "__main__":
27  app.run(main)

通过在配有四个 GPU 的虚拟机上运行以下四个命令（每个命令在不同的终端中），在四个进程上运行 example.py。jax.distributed.initialize 的 local_device_ids 参数确保每个进程仅被分配四个 GPU 中的一个。我们稍后会解释 heartbeat_timeout_seconds 参数。

python example.py --i=0 --n=4  # in terminal 1
python example.py --i=1 --n=4  # in terminal 2
python example.py --i=2 --n=4  # in terminal 3
python example.py --i=3 --n=4  # in terminal 4

当您运行这些命令时，您会看到进程每秒都会忠实地打印当前时间。接下来，终止第四个进程：pkill -9 -f 'python example.py --i=3 --n=4'。大约十秒后，其他进程也会终止并抛出如下错误消息：

E0926 17:26:32.075402  157988 coordination_service_agent.cc:332] Polled an error from coordination service (this can be an error from this or another task).
F0926 17:26:32.075587  157988 client.h:77] Terminating process because the JAX distributed service detected fatal errors. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. absl::Status: UNAVAILABLE: The following tasks are unhealthy (stopped sending heartbeats):
/job:jax_worker/replica:0/task:3
The tasks have crashed. Check the task logs for an earlier error, or scheduler events (e.g. preemption, eviction) to debug further.

RPC: /tensorflow.CoordinationService/PollForError [type.googleapis.com/tensorflow.CoordinationServiceError='']

当多控制器 JAX 程序中的某个进程发现对等进程崩溃时，它决定也随之崩溃。这些进程是命运共享 (fate-sharing) 的。jax.distributed.initialize 的 heartbeat_timeout_seconds 参数决定了进程在断定对等进程死亡之前等待的时间。前三个进程在您杀死第四个进程约十秒后崩溃，因为我们将 heartbeat_timeout_seconds=10 作为参数传递给了 jax.distributed.initialize。

容错存活#

我们可以通过添加 --xla_gpu_nccl_terminate_on_error=false 标志以及在 example.py 中设置 jax_enable_recoverability 配置选项来禁用命运共享，如下所示：

 1import os
 2os.environ['XLA_FLAGS'] = '--xla_gpu_nccl_terminate_on_error=false'
 3
 4from absl import app
 5from absl import flags
 6from collections.abc import Sequence
 7import jax
 8import time
 9
10_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
11_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
12
13
14def main(_: Sequence[str]) -> None:
15  jax.config.update("jax_enable_recoverability", True)
16  jax.distributed.initialize(
17      coordinator_address="localhost:9000",
18      num_processes=_NUM_PROCESSES.value,
19      process_id=_PROCESS_ID.value,
20      local_device_ids=[_PROCESS_ID.value],
21      heartbeat_timeout_seconds=10,
22  )
23  print(f'{jax.devices()=}')
24  print(f'{jax.local_devices()=}')
25  while True:
26    print(time.time())
27    time.sleep(1)
28
29
30if __name__ == "__main__":
31  app.run(main)

再次在四个进程上运行脚本，然后杀死第四个进程。请注意，现在其他三个进程可以愉快地继续执行。

接下来尝试终止进程 0。请注意，即使禁用了命运共享，所有四个进程仍会终止并出现如下错误消息：

E0929 17:42:48.594192 1044529 coordination_service_agent.cc:332] Polled an error from coordination service (this can be an error from this or another task).
F0929 17:42:48.594200 1044529 client.h:77] Terminating process because the JAX distributed service detected fatal errors. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. absl::Status: UNAVAILABLE: Failed to send RPC to coordination service. Either the leader task was preempted/died/restarted unexpectedly or this task is experiencing network issues. Check earlier logs from 1) this task, 2) the leader (usually slice 0 task 0), and 3) cluster scheduler to debug further.
Additional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/PollForError:
:UNKNOWN:Error received from peer  {grpc_message:"Socket closed", grpc_status:14}

进程 0 很特殊。如果进程 0 失败，即使禁用了命运共享，每个进程也会失败。为什么？进程 0 运行一个名为“协调服务”的 RPC 服务，所有进程都使用该服务进行相互协调。如果协调服务失败，所有其他进程别无选择，只能失败。有关更多详细信息，请参阅第 3 部分：实现细节。

集体通信中的卡顿#

example.py 现在能够存活于故障中，但进程之间根本不进行通信。任何实际的多控制器 JAX 程序都涉及进程间的通信（否则，使用多控制器 JAX 还有什么意义呢？）。让我们编辑 example.py，使进程在循环的每次迭代中执行集体 jnp.sum。

 1import os
 2os.environ['XLA_FLAGS'] = '--xla_gpu_nccl_terminate_on_error=false'
 3
 4from absl import app
 5from absl import flags
 6from collections.abc import Sequence
 7import jax
 8import jax.numpy as jnp
 9import time
10
11_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
12_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
13
14
15def main(_: Sequence[str]) -> None:
16  jax.config.update("jax_enable_recoverability", True)
17  jax.distributed.initialize(
18      coordinator_address="localhost:9000",
19      num_processes=_NUM_PROCESSES.value,
20      process_id=_PROCESS_ID.value,
21      local_device_ids=[_PROCESS_ID.value],
22      heartbeat_timeout_seconds=10,
23  )
24  print(f'{jax.devices()=}')
25  print(f'{jax.local_devices()=}')
26
27  n = jax.device_count()
28  jax.set_mesh(jax.make_mesh((n,), ("i",)))
29  x = jax.device_put(jnp.arange(n), jax.P("i"))
30  while True:
31    print(jnp.sum(x))
32    time.sleep(1)
33
34
35if __name__ == "__main__":
36  app.run(main)

在上面突出显示的代码中，进程创建了一个跨四个进程分片的数组 x，然后执行分布式 jnp.sum。再次运行程序并杀死第四个进程。您会发现前三个进程不会崩溃，但它们会卡住。默认情况下，如果一个进程在参与分布式计算（如 jnp.sum）时失败，则参与计算的其余进程将永远卡住。

取消集体通信#

我们可以通过取消包含失败参与者的集体操作来避免卡住。我们可以通过提供一些额外的标志和环境变量（如下所示）来启用集体取消。

 1import os
 2os.environ['XLA_FLAGS'] = ' '.join([
 3  '--xla_gpu_nccl_terminate_on_error=false',
 4  '--xla_gpu_nccl_async_execution=true',
 5  '--xla_gpu_nccl_blocking_communicators=false',
 6])
 7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'
 8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'
 9
10from absl import app
11from absl import flags
12from collections.abc import Sequence
13import jax
14import jax.numpy as jnp
15import time
16
17_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
18_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
19
20
21def main(_: Sequence[str]) -> None:
22  jax.config.update("jax_enable_recoverability", True)
23  jax.distributed.initialize(
24      coordinator_address="localhost:9000",
25      num_processes=_NUM_PROCESSES.value,
26      process_id=_PROCESS_ID.value,
27      local_device_ids=[_PROCESS_ID.value],
28      heartbeat_timeout_seconds=10,
29  )
30  print(f'{jax.devices()=}')
31  print(f'{jax.local_devices()=}')
32
33  # Don't do this. Use live_devices instead.
34  from jax.experimental.multihost_utils import _live_devices
35  _live_devices(jax._src.distributed.global_state.client, jax.devices())
36
37  n = jax.device_count()
38  jax.set_mesh(jax.make_mesh((n,), ("i",)))
39  x = jax.device_put(jnp.arange(n), jax.P("i"))
40  while True:
41    print(jnp.sum(x))
42    time.sleep(1)
43
44
45if __name__ == "__main__":
46  app.run(main)

我们还需要插入一个对 jax.experimental.multihost_utils._live_devices 的调用才能使脚本正常工作。通常您不应该这样做。相反，您应该使用我们稍后将介绍的 live_devices API。目前，_live_devices 是在我们解释正确的 API 之前让脚本运行的一种技巧。

再次运行脚本并杀死第四个进程。前三个进程会卡在 jnp.sum 调用中，但约十秒后，调用将被取消，jnp.sum 将抛出一个类似于此的异常：

jaxlib._jax.XlaRuntimeError: FAILED_PRECONDITION: Task with incarnation id 3446767950926952685 is not connected

识别存活进程#

进程死亡后，剩余的存活进程需要了解谁死了，谁还活着。为此，我们可以使用 JAX 的核心容错 API：live_devices。live_devices 是一个上下文管理器，它接受设备列表作为参数，并返回这些设备中处于存活状态的子集。下面，我们将编辑 example.py 以调用 live_devices。

 1import os
 2os.environ['XLA_FLAGS'] = ' '.join([
 3  '--xla_gpu_nccl_terminate_on_error=false',
 4  '--xla_gpu_nccl_async_execution=true',
 5  '--xla_gpu_nccl_blocking_communicators=false',
 6])
 7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'
 8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'
 9
10from absl import app
11from absl import flags
12from collections.abc import Sequence
13from jax.experimental.multihost_utils import live_devices
14import jax
15import jax.numpy as jnp
16import time
17
18_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
19_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
20
21
22def main(_: Sequence[str]) -> None:
23  jax.config.update("jax_enable_recoverability", True)
24  jax.distributed.initialize(
25      coordinator_address="localhost:9000",
26      num_processes=_NUM_PROCESSES.value,
27      process_id=_PROCESS_ID.value,
28      local_device_ids=[_PROCESS_ID.value],
29      heartbeat_timeout_seconds=10,
30  )
31  print(f'{jax.devices()=}')
32  print(f'{jax.local_devices()=}')
33
34  while True:
35    try:
36      with live_devices(jax.devices()) as devices:
37        print(f'{devices=}')
38        n = len(devices)
39        jax.set_mesh(jax.make_mesh((n,), ("i",), devices=devices))
40        x = jax.device_put(jnp.arange(n), jax.P("i"))
41        print(jnp.sum(x))
42    except Exception as e:
43      print('FAIL:', e)
44    else:
45      print('PASS')
46    time.sleep(1)
47
48
49if __name__ == "__main__":
50  app.run(main)

在上面突出显示的代码中，我们使用所有设备（jax.devices()）调用 live_devices 以获取存活设备集 devices。然后，我们将数组 x 分片到这些设备上并执行 jnp.sum。如果进程在执行 jnp.sum 时失败，则 jnp.sum 将被取消并在剩余的存活设备上引发异常。从技术上讲，不能保证集体操作一定会失败。我们将在原子性中重新讨论这一点。目前，假设它会失败。

注意

jax.devices() 始终返回所有设备的集合，即使其中一些设备处于已失败的进程上。请使用 jax.experimental.multihost_utils.live_devices 来了解其中哪些设备是存活的。

再次运行脚本并杀死第四个进程。请注意，其余三个存活进程捕获了 jnp.sum 引发的异常，并继续执行 while 循环的下一次迭代。在下一次迭代中，devices 不包含已失败的第四个进程上的设备。即使第四个进程已死亡，三个存活进程也能继续正确执行。

接下来，重启第四个进程。请注意，在第四个进程重启后，其设备再次包含在 live_devices 返回的存活设备集合中。然后所有四个进程继续正常执行。

乍一看，live_devices 似乎很简单。你给它一个设备列表，它返回存活的那些。这能有多复杂？不幸的是，正如分布式系统中的许多事物一样，这里有很多微妙之处需要解决。接下来，我们将解释 live_devices 的屏障 (barrier) 语义和原子性属性。

屏障语义#

回想一下，多控制器 JAX 程序中的每个进程都应该同步运行。进程应该以相同的顺序执行相同的指令。否则，几乎肯定会导致死锁、崩溃或异常行为。

在 live_devices 的上下文中，我们需要确保每个进程就哪些进程当前存活达成一致。这很难保证，因为每个进程都在以不同的速度独立执行，并且进程可能随时失败。再次考虑上述在四个进程上运行的 example.py 脚本。想象一下，进程 1 和 2 调用了 live_devices，然后进程 4 失败，接着进程 3 调用了 live_devices。进程 1 和 2 可能认为进程 4 还活着，而进程 3 认为它已经死亡。

为了避免这种情况，live_devices 保证它向每个进程返回相同的存活设备集。它通过使用屏障来实现这一点。对 live_devices(devices) 的调用会阻塞，直到 devices 中包含设备的每个存活进程也调用了 live_devices。一旦每个存活进程都进入了 live_devices 屏障，live_devices 就会向每个进程返回相同的存活设备集。

重要

live_devices 使用屏障来确保它始终向每个存活进程返回相同的存活设备集。

因为 live_devices 实现了屏障，如果使用不当，它很容易导致死锁。我们建议在程序中只使用一个 with live_devices 块。多次调用 live_devices 难以推理，并可能导致死锁。

有关 live_devices 屏障是如何实现的，以及基于线性化 (linearizability) 的形式化语义，请参阅第 3 部分：实现细节。

原子性#

如果分布式计算的每个参与者都同意操作是成功还是失败，那么该计算就是原子的。在上面的 example.py 脚本中，我们看到当某个进程在 jnp.sum 执行期间失败时，jnp.sum 会中止并在其余存活进程上引发异常。那么 jnp.sum 是原子的吗？

不幸的是，它不是。

当进程在执行集体操作（如 jnp.sum）期间失败时，剩余进程可能会取消操作并引发异常，或者它们可能会成功完成操作。JAX 中的集体操作没有任何固有的原子性属性。

但是，如果集体操作不是原子的，那么多控制器 JAX 进程可能会产生分歧。例如，如果进程在机器学习模型的训练步骤中失败，一些进程可能会检测到故障并将模型回滚到检查点，而其他进程可能认为步骤成功并继续训练。

为了避免非原子执行带来的复杂性，尽管集体操作不是原子的，live_devices 提供了其自己的原子性保证。具体来说，with live_devices 块的主体被保证要么在所有进程上成功完成，要么在所有进程上引发异常。更具体地说，如果我们考虑下面的代码片段，要么每个进程都执行分支 A，要么每个进程都执行分支 B。某些进程执行 A 而其他进程执行 B 是不可能的。

try:
  with live_devices(jax.devices()) as devices:
    ...
except Exception as e:
  ... # Branch A
else:
  ... # Branch B

警告

如果代码块由于除因崩溃进程而失败的集体操作之外的原因而非确定性地引发异常，则 with live_devices 块不保证原子性。例如，如果一个进程因为内存不足而引发异常，则此异常不会传播到其他进程。

回想一下，JAX 使用异步调度。像 jnp.sum 这样的操作在操作完成之前不会阻塞。相反，它们返回充当 Future 的 jax.Arrays。这种异步性可能会以意想不到的方式与 live_devices 交互。例如，考虑以下执行 jnp.sum、将结果分配给 y 然后打印 y 的代码：

x = ...
y = ...
try:
  with live_devices(jax.devices()) as devices:
    y = jnp.sum(x)
except Exception as e:
  ... # Branch A
else:
  ... # Branch B
print(y)

想象一下，with live_devices 块在所有进程上都成功执行。也就是说，所有进程都执行分支 B。这只能保证每个进程都成功创建了一个 Future 并将其分配给了 y。jnp.sum 的实际计算可能会延迟到块之外。因此，一些进程可能会成功完成 jnp.sum 并打印 y 的值，而其他进程可能无法完成 jnp.sum，并在尝试打印 y 时引发异常。

为了避免这种情况，请使用 jax.block_until_ready 来确保计算在 with live_devices 块内执行。下面的代码片段在分配给 y 时调用了 jax.block_until_ready，它保证每个进程都会成功执行 jnp.sum，或者每个进程都会引发异常。

x = ...
y = ...
try:
  with live_devices(jax.devices()) as devices:
    y = jax.block_until_ready(jnp.sum(x))
except Exception as e:
  ... # Branch A
else:
  ... # Branch B
print(y)

有关如何实现原子性的详细信息，请参阅第 3 部分：实现细节。

第 2 部分：示例#

live_devices 不是万能药；它是一个工具。它并不能神奇地使多控制器 JAX 程序具备容错能力。相反，它允许您以最适合您应用程序的方式亲自实现容错。

实现容错的具体细节将根据应用程序的性质而有很大差异。在本节中，我们介绍了一些如何使用 live_devices 的示例。这些示例旨在说明问题，而非规定性的。实现容错的方法还有很多。

示例 1：容错数据并行训练#

在此示例中，我们跨四个进程通过数据并行训练一个简单的单参数线性模型（\(y = \alpha x\)）。这个例子是人为设计的——您永远不会跨四台机器训练一个单参数模型——但我们故意保持模型简单以专注于容错。

数据并行使实现容错变得相对简单。因为每个进程都有模型权重的完整副本，如果进程失败，我们可以简单地忽略它并继续训练。此示例容忍任意数量的进程失败（进程 0 除外），但一旦进程失败，我们假设它不会恢复。下一个示例将展示如何处理进程恢复。

首先，我们设置一些标志来禁用命运共享并启用集体取消。我们还进行必要的导入并定义一些标志。

 1import os
 2os.environ['XLA_FLAGS'] = ' '.join([
 3    '--xla_gpu_nccl_terminate_on_error=false',
 4    '--xla_gpu_nccl_async_execution=true',
 5    '--xla_gpu_nccl_blocking_communicators=false',
 6])
 7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'
 8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'
 9
10from absl import app
11from absl import flags
12from collections.abc import Sequence
13from jax.experimental.multihost_utils import live_devices
14import jax
15import jax.numpy as jnp
16import time
17
18_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
19_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")

接下来，我们定义一个 replicated 函数，该函数返回一个在设备集上复制的数组。请注意，replicated 实际上并不移动任何数据。它假设参数 x 在所有进程中已经具有相同的值。它只是返回该数据的新视图，即一个具有复制分片的、跨进程的 jax.Array。

21def replicated(x: jax.Array, devices: list[jax.Device]):
22  """Return x replicated across the provided devices.
23
24  Note that replicated(x) doesn't actually move any data. It simply creates a
25  logically replicated array with x as the local replica.
26  """
27  n = len(devices)
28  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
29  spec = jax.sharding.PartitionSpec(None)
30  sharding = jax.sharding.NamedSharding(mesh, spec)
31  shards = [
32      jax.device_put(x.addressable_shards[0].data, d) for d in devices
33      if d.process_index == jax.process_index()
34  ]
35  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)

我们定义了一个类似的 sharded 函数，该函数返回一个在设备集上分片的数组。同样，sharded 实际上不会在进程间移动任何数据。

38def sharded(x: jax.Array, devices: list[jax.Device]):
39  """Return x sharded across the provided devices.
40
41  Note that sharded(x) doesn't actually move any data. It simply creates a
42  logically sharded array. x should have the same shape as the global array.
43  """
44  n = len(devices)
45  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
46  spec = jax.sharding.PartitionSpec("i")
47  sharding = jax.sharding.NamedSharding(mesh, spec)
48  m = sharding.addressable_devices_indices_map(x.shape)
49  shards = [jax.device_put(x[m[d]], d) for d in jax.local_devices()]
50  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)

现在，我们准备开始编写训练循环。我们首先通过调用 jax.distributed.initialize 初始化多控制器 JAX。

53def main(_: Sequence[str]) -> None:
54  # Parse command line arguments and initialize multi-controller JAX.
55  jax.config.update("jax_enable_recoverability", True)
56  jax.distributed.initialize(coordinator_address="localhost:8000",
57                             process_id=_PROCESS_ID.value,
58                             num_processes=_NUM_PROCESSES.value,
59                             local_device_ids=[_PROCESS_ID.value],
60                             heartbeat_timeout_seconds=10)
61  print(f'{jax.devices()=}')
62  print(f'{jax.local_devices()=}')

然后，我们定义简单的线性模型，生成一些随机训练数据，并初始化一些基本超参数。

64  # Initialize the model's weights.
65  keys = iter(jax.random.split(jax.random.key(seed=42), num=3))
66  weights = jax.random.normal(next(keys), shape=(1, ))
67
68  # We'll learn a trivial linear model: a*x.
69  def predict(weights, X):
70    return weights * X
71
72  # We'll use mean squared error loss.
73  def loss(weights, X, Y):
74    return jnp.mean((predict(weights, X) - Y)**2)
75
76  # Initialize the (noisy) training data with a=10.
77  X = jax.random.permutation(next(keys), jnp.arange(-300., 300.))
78  Y = 10 * X + jax.random.normal(next(keys), X.shape)
79
80  # Hyperparameters.
81  loss_and_grad = jax.jit(jax.value_and_grad(loss))
82  learning_rate = 1e-6
83  device_batch_size = 10

最后，我们进入主训练循环。

 85  step = 0
 86  while True:
 87    try:
 88      with live_devices(jax.devices()) as devices:
 89        print(f'=== Running step {step} with live devices = {devices} ===')
 90
 91        # Replicate the model weights.
 92        weights = replicated(weights, devices)
 93
 94        # Shard the batch.
 95        batch_size = device_batch_size * len(devices)
 96        start = (step * batch_size) % len(X)
 97        stop = start + batch_size
 98        X_batch = sharded(X[start:stop], devices)
 99        Y_batch = sharded(Y[start:stop], devices)
100
101        # Compute gradients and update weights.
102        l, grad = loss_and_grad(weights, X_batch, Y_batch)
103        new_weights = jax.block_until_ready(weights - learning_rate * grad)
104    except Exception as e:
105      print(f'Step {step} failed: {e}')
106    else:
107      print(f'Step {step} succeeded: loss = {l}')
108      step += 1
109      weights = new_weights
110
111    time.sleep(1)

  • 循环的每次迭代，我们都会调用 live_devices 来了解当前哪些设备是存活的。

  • 然后，我们确保模型权重在这些设备上进行复制，并确保训练数据在这些设备上分片。请注意，这实际上不会在设备之间移动任何数据；它只是创建具有适当复制和分片元数据的 JAX 数组。

  • 我们调用 loss_and_grad 来计算权重相对于当前数据批次的梯度，然后计算新权重。请注意，如果训练步骤失败，我们将新权重分配给 new_weights，而不是直接分配给 weights。我们还调用 jax.block_until_ready 以确保我们在退出 live_devices 块时每个进程都已计算出新权重。

  • 如果训练步骤执行期间没有进程失败，则执行 else 分支。步骤增加，weights 更新。否则，将引发异常并执行 except 分支。在这种情况下，我们不会更新 step 或 weights，并在下一次迭代中使用新的存活设备集重试该步骤。

这是完整示例：

  1import os
  2os.environ['XLA_FLAGS'] = ' '.join([
  3    '--xla_gpu_nccl_terminate_on_error=false',
  4    '--xla_gpu_nccl_async_execution=true',
  5    '--xla_gpu_nccl_blocking_communicators=false',
  6])
  7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'
  8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'
  9
 10from absl import app
 11from absl import flags
 12from collections.abc import Sequence
 13from jax.experimental.multihost_utils import live_devices
 14import jax
 15import jax.numpy as jnp
 16import time
 17
 18_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
 19_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
 20
 21def replicated(x: jax.Array, devices: list[jax.Device]):
 22  """Return x replicated across the provided devices.
 23
 24  Note that replicated(x) doesn't actually move any data. It simply creates a
 25  logically replicated array with x as the local replica.
 26  """
 27  n = len(devices)
 28  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
 29  spec = jax.sharding.PartitionSpec(None)
 30  sharding = jax.sharding.NamedSharding(mesh, spec)
 31  shards = [
 32      jax.device_put(x.addressable_shards[0].data, d) for d in devices
 33      if d.process_index == jax.process_index()
 34  ]
 35  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)
 36
 37
 38def sharded(x: jax.Array, devices: list[jax.Device]):
 39  """Return x sharded across the provided devices.
 40
 41  Note that sharded(x) doesn't actually move any data. It simply creates a
 42  logically sharded array. x should have the same shape as the global array.
 43  """
 44  n = len(devices)
 45  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
 46  spec = jax.sharding.PartitionSpec("i")
 47  sharding = jax.sharding.NamedSharding(mesh, spec)
 48  m = sharding.addressable_devices_indices_map(x.shape)
 49  shards = [jax.device_put(x[m[d]], d) for d in jax.local_devices()]
 50  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)
 51
 52
 53def main(_: Sequence[str]) -> None:
 54  # Parse command line arguments and initialize multi-controller JAX.
 55  jax.config.update("jax_enable_recoverability", True)
 56  jax.distributed.initialize(coordinator_address="localhost:8000",
 57                             process_id=_PROCESS_ID.value,
 58                             num_processes=_NUM_PROCESSES.value,
 59                             local_device_ids=[_PROCESS_ID.value],
 60                             heartbeat_timeout_seconds=10)
 61  print(f'{jax.devices()=}')
 62  print(f'{jax.local_devices()=}')
 63
 64  # Initialize the model's weights.
 65  keys = iter(jax.random.split(jax.random.key(seed=42), num=3))
 66  weights = jax.random.normal(next(keys), shape=(1, ))
 67
 68  # We'll learn a trivial linear model: a*x.
 69  def predict(weights, X):
 70    return weights * X
 71
 72  # We'll use mean squared error loss.
 73  def loss(weights, X, Y):
 74    return jnp.mean((predict(weights, X) - Y)**2)
 75
 76  # Initialize the (noisy) training data with a=10.
 77  X = jax.random.permutation(next(keys), jnp.arange(-300., 300.))
 78  Y = 10 * X + jax.random.normal(next(keys), X.shape)
 79
 80  # Hyperparameters.
 81  loss_and_grad = jax.jit(jax.value_and_grad(loss))
 82  learning_rate = 1e-6
 83  device_batch_size = 10
 84
 85  step = 0
 86  while True:
 87    try:
 88      with live_devices(jax.devices()) as devices:
 89        print(f'=== Running step {step} with live devices = {devices} ===')
 90
 91        # Replicate the model weights.
 92        weights = replicated(weights, devices)
 93
 94        # Shard the batch.
 95        batch_size = device_batch_size * len(devices)
 96        start = (step * batch_size) % len(X)
 97        stop = start + batch_size
 98        X_batch = sharded(X[start:stop], devices)
 99        Y_batch = sharded(Y[start:stop], devices)
100
101        # Compute gradients and update weights.
102        l, grad = loss_and_grad(weights, X_batch, Y_batch)
103        new_weights = jax.block_until_ready(weights - learning_rate * grad)
104    except Exception as e:
105      print(f'Step {step} failed: {e}')
106    else:
107      print(f'Step {step} succeeded: loss = {l}')
108      step += 1
109      weights = new_weights
110
111    time.sleep(1)
112
113
114if __name__ == "__main__":
115  app.run(main)

示例 2：具备恢复功能的容错数据并行训练#

现在，我们修改上面的示例以允许失败的进程恢复。当进程恢复时，它需要接收当前的步骤和模型权重。因为我们假设进程 0 永远不会失败（回想一下，如果进程 0 失败，所有进程都会失败），我们让进程 0 将当前的步骤和权重发送给恢复中的进程。

首先，我们定义 send 和 recv 函数，它们使用 shard_map 将数据从一个设备发送到另一个设备。发送方调用 send，接收方调用 recv。

55def send(x: jax.Array, from_device: jax.Device, to_device: jax.Device):
56  """Sends x from one device to another."""
57  assert isinstance(x, jax.Array)
58  devices = [from_device, to_device]
59  psum = lambda x: jax.lax.psum(x, "i")
60  mesh = jax.make_mesh((2, ), ("i", ), devices=devices)
61  spec = jax.sharding.PartitionSpec(None)
62  x = replicated(x, [from_device, to_device])
63  shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)
64
65
66def recv(x: jax.Array, from_device: jax.Device, to_device: jax.Device):
67  """Receives x from a matching send."""
68  assert isinstance(x, jax.Array)
69  to_device = jax.local_devices()[0]
70  devices = [from_device, to_device]
71  psum = lambda x: jax.lax.psum(x, "i")
72  mesh = jax.make_mesh((2, ), ("i", ), devices=devices)
73  spec = jax.sharding.PartitionSpec(None)
74  x = jnp.zeros_like(x)
75  x = replicated(x, [from_device, to_device])
76  return shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)

allgather 在设备集上执行单个浮点数的 AllGather。

79def allgather(x: float, devices: list[jax.Device]) -> list[float]:
80  """Performs an AllGather across the provided devices."""
81  n = len(devices)
82  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
83  spec = jax.sharding.PartitionSpec('i')
84  p = lambda x: jax.lax.all_gather(x, "i", tiled=True)
85  f = jax.shard_map(p, mesh=mesh, in_specs=spec, out_specs=spec)
86  return jax.block_until_ready(f(np.array([x] * len(devices)))).addressable_shards[0].data

最后，我们修改训练循环以处理恢复中的进程，如下面突出显示的代码所示。

121  step = 0
122  while True:
123    try:
124      with live_devices(jax.devices()) as devices:
125        print(f'=== Running step {step} with live devices = {devices} ===')
126
127        # Handle recovering devices. A device is recovering if its step doesn't
128        # match process 0's step. We assume process 0 never fails.
129        print('all gathering steps...')
130        steps = allgather(step, devices)
131        print(f'{steps=}')
132        recovering = [d for d, s in zip(devices, steps) if s != steps[0]]
133        for d in recovering:
134          # Process 0 sends weights and step to the recovering devices.
135          if jax.process_index() == 0:
136            print('sending...')
137            send(weights, jax.devices()[0], d)
138            send(jnp.array([step]), jax.devices()[0], d)
139          elif d.process_index == jax.process_index():
140            print('receiving...')
141            weights = recv(weights, jax.devices()[0], d)
142            step = recv(jnp.array([step]), jax.devices()[0], d)[0]
143
144        # Replicate the model weights.
145        weights = replicated(weights, devices)
146
147        # Shard the batch.
148        batch_size = device_batch_size * len(devices)
149        start = (step * batch_size) % len(X)
150        stop = start + batch_size
151        X_batch = sharded(X[start:stop], devices)
152        Y_batch = sharded(Y[start:stop], devices)
153
154        # Compute gradients and update weights.
155        l, grad = loss_and_grad(weights, X_batch, Y_batch)
156        new_weights = jax.block_until_ready(weights - learning_rate * grad)
157    except Exception as e:
158      print(f'Step {step} failed: {e}')
159    else:
160      print(f'Step {step} succeeded: loss = {l}')
161      step += 1
162      weights = new_weights
163
164    time.sleep(1)

恢复是一个两步过程。首先，我们需要检测哪些进程正在恢复。其次，我们需要进程 0 将步骤和权重发送给正在恢复的进程。

  1. 为了检测哪些进程正在恢复，我们对所有存活进程的步骤执行 AllGather。当失败的进程恢复时，其 step 将为 0，而进程 0 上的 step 将是一个正数，因此如果一个进程的 step 不等于进程 0 的 step，则它正在恢复。

  2. 然后，我们调用上面定义的 send 和 recv 函数，将当前的 step 和模型权重从进程 0 传输到恢复中的进程。

这是完整示例：

  1import os
  2os.environ['XLA_FLAGS'] = ' '.join([
  3    '--xla_gpu_nccl_terminate_on_error=false',
  4    '--xla_gpu_nccl_async_execution=true',
  5    '--xla_gpu_nccl_blocking_communicators=false',
  6])
  7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'
  8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'
  9
 10from absl import app
 11from absl import flags
 12from collections.abc import Sequence
 13from jax.experimental.multihost_utils import live_devices
 14from jax.experimental import shard_map
 15import jax
 16import jax.numpy as jnp
 17import numpy as np
 18import time
 19
 20_PROCESS_ID = flags.DEFINE_integer("i", -1, "Process id")
 21_NUM_PROCESSES = flags.DEFINE_integer("n", -1, "Number of processes")
 22
 23def replicated(x: jax.Array, devices: list[jax.Device]):
 24  """Return x replicated across the provided devices.
 25
 26  Note that replicated(x) doesn't actually move any data. It simply creates a
 27  logically replicated array with x as the local replica.
 28  """
 29  n = len(devices)
 30  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
 31  spec = jax.sharding.PartitionSpec(None)
 32  sharding = jax.sharding.NamedSharding(mesh, spec)
 33  shards = [
 34      jax.device_put(x.addressable_shards[0].data, d) for d in devices
 35      if d.process_index == jax.process_index()
 36  ]
 37  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)
 38
 39
 40def sharded(x: jax.Array, devices: list[jax.Device]):
 41  """Return x sharded across the provided devices.
 42
 43  Note that sharded(x) doesn't actually move any data. It simply creates a
 44  logically sharded array. x should have the same shape as the global array.
 45  """
 46  n = len(devices)
 47  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
 48  spec = jax.sharding.PartitionSpec("i")
 49  sharding = jax.sharding.NamedSharding(mesh, spec)
 50  m = sharding.addressable_devices_indices_map(x.shape)
 51  shards = [jax.device_put(x[m[d]], d) for d in jax.local_devices()]
 52  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)
 53
 54
 55def send(x: jax.Array, from_device: jax.Device, to_device: jax.Device):
 56  """Sends x from one device to another."""
 57  assert isinstance(x, jax.Array)
 58  devices = [from_device, to_device]
 59  psum = lambda x: jax.lax.psum(x, "i")
 60  mesh = jax.make_mesh((2, ), ("i", ), devices=devices)
 61  spec = jax.sharding.PartitionSpec(None)
 62  x = replicated(x, [from_device, to_device])
 63  shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)
 64
 65
 66def recv(x: jax.Array, from_device: jax.Device, to_device: jax.Device):
 67  """Receives x from a matching send."""
 68  assert isinstance(x, jax.Array)
 69  to_device = jax.local_devices()[0]
 70  devices = [from_device, to_device]
 71  psum = lambda x: jax.lax.psum(x, "i")
 72  mesh = jax.make_mesh((2, ), ("i", ), devices=devices)
 73  spec = jax.sharding.PartitionSpec(None)
 74  x = jnp.zeros_like(x)
 75  x = replicated(x, [from_device, to_device])
 76  return shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)
 77
 78
 79def allgather(x: float, devices: list[jax.Device]) -> list[float]:
 80  """Performs an AllGather across the provided devices."""
 81  n = len(devices)
 82  mesh = jax.make_mesh((n, ), ("i", ), devices=devices)
 83  spec = jax.sharding.PartitionSpec('i')
 84  p = lambda x: jax.lax.all_gather(x, "i", tiled=True)
 85  f = jax.shard_map(p, mesh=mesh, in_specs=spec, out_specs=spec)
 86  return jax.block_until_ready(f(np.array([x] * len(devices)))).addressable_shards[0].data
 87
 88
 89def main(_: Sequence[str]) -> None:
 90  # Parse command line arguments and initialize multi-controller JAX.
 91  jax.config.update("jax_enable_recoverability", True)
 92  jax.distributed.initialize(coordinator_address="localhost:8000",
 93                             process_id=_PROCESS_ID.value,
 94                             num_processes=_NUM_PROCESSES.value,
 95                             local_device_ids=[_PROCESS_ID.value],
 96                             heartbeat_timeout_seconds=10)
 97  print(f'{jax.devices()=}')
 98  print(f'{jax.local_devices()=}')
 99
100  # Initialize the model's weights.
101  keys = iter(jax.random.split(jax.random.key(seed=42), num=3))
102  weights = jax.random.normal(next(keys), shape=(1, ))
103
104  # We'll learn a trivial linear model: a*x.
105  def predict(weights, X):
106    return weights * X
107
108  # We'll use mean squared error loss.
109  def loss(weights, X, Y):
110    return jnp.mean((predict(weights, X) - Y)**2)
111
112  # Initialize the (noisy) training data with a=10.
113  X = jax.random.permutation(next(keys), jnp.arange(-300., 300.))
114  Y = 10 * X + jax.random.normal(next(keys), X.shape)
115
116  # Hyperparameters.
117  loss_and_grad = jax.jit(jax.value_and_grad(loss))
118  learning_rate = 1e-6
119  device_batch_size = 10
120
121  step = 0
122  while True:
123    try:
124      with live_devices(jax.devices()) as devices:
125        print(f'=== Running step {step} with live devices = {devices} ===')
126
127        # Handle recovering devices. A device is recovering if its step doesn't
128        # match process 0's step. We assume process 0 never fails.
129        print('all gathering steps...')
130        steps = allgather(step, devices)
131        print(f'{steps=}')
132        recovering = [d for d, s in zip(devices, steps) if s != steps[0]]
133        for d in recovering:
134          # Process 0 sends weights and step to the recovering devices.
135          if jax.process_index() == 0:
136            print('sending...')
137            send(weights, jax.devices()[0], d)
138            send(jnp.array([step]), jax.devices()[0], d)
139          elif d.process_index == jax.process_index():
140            print('receiving...')
141            weights = recv(weights, jax.devices()[0], d)
142            step = recv(jnp.array([step]), jax.devices()[0], d)[0]
143
144        # Replicate the model weights.
145        weights = replicated(weights, devices)
146
147        # Shard the batch.
148        batch_size = device_batch_size * len(devices)
149        start = (step * batch_size) % len(X)
150        stop = start + batch_size
151        X_batch = sharded(X[start:stop], devices)
152        Y_batch = sharded(Y[start:stop], devices)
153
154        # Compute gradients and update weights.
155        l, grad = loss_and_grad(weights, X_batch, Y_batch)
156        new_weights = jax.block_until_ready(weights - learning_rate * grad)
157    except Exception as e:
158      print(f'Step {step} failed: {e}')
159    else:
160      print(f'Step {step} succeeded: loss = {l}')
161      step += 1
162      weights = new_weights
163
164    time.sleep(1)
165
166
167if __name__ == "__main__":
168  app.run(main)

第 3 部分：实现细节#

现在，我们深入探讨多控制器 JAX 的体系结构以及 live_devices 的语义和实现。如果您只对编写容错多控制器 JAX 程序感兴趣，本文的前两个部分就足够了。

协调服务#

当您启动一个多控制器 JAX 程序时，第一个进程（即进程 0）运行一个名为协调服务 (coordination service) 的独立 RPC 服务器。此外，所有进程（包括进程 0）都会创建一个连接到协调服务的 RPC 客户端。具体而言，jax.distributed.initialize() 的 coordinator_address 参数是协调服务的地址。此参数让进程 0 知道在哪个地址运行服务器，并让所有进程知道要连接到哪个地址。

协调服务实现了多控制器 JAX 控制平面 (control plane)。例如，它可以跨所有进程执行分布式屏障，并实现一个键值存储，供进程交换少量元数据。但请注意，数据平面 (data plane)（例如，程序数据上的所有集体操作）是在进程之间直接实现的，不涉及协调服务。

协调服务最重要的功能之一是健康检查。每个进程都会定期向协调服务发送心跳。如果进程失败，它将停止发送心跳。如果协调服务在一段时间内没有收到某个进程的心跳，它就会认定该进程已失败。

这显示在下面的交互式可视化中。协调服务显示在顶部，三个多控制器 JAX 进程显示在底部。请注意，进程如何定期向控制器发送心跳，控制器如何根据上次收到心跳的时间来跟踪每个进程的健康状况。尝试通过单击“失败”按钮终止进程 2。观察进程如何停止发送心跳，协调服务最终如何将进程视为死亡。

默认情况下，当协调服务检测到某个进程失败时，它会向所有其他进程发送消息，要求它们自我终止。换句话说，多控制器 JAX 程序中的所有进程都是命运共享的。再次在下面的可视化中通过单击“失败”按钮终止进程 2，观察协调服务如何通知其他进程失败。

这种命运共享意味着多控制器 JAX 程序根本不具备容错能力。它们是不容错的。为了启用容错，我们需要做两件事：

  • 首先，我们需要消除命运共享，并允许进程即使在对等进程死亡时也能继续执行。这可以使用 jax_enable_recoverability 选项启用，如第 1 部分：容错基础中所述。我们将假设此选项已设置。

  • 其次，我们需要提供一个 API，供进程了解哪些进程存活，哪些进程已失败。这就是在第 1 部分：容错基础中介绍的 live_devices API。

实现 live_devices API 涉及相当多的技术深度和微妙之处。我们将逐步介绍该 API 的设计和实现。我们将首先介绍一个更简单的 live_processes API，并逐步改进它，直到我们得到 live_devices API。

存活进程#

让我们尝试设计一个新的假设性 JAX API：jax.live_processes。顾名思义，我们希望 jax.live_processes() 返回所有当前存活进程的集合。下面是一个简单但（我们稍后会看到）不正确的实现。当进程调用 jax.live_processes() 时，它会向协调服务发送 RPC 请求。记住，协调服务已经使用心跳来跟踪哪些进程已死亡，哪些进程存活，因此当它收到 jax.live_processes 请求时，它会响应其认为存活的进程集合。

这在下面进行了说明。每个进程下方都有一个“调用 live_processes”按钮。您可以单击此按钮使进程调用 jax.live_processes。注意协调服务如何用存活进程集回复 live_processes 请求。通过单击“失败”按钮终止进程 2，看看它如何影响后续对 jax.live_processes 的调用。

这种简单的实现很简单，但不正确。多控制器 JAX 作业中的所有进程都必须以相同的顺序执行相同的指令，这一点至关重要。如果进程开始分歧（通过执行 JAX 程序中的不同代码路径），作业的行为将变得不规律。最可能的情况是它会崩溃、挂起或产生垃圾值，而且可以肯定的是，它将很难进行推理。

我们对 jax.live_processes 的简单实现很容易导致分歧。例如，考虑一个包含三个进程的多控制器 JAX 作业。如果进程 0 和 1 都在进程 2 失败的同时调用 jax.live_processes，协调服务可能会向进程 0 报告所有进程都存活，但向进程 1 报告只有进程 0 和 1 存活。尝试在下面的可视化中重现这种情况：

如果进程对于哪些进程存活存在分歧，它们几乎肯定会产生分歧。值得庆幸的是，我们可以通过利用屏障语义增强 jax.live_processes 来避免这种分歧。

屏障语义#

让我们更改 jax.live_processes 的实现，以便当协调服务收到 jax.live_processes() 请求时，它不会立即回复。相反，协调服务仅在每个存活进程都调用了 jax.live_processes() 后才会回复。一旦每个存活进程都进入了 jax.live_processes() 屏障，协调服务就会返回存活进程集。至关重要的是，协调服务向所有进程返回相同的存活进程集，这防止了进程分歧。

这在下面进行了说明。注意，协调服务器现在会跟踪哪些设备处于 live_processes 屏障中。尝试从每个进程调用 live_processes。注意协调服务在每个进程进入屏障之前是如何不响应的。然后终止进程 2 并从进程 0 和进程 1 调用 live_processes。

形式化语义#

分布式系统极其复杂。机器可能在任意时间失败，网络消息可能会丢失、延迟和乱序。在本节中，我们引入 jax.live_processes API 的形式化语义，以帮助应对这种复杂性。严谨地思考 jax.live_processes 的语义将帮助我们理解该 API 即使在病态执行情况下的行为。

我们将 jax.live_processes 的形式化语义建立在线性化 (linearizability) 基础之上：这是一种流行的形式化方法，用于定义许多分布式 API 的语义。具体来说，我们将分布式系统建模为若干进程。每个进程串行执行若干事件。事件分为四种类型：

  1. 进程可以启动 (👶)。我们假设当进程启动时，它会连接到协调服务，因此协调服务知道它已启动。

  2. 进程可以失败 (💀)。与启动不同，协调服务可能不会立即意识到进程已失败。

  3. 进程可以向协调服务发送 jax.live_processes 请求。

  4. 进程可以从协调服务接收 jax.live_processes 请求的回复。

下面是三个进程 0、1 和 2 的执行图。时间从左向右推进。首先，所有三个进程都启动了。这用婴儿表情符号显示。然后所有三个进程都向协调服务发送 jax.live_processes 请求。这显示为粗体彩色区域的开始。稍后，所有三个进程都从协调服务收到回复，其中 0,1,2 是存活设备的集合。

0 1 2 👶 0,1,2 👶 0,1,2 👶 0,1,2

在这种简单的执行中，很明显 jax.live_processes 的行为是正确的。我们可以用以下形式化语义来形式化这种直觉。

注意

如果每当 jax.live_processes 返回存活进程集 P 时，都存在一个瞬间，在该时刻 P 中的每个进程都处于 live_processes 屏障中，而其他所有进程都已死亡，则执行是有效的。如果一个 live_processes 的实现仅允许有效执行，则它是正确的。

稍后，我们将修改这些形式化语义以涵盖一些微妙的极端情况，但现在请假设这种简化的语义。

在上面的示例中，live_processes 返回 0,1,2。在下面的可视化中，我们表明确实存在一个瞬间，在该时刻进程 0、1 和 2 都处于屏障中，而所有其他进程（没有）都已死亡。该瞬间被绘制为一条垂直的红线。

0 1 2 👶 0,1,2 👶 0,1,2 👶 0,1,2

我们在上面的可视化中选择的特定瞬间并没有什么特别之处。重要的是存在某个瞬间，在该时刻 P 中的所有进程都处于屏障中，而其他所有进程都已死亡。满足此属性的瞬间有很多，如下所示。

0 1 2 👶 0,1,2 👶 0,1,2 👶 0,1,2

在下一个示例中，进程 0 和 1 启动，调用 live_devices，并收到 0,1 作为回复。进程 2 在整个执行过程中一直处于死亡状态。

0 1 2 👶 0,1 👶 0,1 💀

根据我们的形式化语义，这是一个有效的执行，因为存在一个瞬间，在该时刻进程 0 和 1 处于屏障中，而进程 2 已死亡。

0 1 2 👶 0,1 👶 0,1 💀

在接下来的执行中，进程 0 调用 jax.live_processes 并收到 0 的回复。进程 1 调用 jax.live_processes，但在收到回复之前死亡。

0 1 👶 0 👶 💀

这是一个有效的执行吗？是的。存在一个瞬间，在该时刻进程 0 处于屏障中，而进程 1 已死亡，如下所示。即使进程 1 调用了 jax.live_processes，也不能保证进程 1 会包含在协调服务的响应中。

例如，进程 1 的 jax.live_processes 请求可能已被网络丢弃，协调服务从未收到过。因此，从协调服务的角度来看，进程 1 已彻底死亡，甚至从未进入 live_processes 屏障。

0 1 👶 0 👶 💀

那么同样的执行，除了进程 0 现在从协调服务收到回复 0,1 呢？

0 1 👶 0,1 👶 💀

这同样是一个有效的执行，见下文。直观地，协调服务可以从进程 0 和 1 接收 jax.live_processes 请求，并将回复 0,1 发送给两者。当此回复在网络中时，进程 1 失败了。因此，即使进程 1 在进程 0 收到回复时已死亡，该执行仍然有效。

0 1 👶 0,1 👶 💀

这一点值得重复。如果 jax.live_processes 返回进程集 P，这并不意味着 P 中的所有进程目前都存活，而其他所有进程目前都已死亡。它只意味着曾存在过一个时间点，在那时这是正确的。

在接下来的执行中，进程 1 调用 jax.live_processes 并失败。后来，进程 0 启动，调用 jax.live_processes，并收到 0,1 作为回复。

0 1 👶 0,1 👶 💀

使用迄今为止描述的形式化语义，这不是一个有效的执行。从来没有一个时间点进程 0 和 1 都存活。然而，这应该是一个有效的执行。

原因在于一个不可避免的事实：在分布式系统中，不可能以 100% 的准确率检测故障。如果协调服务在一段时间内没有收到某个进程的心跳，它就会认定该进程死亡。但是，协调服务无法 100% 确定进程何时死亡，或者进程是否真的死亡了。也许该进程很久以前就死了，或者也许它刚刚死掉，或者也许它活着，但在网络分区的另一侧。

让我们回到上面的执行以获取具体的例子。想象协调服务成功收到了进程 1 的 live_processes 请求。然后，进程 1 失败了，但协调服务没有立即检测到故障。与此同时，协调服务收到了进程 0 的 live_processes 请求。此时，协调服务认为两个进程都存活，并看到两个进程都处于屏障中，因此它自然地向两个进程返回了 0,1（尽管只有进程 0 收到了回复，因为进程 1 已经死亡）。

协调服务认为进程 1 存活时，它实际上已经死了。有时协调服务可能认为一个进程死了，而它实际上还活着。虽然不是理想的，但我们需要适应这样的执行，因为它们是不可避免的。

我们修改了形式化语义，允许我们将故障向前或向后移动，尽管我们不能将故障移动到同一进程的不同事件之后。直观地，我们可以将故障从它实际发生的时间移动到协调服务认为它发生的时间点。继续上面的例子，我们可以延迟进程 1 的故障，以创造一个瞬间，在该瞬间进程 0 和 1 都处于屏障中，证明了该执行是有效的。

0 1 👶 0,1 👶 💀

考虑下面类似的执行：

0 1 👶 0 👶 💀

就目前而言，不存在一个进程 0 存活而进程 1 死亡的时间点。但是，如果我们向左移动进程 1 的故障，则存在。这样的执行是如何产生的？想象进程 1 与协调服务分区了。协调服务没有收到来自进程 1 的任何消息，包括其心跳。这导致协调服务得出结论，进程 1 已死亡，即使事实并非如此。然后，协调服务接收进程 0 的 live_processes 请求并以 0 进行响应。

0 1 👶 0 👶 💀

但是，我们不能将进程故障移动到进程的其他事件之后。例如，以下执行是无效的，因为无论我们将进程 1 的故障移动到哪里，都不会有一个时间点两个进程都处于屏障中。

0 1 👶 0,1 👶 👶 💀

有了这些形式化语义，我们甚至可以理解复杂的执行。例如，考虑以下执行：

0 1 2 👶 0 0,2 👶 💀 👶 💀 👶 💀

在移动了一些进程故障后，我们看到执行是有效的。

0 1 2 👶 0 0,2 👶 💀 👶 💀 👶 💀

另一方面，以下执行是无效的：

0 1 2 👶 0,2 👶 1 💀 👶 💀

原子性#

有了 jax.live_processes，让我们尝试编写一些容错的多控制器 JAX 代码。

step = 0
while True:
    # Get the devices on all live processes.
    procs = jax.live_processes()
    devices = [d for d in jax.devices() if d.process_index in procs]

    # Shard array x over these devices.
    mesh = jax.make_mesh((len(devices),), ("i",), devices=devices)
    spec = jax.sharding.PartitionSpec("i")
    sharding = jax.sharding.NamedSharding(mesh, spec)
    x = jax.make_array_from_process_local_data(sharding, np.ones(1))

    # Try to perform a jnp.sum.
    try:
        print(jnp.sum(x))
    except:
        # jnp.sum failed.
        pass
    else:
        # jnp.sum succeeded.
        step += 1

该代码反复：

  • 调用 jax.live_processes 以了解哪些进程存活，

  • 计算健康进程上的设备集，

  • 将数组分片到这些健康设备上，

  • 对数组执行 jnp.sum（即 AllReduce），以及

  • 如果 jnp.sum 成功，则递增 step。

这段代码看起来是正确的，但它有一个非常微妙的 Bug。假设 jnp.sum 正在一组进程 P 上执行。如果 P 中的一个（或多个）进程在 jnp.sum 执行期间失败，那么 jnp.sum 在不同进程上的表现可能会有所不同。P 中的一些进程可能看到 jnp.sum 返回正确结果。其他进程可能看到 jnp.sum 引发异常。还有一些可能看到 jnp.sum 返回不正确的结果。

警告

如果进程在集体操作期间失败，该操作在不同进程上的表现可能会有所不同。

这意味着执行上述代码示例的进程可能会分歧。一些可能会递增 step，而另一些则不会。在上面的简单代码示例中，这种分歧是良性的，但在实际程序中，这种分歧很可能导致崩溃、死锁或垃圾输出。例如，如果多控制器 JAX 程序正在通过数据并行训练模型并开始分歧，一些进程可能会将其模型权重回滚到先前的检查点，而其他进程则继续训练，导致产生一个“科学怪人模型”，没有人同意模型权重应该是什么。

为了编写不会分歧的容错代码，我们需要原子性。在执行代码块（如上面的 jnp.sum）时，我们希望每个进程都成功运行代码，或者每个进程都知道代码执行失败。我们不希望一些进程成功，而另一些进程失败。

值得庆幸的是，我们可以通过一个非常简单的技巧实现原子性：调用 live_processes 两次，一次在代码块之前，一次在代码块之后。如果块之前存活的所有进程在块之后也存活，则代码块在所有存活进程上都执行成功。另一方面，如果任何进程死亡，则所有剩余进程都可以同意代码块未能正确执行。以下是可能的样子：

# Get the set of live processes before the code block.
procs_before = jax.live_processes()

# Execute the code block.
...

# Get the set of live processes after the code block
procs_after = jax.live_processes()
if procs_before == procs_after:
    # The code block executed successfully on all processes in
    # procs_before.
    pass
else:
    # The code block did not execute successfully. All processes will
    # agree it failed.
    pass

上面的代码应该让您初步了解如何使用两次调用 live_processes 来实现原子性，但在它完全正确之前，我们还需要解决少数小问题。例如：

  • 如果代码块抛出异常怎么办？我们需要捕获异常，并在第二次调用 live_processes 后重新引发异常。

  • 如果进程在第一次调用 live_processes 之后失败，并在第二次调用之前恢复怎么办？代码块难道不会失败，但前后进程集是一样的吗？每次进程启动时，它都会生成一个随机的化身 ID (incarnation id)。除了检查进程集没有改变之外，我们还检查它们的化身 ID 也没有改变。

  • 如果一个进程恢复，并且它的第一次 live_processes 调用与另一个进程的第二次 live_processes 调用匹配怎么办？难道这不会导致死锁吗？是的。我们可以通过仅在单个程序点调用 live_processes 来避免这个问题。我们可以巧妙地使用单次 live_processes 调用来实现两个目的。它可用于检查进程集自上次调用 live_processes 以来没有改变，并可用于生成在下次执行原子代码块时应使用的存活进程集。

所有这些细节都由第 1 部分：容错基础中介绍的 live_devices API 处理和抽象化。live_devices 是一个确保代码块原子执行的上下文管理器。在下面的代码片段中，devices 是所有存活进程上的设备列表。代码块 A 将在这些进程上原子执行。也就是说，要么每个进程都会看到代码引发异常（分支 B），要么每个进程都会看到代码成功（分支 C）。

try:
  with live_devices() as devices:
    pass # A
except Exception as e:
  pass # B
else:
  pass # C

取消集体通信#

如取消集体通信中所述，如果参与集体通信的进程失败，则其他参与进程将永远卡住。我们需要显式取消这些集体通信，以允许存活的参与者取得进展。虽然 live_devices API 在所有 JAX 后端（即 CPU、GPU、TPU）上都受支持，但取消集体通信仅由 GPU 后端支持。在这里，我们简要解释集体取消背后的一些实现细节。

GPU 后端使用NCCL（NVIDIA 的集体通信库）实现集体通信。当一组进程想要执行集体通信时，它们会形成一个 NCCL 通信器 (NCCL communicator)。然后，进程可以使用此通信器重复执行集体操作。创建通信器很昂贵——它需要网络通信——因此 JAX 后端会缓存以参与进程集及其化身 ID 为键的通信器。

在内部，JAX 客户端会轮询协调服务以获取每个进程的当前状态。如果客户端检测到进程已死亡或以新的化身 ID 重启，则客户端会中止其缓存键中具有失败化身 ID 的所有通信器。