分布式数组和自动并行化

本教程讨论通过 jax.Array 的并行化，这是 JAX v0.4.1 及更高版本中可用的统一数组对象模型。

from typing import Optional

import numpy as np

import jax
import jax.numpy as jnp

⚠️ 警告：笔记本需要 8 个设备才能运行。

if len(jax.local_devices()) < 8:
  raise Exception("Notebook requires 8 devices to run")

简介和一个快速示例

通过阅读本教程笔记本，您将了解 jax.Array，这是一种用于表示数组的统一数据类型，即使物理存储跨越多个设备。您还将了解如何使用 jax.Array 与 jax.jit 结合使用来提供基于编译器的自动并行化。

在我们逐步思考之前，这里有一个快速示例。首先，我们将创建一个跨多个设备分片的 jax.Array

from jax.experimental import mesh_utils
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding

# Create a Sharding object to distribute a value across devices:
mesh = Mesh(devices=mesh_utils.create_device_mesh((4, 2)),
            axis_names=('x', 'y'))

# Create an array of random values:
x = jax.random.normal(jax.random.key(0), (8192, 8192))
# and use jax.device_put to distribute it across devices:
y = jax.device_put(x, NamedSharding(mesh, P('x', 'y')))
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

接下来，我们将对它应用一个计算，并可视化结果值是如何跨多个设备存储的

z = jnp.sin(y)
jax.debug.visualize_array_sharding(z)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

对 jnp.sin 应用的评估是自动并行化到存储输入值（和输出值）的设备上的。

# `x` is present on a single device
%timeit -n 5 -r 5 jnp.sin(x).block_until_ready()

The slowest run took 8.96 times longer than the fastest. This could mean that an intermediate result is being cached.
25.2 ms ± 30.9 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

# `y` is sharded across 8 devices.
%timeit -n 5 -r 5 jnp.sin(y).block_until_ready()

2.4 ms ± 61.4 µs per loop (mean ± std. dev. of 5 runs, 5 loops each)

现在让我们更详细地看看这些部分!

分片 描述了数组值如何在跨设备的内存中布局

分片基础和 NamedSharding 子类

为了跨多个设备并行化计算，我们首先必须将输入数据布局在多个设备上。

在 JAX 中，分片 对象描述了分布式内存布局。它们可以与 jax.device_put 一起使用，以生成具有分布式布局的值。

例如，这里有一个具有单设备 分片 的值

import jax
x = jax.random.normal(jax.random.key(0), (8192, 8192))

jax.debug.visualize_array_sharding(x)

┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│         TPU 0         │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘

这里，我们使用 jax.debug.visualize_array_sharding 函数来显示值 x 在内存中的存储位置。所有 x 都存储在一个设备上，所以可视化相当无聊!

但是我们可以使用 jax.device_put 和 分片 对象跨多个设备分片 x。首先，我们使用 mesh_utils.create_device_mesh 创建一个 numpy.ndarray 的 设备，它会考虑硬件拓扑结构以确定 设备 顺序

from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental import mesh_utils

P = PartitionSpec

devices = mesh_utils.create_device_mesh((4, 2))
mesh = Mesh(devices, axis_names=('a', 'b'))
y = jax.device_put(x, NamedSharding(mesh, P('a', 'b')))
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

我们可以定义一个辅助函数来简化操作

devices = mesh_utils.create_device_mesh((4, 2))
default_mesh = Mesh(devices, axis_names=('a', 'b'))

def mesh_sharding(
    pspec: PartitionSpec, mesh: Optional[Mesh] = None,
  ) -> NamedSharding:
  if mesh is None:
    mesh = default_mesh
  return NamedSharding(mesh, pspec)

y = jax.device_put(x, mesh_sharding(P('a', 'b')))
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

这里，我们使用 P('a', 'b') 来表示 x 的第一个和第二个轴应该分别在设备网格轴 'a' 和 'b' 上分片。我们可以轻松地切换到 P('b', 'a')，以便将 x 的轴在不同的设备上分片

y = jax.device_put(x, mesh_sharding(P('b', 'a')))
jax.debug.visualize_array_sharding(y)

┌───────┬───────┬───────┬───────┐
│       │       │       │       │
│ TPU 0 │ TPU 2 │ TPU 6 │ TPU 4 │
│       │       │       │       │
│       │       │       │       │
├───────┼───────┼───────┼───────┤
│       │       │       │       │
│ TPU 1 │ TPU 3 │ TPU 7 │ TPU 5 │
│       │       │       │       │
│       │       │       │       │
└───────┴───────┴───────┴───────┘

# This `None` means that `x` is not sharded on its second dimension,
# and since the Mesh axis name 'b' is not mentioned, shards are
# replicated across it.
y = jax.device_put(x, mesh_sharding(P('a', None)))
jax.debug.visualize_array_sharding(y)

┌───────────────────────┐
│        TPU 0,1        │
├───────────────────────┤
│        TPU 2,3        │
├───────────────────────┤
│        TPU 6,7        │
├───────────────────────┤
│        TPU 4,5        │
└───────────────────────┘

这里，因为 P('a', None) 没有提到 网格 轴名 'b'，所以我们在轴 'b' 上获得了复制。这里的 None 只是充当占位符，与值 x 的第二个轴对齐，而不表示在任何网格轴上进行分片。（作为一种简写，尾部的 None 可以省略，因此 P('a', None) 与 P('a') 的含义相同。但是明确地写出来并不会有任何坏处！）

为了只在 x 的第二个轴上进行分片，我们可以使用 分片规范 中的 None 占位符

y = jax.device_put(x, mesh_sharding(P(None, 'b')))
jax.debug.visualize_array_sharding(y)

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

y = jax.device_put(x, mesh_sharding(P(None, 'a')))
jax.debug.visualize_array_sharding(y)

┌───────┬───────┬───────┬───────┐
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
│TPU 0,1│TPU 2,3│TPU 6,7│TPU 4,5│
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
│       │       │       │       │
└───────┴───────┴───────┴───────┘

对于固定的网格，我们甚至可以将 x 的一个逻辑轴在多个设备网格轴上进行分片

y = jax.device_put(x, mesh_sharding(P(('a', 'b'), None)))
jax.debug.visualize_array_sharding(y)

┌───────────────────────┐
│         TPU 0         │
├───────────────────────┤
│         TPU 1         │
├───────────────────────┤
│         TPU 2         │
├───────────────────────┤
│         TPU 3         │
├───────────────────────┤
│         TPU 6         │
├───────────────────────┤
│         TPU 7         │
├───────────────────────┤
│         TPU 4         │
├───────────────────────┤
│         TPU 5         │
└───────────────────────┘

使用 NamedSharding 使得在定义完设备网格并命名其轴后，很容易只在每个 device_put 的 分片规范 中引用这些名称。

计算遵循数据分片并自动并行化

有了分片输入数据，编译器就可以为我们提供并行计算。特别是，用 jax.jit 装饰的函数可以在分片数组上操作，而无需将数据复制到单个设备上。相反，计算遵循分片：根据输入数据的分片，编译器决定中间值和输出值的碎片，并并行化它们的评估，甚至在必要时插入通信操作。

例如，最简单的计算是逐元素计算

devices = mesh_utils.create_device_mesh((4, 2))
mesh = Mesh(devices, axis_names=('a', 'b'))

x = jax.device_put(x, NamedSharding(mesh, P('a', 'b')))
print('input sharding:')
jax.debug.visualize_array_sharding(x)

y = jnp.sin(x)
print('output sharding:')
jax.debug.visualize_array_sharding(y)

input sharding:
output sharding:

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

这里，对于逐元素操作 jnp.sin，编译器选择了与输入相同的输出分片。此外，编译器自动并行化了计算，因此每个设备并行地从其输入分片计算其输出分片。

换句话说，即使我们编写 jnp.sin 计算就好像一台机器要执行它一样，编译器也会为我们分割计算并将其在多个设备上执行。

我们也可以对不仅仅是逐元素操作进行相同的操作。考虑一个具有分片输入的矩阵乘法

y = jax.device_put(x, NamedSharding(mesh, P('a', None)))
z = jax.device_put(x, NamedSharding(mesh, P(None, 'b')))
print('lhs sharding:')
jax.debug.visualize_array_sharding(y)
print('rhs sharding:')
jax.debug.visualize_array_sharding(z)

w = jnp.dot(y, z)
print('out sharding:')
jax.debug.visualize_array_sharding(w)

lhs sharding:
rhs sharding:
out sharding:

┌───────────────────────┐
│        TPU 0,1        │
├───────────────────────┤
│        TPU 2,3        │
├───────────────────────┤
│        TPU 6,7        │
├───────────────────────┤
│        TPU 4,5        │
└───────────────────────┘

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

这里编译器选择了输出分片，以便它能够最大程度地并行化计算：无需通信，每个设备已经拥有它需要计算其输出分片所需的输入分片。

我们如何确定它是否真的在并行运行？我们可以做一个简单的计时实验

x_single = jax.device_put(x, jax.devices()[0])
jax.debug.visualize_array_sharding(x_single)

┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│         TPU 0         │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘

np.allclose(jnp.dot(x_single, x_single),
            jnp.dot(y, z))

True

%timeit -n 5 -r 5 jnp.dot(x_single, x_single).block_until_ready()

49.7 ms ± 349 µs per loop (mean ± std. dev. of 5 runs, 5 loops each)

%timeit -n 5 -r 5 jnp.dot(y, z).block_until_ready()

7.47 ms ± 44.8 µs per loop (mean ± std. dev. of 5 runs, 5 loops each)

即使复制分片 数组 也会生成一个具有输入分片的输出

w_copy = jnp.copy(w)
jax.debug.visualize_array_sharding(w_copy)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

因此，计算遵循数据放置：当我们使用 jax.device_put 显式地分片数据，并将函数应用于该数据时，编译器会尝试并行化计算并决定输出分片。这种针对分片数据的策略是 JAX 遵循显式设备放置策略的泛化。

当显式分片不一致时，JAX 会报错

但是如果一个计算的两个参数被显式地放置在不同的设备集上，或者具有不兼容的设备顺序？在这些模棱两可的情况下，会抛出错误

import textwrap
from termcolor import colored

def print_exception(e):
  name = colored(f'{type(e).__name__}', 'red', force_color=True)
  print(textwrap.fill(f'{name}: {str(e)}'))

sharding1 = NamedSharding(Mesh(jax.devices()[:4], 'x'), P('x'))
sharding2 = NamedSharding(Mesh(jax.devices()[4:], 'x'), P('x'))

y = jax.device_put(x, sharding1)
z = jax.device_put(x, sharding2)
try: y + z
except ValueError as e: print_exception(e)

ValueError: Received incompatible devices for jitted
computation. Got argument x1 of jax.numpy.add with shape int32[24] and
device ids [0, 1, 2, 3] on platform TPU and argument x2 of
jax.numpy.add with shape int32[24] and device ids [4, 5, 6, 7] on
platform TPU

devices = jax.devices()
permuted_devices = [devices[i] for i in [0, 1, 2, 3, 6, 7, 4, 5]]

sharding1 = NamedSharding(Mesh(devices, 'x'), P('x'))
sharding2 = NamedSharding(Mesh(permuted_devices, 'x'), P('x'))

y = jax.device_put(x, sharding1)
z = jax.device_put(x, sharding2)
try: y + z
except ValueError as e: print_exception(e)

ValueError: Received incompatible devices for jitted
computation. Got argument x1 of jax.numpy.add with shape int32[24] and
device ids [0, 1, 2, 3, 4, 5, 6, 7] on platform TPU and argument x2 of
jax.numpy.add with shape int32[24] and device ids [0, 1, 2, 3, 6, 7,
4, 5] on platform TPU

我们说使用 jax.device_put 显式放置或分片过的数组已“绑定”到其设备(s)，因此不会被自动移动。有关更多信息，请参阅 设备放置常见问题解答。

当数组没有使用 jax.device_put 显式放置或分片时，它们会未绑定地放置在默认设备上。与绑定数组不同，未绑定数组可以被自动移动和重新分片：也就是说，即使其他参数被显式地放置在不同的设备上，未绑定数组也可以作为计算的参数。

例如，jnp.zeros、jnp.arange 和 jnp.array 的输出是未绑定的

y = jax.device_put(x, sharding1)
y + jnp.ones_like(y)
y + jnp.arange(y.size).reshape(y.shape)
print('no error!')

no error!

在 jit 代码中约束中间值的碎片

虽然编译器会尝试决定一个函数的中间值和输出应该如何分片，但我们也可以使用 jax.lax.with_sharding_constraint 向它提供提示。使用 jax.lax.with_sharding_constraint 与 jax.device_put 非常相似，不同的是我们在分阶段执行的（即 jit 装饰的）函数内部使用它

mesh = Mesh(mesh_utils.create_device_mesh((4, 2)), ('x', 'y'))

x = jax.random.normal(jax.random.key(0), (8192, 8192))
x = jax.device_put(x, NamedSharding(mesh, P('x', 'y')))

@jax.jit
def f(x):
  x = x + 1
  y = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P('y', 'x')))
  return y

jax.debug.visualize_array_sharding(x)
y = f(x)
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

┌───────┬───────┬───────┬───────┐
│       │       │       │       │
│ TPU 0 │ TPU 2 │ TPU 6 │ TPU 4 │
│       │       │       │       │
│       │       │       │       │
├───────┼───────┼───────┼───────┤
│       │       │       │       │
│ TPU 1 │ TPU 3 │ TPU 7 │ TPU 5 │
│       │       │       │       │
│       │       │       │       │
└───────┴───────┴───────┴───────┘

@jax.jit
def f(x):
  x = x + 1
  y = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P()))
  return y

jax.debug.visualize_array_sharding(x)
y = f(x)
jax.debug.visualize_array_sharding(y)

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│  TPU 0,1,2,3,4,5,6,7  │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘

通过添加 with_sharding_constraint，我们已经约束了输出的分片。除了尊重特定中间值的注释之外，编译器还会使用注释来决定其他值的碎片。

注释计算的输出通常是一个好习惯，例如，根据这些值最终的消费方式。

示例：神经网络

⚠️ 警告：以下内容旨在简单地演示使用 jax.Array 的自动分片传播，但它可能不反映真实示例的最佳实践。 例如，真实示例可能需要更多地使用 with_sharding_constraint。

我们可以使用 jax.device_put 和 jax.jit 的计算遵循分片的特性来并行化神经网络中的计算。以下是一些简单的示例，基于这个基本的神经网络

import jax
import jax.numpy as jnp

def predict(params, inputs):
  for W, b in params:
    outputs = jnp.dot(inputs, W) + b
    inputs = jnp.maximum(outputs, 0)
  return outputs

def loss(params, batch):
  inputs, targets = batch
  predictions = predict(params, inputs)
  return jnp.mean(jnp.sum((predictions - targets)**2, axis=-1))

loss_jit = jax.jit(loss)
gradfun = jax.jit(jax.grad(loss))

def init_layer(key, n_in, n_out):
  k1, k2 = jax.random.split(key)
  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)
  b = jax.random.normal(k2, (n_out,))
  return W, b

def init_model(key, layer_sizes, batch_size):
  key, *keys = jax.random.split(key, len(layer_sizes))
  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))

  key, *keys = jax.random.split(key, 3)
  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))
  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))

  return params, (inputs, targets)

layer_sizes = [784, 8192, 8192, 8192, 10]
batch_size = 8192

params, batch = init_model(jax.random.key(0), layer_sizes, batch_size)

8 路批次数据并行

mesh = Mesh(mesh_utils.create_device_mesh((8,)), 'batch')

sharding = NamedSharding(mesh, P('batch'))
replicated_sharding = NamedSharding(mesh, P())

batch = jax.device_put(batch, sharding)
params = jax.device_put(params, replicated_sharding)

loss_jit(params, batch)

Array(23.469475, dtype=float32)

step_size = 1e-5

for _ in range(30):
  grads = gradfun(params, batch)
  params = [(W - step_size * dW, b - step_size * db)
            for (W, b), (dW, db) in zip(params, grads)]

print(loss_jit(params, batch))

10.760109

%timeit -n 5 -r 5 gradfun(params, batch)[0][0].block_until_ready()

53.8 ms ± 1.14 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

batch_single = jax.device_put(batch, jax.devices()[0])
params_single = jax.device_put(params, jax.devices()[0])

%timeit -n 5 -r 5 gradfun(params_single, batch_single)[0][0].block_until_ready()

351 ms ± 81.2 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)

4 路批次数据并行和 2 路模型张量并行

mesh = Mesh(mesh_utils.create_device_mesh((4, 2)), ('batch', 'model'))

batch = jax.device_put(batch, NamedSharding(mesh, P('batch', None)))
jax.debug.visualize_array_sharding(batch[0])
jax.debug.visualize_array_sharding(batch[1])

┌───────┐
│TPU 0,1│
├───────┤
│TPU 2,3│
├───────┤
│TPU 6,7│
├───────┤
│TPU 4,5│
└───────┘

┌───────┐
│TPU 0,1│
├───────┤
│TPU 2,3│
├───────┤
│TPU 6,7│
├───────┤
│TPU 4,5│
└───────┘

replicated_sharding = NamedSharding(mesh, P())

(W1, b1), (W2, b2), (W3, b3), (W4, b4) = params

W1 = jax.device_put(W1, replicated_sharding)
b1 = jax.device_put(b1, replicated_sharding)

W2 = jax.device_put(W2, NamedSharding(mesh, P(None, 'model')))
b2 = jax.device_put(b2, NamedSharding(mesh, P('model')))

W3 = jax.device_put(W3, NamedSharding(mesh, P('model', None)))
b3 = jax.device_put(b3, replicated_sharding)

W4 = jax.device_put(W4, replicated_sharding)
b4 = jax.device_put(b4, replicated_sharding)

params = (W1, b1), (W2, b2), (W3, b3), (W4, b4)

jax.debug.visualize_array_sharding(W2)

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

jax.debug.visualize_array_sharding(W3)

┌───────────────────────┐
│                       │
│      TPU 0,2,4,6      │
│                       │
│                       │
├───────────────────────┤
│                       │
│      TPU 1,3,5,7      │
│                       │
│                       │
└───────────────────────┘

print(loss_jit(params, batch))

10.760109

step_size = 1e-5

for _ in range(30):
    grads = gradfun(params, batch)
    params = [(W - step_size * dW, b - step_size * db)
              for (W, b), (dW, db) in zip(params, grads)]

print(loss_jit(params, batch))

10.752513

(W1, b1), (W2, b2), (W3, b3), (W4, b4) = params
jax.debug.visualize_array_sharding(W2)
jax.debug.visualize_array_sharding(W3)

┌───────────┬───────────┐
│           │           │
│           │           │
│           │           │
│           │           │
│TPU 0,2,4,6│TPU 1,3,5,7│
│           │           │
│           │           │
│           │           │
│           │           │
└───────────┴───────────┘

┌───────────────────────┐
│                       │
│      TPU 0,2,4,6      │
│                       │
│                       │
├───────────────────────┤
│                       │
│      TPU 1,3,5,7      │
│                       │
│                       │
└───────────────────────┘

%timeit -n 10 -r 10 gradfun(params, batch)[0][0].block_until_ready()

51.4 ms ± 454 µs per loop (mean ± std. dev. of 10 runs, 10 loops each)

重要细节

生成随机数

JAX 带有一个功能性的、确定性的 随机数生成器。它构成了 jax.random 模块 中的各种采样函数的基础，例如 jax.random.uniform。

JAX 的随机数是由基于计数器的 PRNG 生成的，因此原则上，随机数生成应该是一个对计数器值进行纯映射的操作。纯映射原则上是一个容易分割的操作。它不应该需要跨设备通信，也不应该需要跨设备进行任何冗余计算。

然而，现有的稳定 RNG 实现并非自动可分割的，这是由于历史原因。

考虑以下示例，其中一个函数绘制随机均匀数并逐元素地将它们添加到输入中

@jax.jit
def f(key, x):
  numbers = jax.random.uniform(key, x.shape)
  return x + numbers

key = jax.random.key(42)
mesh = Mesh(jax.devices(), 'x')
x_sharding = NamedSharding(mesh, P('x'))
x = jax.device_put(jnp.arange(24), x_sharding)

在分片输入上，函数 f 生成的输出也是分片的

jax.debug.visualize_array_sharding(f(key, x))

┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ TPU 0 │ TPU 1 │ TPU 2 │ TPU 3 │ TPU 4 │ TPU 5 │ TPU 6 │ TPU 7 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘

但是如果我们检查在该分片输入上对 f 进行编译后的计算，我们会发现它确实涉及一些通信

f_exe = f.lower(key, x).compile()
print('Communicating?', 'collective-permute' in f_exe.as_text())

Communicating? True

解决此问题的一种方法是使用实验性升级标志 jax_threefry_partitionable 配置 JAX。启用该标志后，编译后的计算中不再存在“集体置换”操作

jax.config.update('jax_threefry_partitionable', True)
f_exe = f.lower(key, x).compile()
print('Communicating?', 'collective-permute' in f_exe.as_text())

Communicating? False

输出仍然是分片的

jax.debug.visualize_array_sharding(f(key, x))

┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ TPU 0 │ TPU 1 │ TPU 2 │ TPU 3 │ TPU 4 │ TPU 5 │ TPU 6 │ TPU 7 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘

然而，jax_threefry_partitionable 选项的一个注意事项是，生成的随机值可能与未设置标志时不同，即使它们是由相同的随机键生成的

jax.config.update('jax_threefry_partitionable', False)
print('Stable:')
print(f(key, x))
print()

jax.config.update('jax_threefry_partitionable', True)
print('Partitionable:')
print(f(key, x))

Stable:
[ 0.72503686  1.8532515   2.983416    3.083253    4.0332246   5.4782867
  6.1720605   7.6900277   8.602836    9.810046   10.861367   11.907651
 12.330483   13.456195   14.808557   15.960099   16.067581   17.739723
 18.335474   19.46401    20.390276   21.116539   22.858128   23.223194  ]

Partitionable:
[ 0.48870957  1.6797972   2.6162715   3.561016    4.4506445   5.585866
  6.0748096   7.775133    8.698959    9.818634   10.350306   11.87282
 12.925881   13.86013    14.477554   15.818481   16.711355   17.586697
 18.073738   19.777622   20.404566   21.119123   22.026257   23.63918   ]

在 jax_threefry_partitionable 模式下，JAX PRNG 仍然是确定性的，但其实现是新的（并且正在开发中）。对于给定键生成的随机值在给定的 JAX 版本（或 main 分支上的给定提交）下将保持一致，但可能在不同的版本之间有所不同。