多控制器 JAX 简介（即多进程/多主机 JAX）

多控制器 JAX（又称多进程/多主机 JAX）简介#

通过阅读本教程，你将学习如何将 JAX 计算扩展到比单台主机所能容纳的设备更多的设备上，例如在 GPU 集群、Cloud TPU pod 或多台仅有 CPU 的机器上运行时的情况。

核心思想

  • 运行多个 Python 进程，我们有时将其称为“控制器”。每台主机可以运行一个（或多个）进程。

  • 使用 jax.distributed.initialize() 初始化集群。.

  • 一个 jax.Array 可以跨越所有进程，如果每个进程都对它应用相同的 JAX 函数，那么编程体验就像是在面对一台大型设备一样。

  • 使用与单控制器 JAX 相同的 统一分片机制 来控制数据如何分布以及计算如何并行化。当有 TPU ICI 或 NVLink 等高速网络链接可用时，XLA 会自动利用它们，否则将使用可用的主机网络（如以太网、InfiniBand）。

  • 所有进程（通常）运行相同的 Python 脚本。编写 Python 代码的方式与单进程几乎完全相同——只需运行它的多个实例，剩下的交由 JAX 处理即可。换句话说，除了数组创建之外，你可以编写 JAX 代码，就好像有一台连接了所有设备的大型机器一样。

本教程假设你已经阅读过 分布式数组和自动并行化，该文档主要针对单控制器 JAX。

Illustration of a multi-host TPU pod. Each host in the pod is attached via PCI to a board of four TPU chips. The TPUs chips themselves are connected via high-speed inter-chip interconnects.

多主机 TPU pod 的示意图。Pod 中的每台主机（绿色）通过 PCI 连接到一块包含四个 TPU 芯片（蓝色）的板卡。TPU 芯片本身通过高速芯片间互连 (ICI) 连接。JAX Python 代码运行在每台主机上（例如通过 ssh）。每台主机上的 JAX 进程彼此感知，允许你协调整个 pod 芯片组的计算。对于 GPU、CPU 以及其他支持 JAX 的平台，原则也是一样的！#

简易示例#

在我们定义术语并深入细节之前，先来看一个简易示例：创建一个跨进程的 jax.Array 并对其应用 jax.numpy 函数。

# call this file toy.py, to be run in each process simultaneously

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P
import numpy as np

# in this example, get multi-process parameters from sys.argv
import sys
proc_id = int(sys.argv[1])
num_procs = int(sys.argv[2])

# initialize the distributed system
jax.distributed.initialize('localhost:10000', num_procs, proc_id)

# this example assumes 8 devices total
assert jax.device_count() == 8

# make a 2D mesh that refers to devices from all processes
mesh = jax.make_mesh((4, 2), ('i', 'j'))

# create some toy data
global_data = np.arange(32).reshape((4, 8))

# make a process- and device-spanning array from our toy data
sharding = NamedSharding(mesh, P('i', 'j'))
global_array = jax.device_put(global_data, sharding)
assert global_array.shape == global_data.shape

# each process has different shards of the global array
for shard in global_array.addressable_shards:
  print(f"device {shard.device} has local data {shard.data}")

# apply a simple computation, automatically partitioned
global_result = jnp.sum(jnp.sin(global_array))
print(f'process={proc_id} got result: {global_result}')

在这里，mesh 包含了来自所有进程的设备。我们使用它来创建 global_array，在逻辑上这是一个单一的共享数组，存储在所有进程的设备上。

每个进程必须对 global_array 以相同的顺序应用相同的操作。XLA 会自动对这些计算进行分区，例如插入通信集合操作来计算整个数组的 jnp.sum。我们可以打印最终结果，因为它在各个进程中是复制的。

我们可以在本地 CPU 上运行此代码，例如使用 4 个进程，每个进程拥有 2 个 CPU 设备。

export JAX_NUM_CPU_DEVICES=2
num_processes=4

range=$(seq 0 $(($num_processes - 1)))

for i in $range; do
  python toy.py $i $num_processes > /tmp/toy_$i.out &
done

wait

for i in $range; do
  echo "=================== process $i output ==================="
  cat /tmp/toy_$i.out
  echo
done

输出

=================== process 0 output ===================
device TFRT_CPU_0 has local data [[0 1 2 3]]
device TFRT_CPU_1 has local data [[4 5 6 7]]
process=0 got result: -0.12398731708526611

=================== process 1 output ===================
device TFRT_CPU_131072 has local data [[ 8  9 10 11]]
device TFRT_CPU_131073 has local data [[12 13 14 15]]
process=1 got result: -0.12398731708526611

=================== process 2 output ===================
device TFRT_CPU_262144 has local data [[16 17 18 19]]
device TFRT_CPU_262145 has local data [[20 21 22 23]]
process=2 got result: -0.12398731708526611

=================== process 3 output ===================
device TFRT_CPU_393216 has local data [[24 25 26 27]]
device TFRT_CPU_393217 has local data [[28 29 30 31]]
process=3 got result: -0.12398731708526611

这看起来可能与单控制器 JAX 代码没什么不同，事实上，这正是你编写同一程序的单控制器版本的方式！（从技术上讲，单控制器不需要调用 jax.distributed.initialize()，但调用它也不会产生副作用。）让我们从单个进程运行相同的代码。

JAX_NUM_CPU_DEVICES=8 python toy.py 0 1

输出

device TFRT_CPU_0 has local data [[0 1 2 3]]
device TFRT_CPU_1 has local data [[4 5 6 7]]
device TFRT_CPU_2 has local data [[ 8  9 10 11]]
device TFRT_CPU_3 has local data [[12 13 14 15]]
device TFRT_CPU_4 has local data [[16 17 18 19]]
device TFRT_CPU_5 has local data [[20 21 22 23]]
device TFRT_CPU_6 has local data [[24 25 26 27]]
device TFRT_CPU_7 has local data [[28 29 30 31]]
process=0 got result: -0.12398731708526611

数据被分片到单进程的 8 个设备上，而不是 4 个进程的 8 个设备上，但除此之外，我们在相同的数据上运行相同的操作。

术语表#

有必要明确一些术语。

我们有时将每个运行 JAX 计算的 Python 进程称为控制器 (controller)，但这两个术语本质上是同义词。

每个进程都有一组本地设备 (local devices)，这意味着它可以与这些设备的内存之间进行数据传输，并在这些设备上运行计算，而无需涉及任何其他进程。本地设备通常物理连接到该进程对应的主机上（例如通过 PCI）。一个设备只能属于一个进程的本地设备；也就是说，本地设备集是不相交的。可以通过评估 jax.local_devices() 来查询进程的本地设备。我们有时使用可寻址 (addressable) 一词来表示相同的意思。

Illustration of how a process/controller and local devices fit into a larger multi-host cluster. The "global devices" are all devices in the cluster.

进程/控制器和本地设备如何融入更大规模多主机集群的示意图。“全局设备”是集群中的所有设备。#

跨越所有进程的设备被称为全局设备 (global devices)。全局设备列表通过 jax.devices() 查询。该列表通过在所有进程上运行 jax.distributed.initialize() 来填充，这会建立一个连接这些进程的简单分布式系统。

我们通常使用全局 (global) 和本地 (local) 来描述跨进程和进程内相关的概念。例如，“本地数组”可能是仅对单个进程可见的 numpy 数组，而 JAX “全局数组”在概念上对所有进程可见。

设置多个 JAX 进程#

在实践中，设置多个 JAX 进程看起来与在单台主机上运行的简易示例略有不同。我们通常在不同的主机上启动每个进程，或者拥有多台主机，每台主机运行多个进程。我们可以直接使用 ssh 或使用 Slurm 或 Kubernetes 等集群管理器来完成。无论哪种方式，你都必须手动在每台主机上运行你的 JAX 程序！ JAX 不会自动从单个程序调用启动多个进程。

无论它们是如何启动的，Python 进程都需要运行 jax.distributed.initialize()。使用 Slurm、Kubernetes 或任何 Cloud TPU 部署时，我们可以无参数地运行 jax.distributed.initialize()，因为参数会自动填充。初始化系统意味着我们可以运行 jax.devices() 来报告跨越所有进程的所有设备。

警告

jax.distributed.initialize() 必须在运行 jax.devices()、jax.local_devices() 或在设备上运行任何计算（例如使用 jax.numpy）之前调用。否则，JAX 进程将不会感知到任何非本地设备。（使用 jax.config() 或其他不涉及访问设备的功能是可以的。）如果你在访问任何设备后不小心调用了 jax.distributed.initialize()，它会抛出错误。

GPU 示例#

我们可以在 GPU 机器 集群上运行多控制器 JAX。例如，在 Google Cloud 上创建四个虚拟机，每个虚拟机配备两个 GPU 后，我们可以在每台虚拟机上运行以下 JAX 程序。在此示例中，我们显式地为 jax.distributed.initialize() 提供参数。协调器地址、进程 ID 和进程数量从命令行读取。

# In file gpu_example.py...

import jax
import sys

# Get the coordinator_address, process_id, and num_processes from the command line.
coord_addr = sys.argv[1]
proc_id = int(sys.argv[2])
num_procs = int(sys.argv[3])

# Initialize the GPU machines.
jax.distributed.initialize(coordinator_address=coord_addr,
                           num_processes=num_procs,
                           process_id=proc_id)
print("process id =", jax.process_index())
print("global devices =", jax.devices())
print("local devices =", jax.local_devices())

例如，如果第一台虚拟机的地址是 192.168.0.1，那么你将在第一台虚拟机上运行 python3 gpu_example.py 192.168.0.1:8000 0 4，在第二台虚拟机上运行 python3 gpu_example.py 192.168.0.1:8000 1 4，依此类推。在所有四台虚拟机上运行 JAX 程序后，第一个进程会打印以下内容。

process id = 0
global devices = [CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3), CudaDevice(id=4), CudaDevice(id=5), CudaDevice(id=6), CudaDevice(id=7)]
local devices = [CudaDevice(id=0), CudaDevice(id=1)]

该进程成功地将所有八个 GPU 视为全局设备，以及它的两个本地设备。类似地，第二个进程打印以下内容。

process id = 1
global devices = [CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3), CudaDevice(id=4), CudaDevice(id=5), CudaDevice(id=6), CudaDevice(id=7)]
local devices = [CudaDevice(id=2), CudaDevice(id=3)]

这台虚拟机看到了相同的全局设备，但有一组不同的本地设备。

TPU 示例#

作为另一个示例，我们可以在 Cloud TPU 上运行。在创建一个 v5litepod-16（它有 4 台主机）后，我们可能想测试是否可以连接这些进程并列出所有设备。

$ TPU_NAME=jax-demo
$ EXTERNAL_IPS=$(gcloud compute tpus tpu-vm describe $TPU_NAME --zone 'us-central1-a' \
                 | grep externalIp | cut -d: -f2)
$ cat << EOF > demo.py
import jax
jax.distributed.initialize()
if jax.process_index() == 0:
  print(jax.devices())
EOF
$ echo $EXTERNAL_IPS | xargs -n 1 -P 0 bash -c '
scp demo.py $0:
ssh $0 "pip -q install -U jax[tpu]"
ssh $0 "python demo.py" '

这里我们使用 xargs 并行运行多个 ssh 命令，每个命令在 TPU 主机之一上运行相同的 Python 程序。在 Python 代码中，我们使用 jax.process_index() 仅在一个进程上打印。它的打印结果如下：

[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=2, process_index=1, coords=(2,0,0), core_on_chip=0), TpuDevice(id=3, process_index=1, coords=(3,0,0), core_on_chip=0), TpuDevice(id=6, process_index=1, coords=(2,1,0), core_on_chip=0), TpuDevice(id=7, process_index=1, coords=(3,1,0), core_on_chip=0), TpuDevice(id=8, process_index=2, coords=(0,2,0), core_on_chip=0), TpuDevice(id=9, process_index=2, coords=(1,2,0), core_on_chip=0), TpuDevice(id=12, process_index=2, coords=(0,3,0), core_on_chip=0), TpuDevice(id=13, process_index=2, coords=(1,3,0), core_on_chip=0), TpuDevice(id=10, process_index=3, coords=(2,2,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=15, process_index=3, coords=(3,3,0), core_on_chip=0)]

呜呼，看看那些 TPU 核心！

Kubernetes 示例#

在 Kubernetes 集群上运行多控制器 JAX 在精神上与上述 GPU 和 TPU 示例几乎相同：每个 pod 运行相同的 Python 程序，JAX 发现其对等方，集群表现得像一台巨大的机器。

  1. 容器镜像 - 从启用了 JAX 的镜像开始，例如 Google Artifact Registry 上的公共 JAX AI 镜像（TPU / GPU）或 NVIDIA（NGC / JAX-Toolbox）。

  2. 工作负载类型 - 使用 JobSet 或 索引 Job (indexed Job)。每个副本对应一个 JAX 进程。

  3. 服务账户 - JAX 需要权限来列出属于该 Job 的 pod，以便进程发现彼此。最小 RBAC 设置可在 examples/k8s/svc-acct.yaml 中找到。

以下是一个 最小 JobSet，它启动了两个副本。将占位符（镜像、GPU 计数和任何私有注册表机密）替换为与你的环境相匹配的值。

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: jaxjob
spec:
  replicatedJobs:
  - name: workers
    template:
      spec:
        parallelism: 2
        completions: 2
        backoffLimit: 0
        template:
          spec:
            serviceAccountName: jax-job-sa  # kubectl apply -f svc-acct.yaml
            restartPolicy: Never
            imagePullSecrets:
              # https://k8s.io/docs/tasks/configure-pod-container/pull-image-private-registry/
            - name: null
            containers:
            - name: main
              image: null  # e.g. ghcr.io/nvidia/jax:jax
              imagePullPolicy: Always
              resources:
                limits:
                  cpu: 1
                  # https://k8s.io/docs/tasks/manage-gpus/scheduling-gpus/
                  nvidia.com/gpu: null
              command: 
                - python
              args:
                - -c
                - |
                  import jax
                  jax.distributed.initialize()
                  print(jax.devices())
                  print(jax.local_devices())
                  assert jax.process_count() > 1
                  assert len(jax.devices()) > len(jax.local_devices())

应用清单并观察 pod 完成情况

$ kubectl apply -f example.yaml
$ kubectl get pods -l jobset.sigs.k8s.io/jobset-name=jaxjob
NAME                       READY   STATUS      RESTARTS   AGE
jaxjob-workers-0-0-xpx8l   0/1     Completed   0          8m32s
jaxjob-workers-0-1-ddkq8   0/1     Completed   0          8m32s

当作业完成后，检查日志以确认每个进程都看到了所有加速器

$ kubectl logs -l jobset.sigs.k8s.io/jobset-name=jaxjob
[CudaDevice(id=0), CudaDevice(id=1)]
[CudaDevice(id=0)]
[CudaDevice(id=0), CudaDevice(id=1)]
[CudaDevice(id=1)]

每个 pod 应该具有相同的全局设备集和不同的本地设备集。此时，你可以将内联脚本替换为你的真实 JAX 程序。

一旦进程设置好，我们就可以开始构建全局 jax.Array 并运行计算。本教程中剩余的 Python 代码示例旨在在运行 jax.distributed.initialize() 后，在所有进程上同步运行。

Mesh、分片 (Sharding) 和计算可以跨越多个进程和主机#

在 JAX 中对多个进程进行编程通常看起来就像编程单个进程一样，只是设备更多了！主要例外在于进入或离开 JAX 的数据，例如从外部数据源加载时。我们将首先在这里介绍多进程计算的基础知识，它们在很大程度上与单进程对应物看起来相同。我们将在本文档的后面部分介绍一些数据加载的基本原理，即如何从非 JAX 源创建 JAX 数组。

回想一下 jax.sharding.Mesh 将一个 jax.Device 数组与一系列名称配对，每个数组轴对应一个名称。通过使用来自多个进程的设备创建 Mesh，然后在 jax.sharding.Sharding 中使用该 mesh，我们可以构建分片在来自多个进程的设备上的 jax.Array。

这是一个直接使用 jax.devices() 从所有进程获取设备来构建 Mesh 的示例：

from jax.sharding import Mesh
mesh = Mesh(jax.devices(), ('a',))

# in this case, the same as
mesh = jax.make_mesh((jax.device_count(),), ('a',))  # use this in practice

你在实践中应该使用 jax.make_mesh() 辅助工具，不仅因为它更简单，而且因为它能自动选择性能更高的设备排序。我们在这里把它写出来。默认情况下，它包含跨进程的所有设备，就像 jax.devices() 一样。

一旦有了 mesh，我们就可以在其上对数组进行分片。高效构建跨进程数组有几种方法，最后一节有详细说明，但现在我们为了简单起见只坚持使用 jax.device_put。

arr = jax.device_put(jnp.ones((32, 32)), NamedSharding(mesh, P('a')))
if jax.process_index() == 0:
  jax.debug.visualize_array_sharding(arr)

在进程 0 上，打印如下：

┌───────────────────────┐
│         TPU 0         │
├───────────────────────┤
│         TPU 1         │
├───────────────────────┤
│         TPU 4         │
├───────────────────────┤
│         TPU 5         │
├───────────────────────┤
│         TPU 2         │
├───────────────────────┤
│         TPU 3         │
├───────────────────────┤
│         TPU 6         │
├───────────────────────┤
│         TPU 7         │
├───────────────────────┤
│         TPU 8         │
├───────────────────────┤
│         TPU 9         │
├───────────────────────┤
│        TPU 12         │
├───────────────────────┤
│        TPU 13         │
├───────────────────────┤
│        TPU 10         │
├───────────────────────┤
│        TPU 11         │
├───────────────────────┤
│        TPU 14         │
├───────────────────────┤
│        TPU 15         │
└───────────────────────┘

让我们尝试一个稍微有趣一点的计算！

mesh = jax.make_mesh((jax.device_count() // 2, 2), ('a', 'b'))

def device_put(x, spec):
  return jax.device_put(x, NamedSharding(mesh, spec))

# construct global arrays by sharding over the global mesh
x = device_put(jnp.ones((4096, 2048)), P('a', 'b'))
y = device_put(jnp.ones((2048, 4096)), P('b', None))

# run a distributed matmul
z = jax.nn.relu(x @ y)

# inspect the sharding of the result
if jax.process_index() == 0:
  jax.debug.visualize_array_sharding(z)
  print()
  print(z.sharding)

在进程 0 上，打印如下：

┌───────────────────────┐
│        TPU 0,1        │
├───────────────────────┤
│        TPU 4,5        │
├───────────────────────┤
│        TPU 8,9        │
├───────────────────────┤
│       TPU 12,13       │
├───────────────────────┤
│        TPU 2,3        │
├───────────────────────┤
│        TPU 6,7        │
├───────────────────────┤
│       TPU 10,11       │
├───────────────────────┤
│       TPU 14,15       │
└───────────────────────┘

NamedSharding(mesh=Mesh('a': 8, 'b': 2), spec=PartitionSpec('a',), memory_kind=device)

在这里，仅仅通过在所有进程上评估 x @ y，XLA 就会自动生成并运行分布式矩阵乘法。结果像 P('a', None) 一样针对 mesh 进行分片，因为在这种情况下，矩阵乘法在 'b' 轴上包含了一个 psum。

警告

将 JAX 计算应用于跨进程数组时，为了避免死锁和挂起，至关重要的是所有拥有参与设备的进程都必须以相同的顺序运行相同的计算。这是因为计算可能涉及集合通信屏障。如果数组分片所在的设备因为其控制器没有发出相同的计算而没有加入集合操作，那么其他设备就会处于等待状态。例如，如果只有前三个进程评估了 x @ y，而最后一个进程评估了 y @ x，计算很可能会无限期挂起。这个假设（跨进程数组上的计算在所有参与进程上以相同顺序运行）在很大程度上是未经检查的。

因此，在多进程 JAX 中避免死锁的最简单方法是在每个进程上运行相同的 Python 代码，并警惕任何依赖于 jax.process_index() 且涉及通信的控制流。

如果一个跨进程数组被分片在不同进程的设备上，那么对该数组执行需要数据在进程本地可用的操作（例如打印）是错误的。例如，如果我们执行 print(z)，在前述示例中，我们会看到：

RuntimeError: Fetching value for `jax.Array` that spans non-addressable (non process local) devices is not possible. You can use `jax.experimental.multihost_utils.process_allgather` to print the global array or use `.addressable_shards` method of jax.Array to inspect the addressable (process local) shards.

要打印完整的数组值，我们必须首先确保它在各进程间已复制（但不必在每个进程的本地设备上复制），例如使用 jax.device_put。在上面的例子中，我们可以在最后写道：

w = device_put(z, P(None, None))
if jax.process_index() == 0:
  print(w)

注意不要将 jax.device_put() 写入 if process_index() == 0 下，因为这会导致死锁，只有进程 0 发起集合通信并无限期等待其他进程。 jax.experimental.multihost_utils 模块有一些函数可以更轻松地处理全局 jax.Array（例如 jax.experimental.multihost_utils.process_allgather()）。

或者，若要仅对进程本地数据进行打印或执行其他 Python 操作，我们可以访问 z.addressable_shards。访问该属性不需要任何通信，因此任何进程子集都可以执行它，而无需其他进程的参与。该属性在 jax.jit() 下不可用。

在部分设备子集上运行#

在上面的示例中，我们使用了跨越所有设备的全局 mesh。多控制器 JAX 还支持仅跨越部分设备的 mesh，这对于在不同设备上同时运行不同计算非常有用。

让我们定义一个仅包含一半全局设备的 mesh，并将一些数据放置在其上。我们将使用 jax.make_mesh 的 devices 参数来指示使用哪些设备。

num_devices = jax.device_count() // 2
mesh = jax.make_mesh((num_devices,), ('a',),
                     devices=jax.devices()[num_devices:],
                     axis_types=(jax.sharding.AxisType.Explicit,))
sharding = NamedSharding(mesh, P('a'))

data = np.arange(64).reshape((8, 8))
x = jax.device_put(data, sharding)

# inspect the sharding of the result
if jax.process_index() == 0:
  jax.debug.visualize_array_sharding(x)
  print()
  print(x.sharding)

# inspect the data local to each host
print(f"Devices attached to process {jax.process_index()}: {jax.local_devices()}")
print(f"Addressable data for process {jax.process_index()}:")
for shard in x.addressable_shards:
  print(f"device {shard.device} has local data {shard.data}")

分片（再次在四主机的 v5litepod-16 上）看起来如下：

┌───────────────────────┐
│         TPU 8         │
├───────────────────────┤
│         TPU 9         │
├───────────────────────┤
│        TPU 10         │
├───────────────────────┤
│        TPU 11         │
├───────────────────────┤
│        TPU 15         │
├───────────────────────┤
│        TPU 14         │
├───────────────────────┤
│        TPU 13         │
├───────────────────────┤
│        TPU 12         │
└───────────────────────┘

NamedSharding(mesh=Mesh('a': 8, axis_types=(Explicit,)), spec=PartitionSpec('a',), memory_kind=device)

只有进程 2 和 3 在分片中有本地设备；进程 0 和 1 不参与。由于进程 0 在数组中没有可寻址的数据，它打印以下内容：

Devices attached to process 0: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0)]

Addressable data for process 0:

进程 1 打印类似的内容。而进程 3 则在其本地设备上拥有该数组一半的数据，因此它打印以下内容：

Devices attached to process 3: [TpuDevice(id=10, process_index=3, coords=(2,2,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=15, process_index=3, coords=(3,3,0), core_on_chip=0)]
Addressable data for process 3
device TPU_10(process=3,(2,2,0,0)) has local data [[16 17 18 19 20 21 22 23]]
device TPU_11(process=3,(3,2,0,0)) has local data [[24 25 26 27 28 29 30 31]]
device TPU_15(process=3,(3,3,0,0)) has local data [[32 33 34 35 36 37 38 39]]
device TPU_14(process=3,(2,3,0,0)) has local data [[40 41 42 43 44 45 46 47]]

现在让我们运行与上述 toy.py 中相同的计算，并将此数组作为输入。这一次，只有连接到进程 2 和 3 的设备参与，并且集合操作 jnp.sum 仅在这些设备上复制结果。

result = jnp.sum(jnp.sin(x))
print(f"process={jax.process_index()} got result: {result}")

进程 2（和 3）可以打印结果：

process=2 got result: Array(0.09658563, dtype=float32)

进程 0 和 1 没有参与的设备，因此它们没有结果的本地副本。进程 0 打印：

process=0 got result: Array(shape=(), dtype=float32)

请记住，每个进程必须在参与分片的进程中以相同的顺序应用相同的操作。在之前的示例中，所有进程都参与了。在这个例子中，我们本可以让计算仅在进程 2 和 3 中运行，但在不参与的进程中应用它也是可以的，并且会产生一个没有本地数据的数组，如上面的进程 0 那样。

使用 jax.device_put 在进程间传输数据#

jax.device_put() 可以在跨进程的设备之间传输数据。与跨进程集合操作一样，数据在可用时通过 TPU ICI 或 NVLink 等高速网络链接传输。跨进程 jax.device_put() 必须在参与源或目标分片的所有主机上调用。

一个示例用例是流水线并行程序，其中每个阶段运行在不同的设备集上。在第一个阶段完成后，使用 jax.device_put() 将结果传输到另一组设备以进行流水线的下一阶段。这是一个简单的示意图：

# Create a sharding that contains half of the global devices for the first
# stage of the pipeline.
num_devices = jax.device_count() // 2
mesh_first_half = jax.make_mesh((num_devices,), ('a',),
                                devices=jax.devices()[:num_devices],
                                axis_types=(jax.sharding.AxisType.Explicit,))
sharding_first_half = NamedSharding(mesh_first_half, P('a'))

# Create a sharding that contains the other half of the devices.
mesh_second_half = jax.make_mesh((num_devices,), ('a',),
                                 devices=jax.devices()[num_devices:],
                                 axis_types=(jax.sharding.AxisType.Explicit,))
sharding_second_half = NamedSharding(mesh_second_half, P('a'))

# Place the input data on the first mesh.
data = np.arange(64).reshape((8, 8))
x = jax.device_put(data, sharding_first_half)

# `f` is the first stage of the pipeline.
@jax.jit
def f(x):
  # Arbitrary JAX computation.
  return x

# `g` is the second stage of the pipeline
@jax.jit
def g(x):
  # More JAX operations.
  return x

# Run the first stage on the first set of devices.
y = f(x)

# Transfer the data to the second set of devices.
# `device_put` must be called in all processes that participate in either
# `y.sharding` or `sharding_second_half`, so it's important to call `y = f(x)`
# in all processes -- not just those that participate in the first stage -- so
# that we always have a reference to `y`.
z = jax.device_put(y, sharding_second_half)

# Run the second stage on the second set of devices.
result = g(z)

得益于 JAX 的 异步分发，如果输入准备就绪，运行在不同设备上的 jax.jit() 函数和/或 jax.device_put() 将会并行运行。我们可以利用这一点来实现一个微批处理流水线并行性的非常简单的示例：

# Pipeline stage functions. Each stage will run on a different device.
pipeline_stages = [f, g, f, g]
devices = jax.devices()[:4]

microbatches = [np.arange(512**2).reshape((512, 512)) for _ in range(12)]

# Each microbatch is enqueued on each device sequentially, but each device
# conceptually has an independent queue of computations and transfers which can
# run in parallel across queues. For example, because there are no data
# dependencies between the microbatches, device 0 will immediately start a new
# microbatch once the previous is finished, overlapping with the `device_put` to
# device 1.
results = []
for mb in microbatches:
  for d, s in zip(devices, pipeline_stages):
    mb = jax.device_put(mb, d)
    mb = s(mb)
  results.append(mb)

跨进程 jax.device_put() 目前仅在源分片和目标分片包含相同数量的设备且具有相同的分片形状时支持。如果你觉得这太严格，请提交一个 Github Issue。

从外部数据创建跨进程数组#

从外部数据源（例如来自数据加载器的 numpy 数组）创建跨进程 jax.Array 主要有三种方式：

  1. 在所有进程上创建或加载完整数组，然后使用 jax.device_put() 将其分片到设备上；

  2. 在每个进程上创建或加载仅表示将本地分片并存储在该进程设备上的数据的数组，然后使用 jax.make_array_from_process_local_data() 将其分片到设备上；

  3. 在每个进程的设备上分别创建或加载单独的数组，每个数组代表要存储在该设备上的数据，然后使用 jax.make_array_from_single_device_arrays() 在没有任何数据移动的情况下组装它们。

后两种在实践中最常使用，因为在每个进程中物化完整的全局数据通常过于昂贵。

上面的简易示例使用了 jax.device_put()。

jax.make_array_from_process_local_data() 通常用于分布式数据加载。它不如 jax.make_array_from_single_device_arrays() 通用，因为它没有直接指定哪个进程本地数据的切片进入每个本地设备。这在加载数据并行批处理时很方便，因为每个微批处理放在哪个设备上并不重要。例如：

# target (micro)batch size across the whole cluster
batch_size = 1024
# how many examples each process should load per batch
per_process_batch_size = batch_size // jax.process_count()
# how many examples each device will process per batch
per_device_batch_size = batch_size // jax.device_count()

# make a data-parallel mesh and sharding
mesh = jax.make_mesh((jax.device_count(),), ('batch'))
sharding = NamedSharding(mesh, P('batch'))

# our "data loader". each process loads a different set of "examples".
process_batch = np.random.rand(per_process_batch_size, 2048, 42)

# assemble a global array containing the per-process batches from all processes
global_batch = jax.make_array_from_process_local_data(sharding, process_batch)

# sanity check that everything got sharded correctly
assert global_batch.shape[0] == batch_size
assert process_batch.shape[0] == per_process_batch_size
assert global_batch.addressable_shards[0].data.shape[0] == per_device_batch_size

jax.make_array_from_single_device_arrays() 是构建跨进程数组的最通用方法。它通常在执行 jax.device_put() 以将所需数据发送到每个设备后使用。这是最低级别的选项，因为所有数据移动都是手动执行的（例如通过 jax.device_put()）。示例：

shape = (jax.process_count(), jax.local_device_count())
mesh = jax.make_mesh(shape, ('i', 'j'))
sharding = NamedSharding(mesh, P('i', 'j'))

# manually create per-device data equivalent to np.arange(jax.device_count())
# i.e. each device will get a single scalar value from 0..N
local_arrays = [
    jax.device_put(
        jnp.array([[jax.process_index() * jax.local_device_count() + i]]),
        device)
    for i, device in enumerate(jax.local_devices())
]

# assemble a global array from the local_arrays across all processes
global_array = jax.make_array_from_single_device_arrays(
    shape=shape,
    sharding=sharding,
    arrays=local_arrays)

# sanity check
assert (np.all(
    jax.experimental.multihost_utils.process_allgather(global_array) ==
    np.arange(jax.device_count()).reshape(global_array.shape)))

所有这些方法也可用于创建仅跨越部分进程的数组。例如，我们可以使用 jax.make_array_from_single_device_arrays() 来创建一个跨越进程 0 和 1 上设备的数组：

num_participating_processes = 2
shape = (num_participating_processes, jax.local_device_count())
devices = (jax.local_devices(process_index=0) +
           jax.local_devices(process_index=1))
mesh = jax.make_mesh(shape, ('i', 'j'),
                     axis_types=(jax.sharding.AxisType.Explicit,) * 2,
                     devices=devices)
sharding = NamedSharding(mesh, P('i', 'j'))

# manually create per-device data in processes 0 and 1.
if jax.process_index() in (0, 1):
  local_arrays = [
      jax.device_put(
          jnp.array([[jax.process_index() * jax.local_device_count() + i]]),
          device)
      for i, device in enumerate(jax.local_devices())
  ]
else:
  local_arrays = []

# assemble an array from the local_arrays across processes 0 and 1
array = jax.make_array_from_single_device_arrays(
    shape=shape,
    sharding=sharding,
    arrays=local_arrays,
    dtype=jnp.int32)

# sanity check
if jax.process_index() in (0, 1):
  for shard in array.addressable_shards:
    assert shard.data.size == 1
else:
  assert not array.addressable_shards